Automated PR - 2026-07-07

packages/ltx-kernels/src/ltx_kernels/__init__.py

@@ -0,0 +1,6 @@
+"""ltx-kernels: High-performance CUDA kernels for distributed attention operations."""
+
+from ltx_kernels.all_to_all import All2All
+
+__all__ = ["All2All"]
+__version__ = "0.1.0"

packages/ltx-kernels/src/ltx_kernels/all_to_all.py

@@ -0,0 +1,183 @@
+"""All2All communication primitives for distributed attention.
+The C++ kernels are exposed via ``torch.library.custom_op`` so that
+``torch.compile`` (and CUDA Graph capture under ``mode="reduce-overhead"``)
+can trace through them without a graph break. Each :class:`All2All`
+instance registers itself in a class-level registry indexed by an integer
+``comm_id``; the custom op takes that id plus an input tensor and dispatches
+to the appropriate C++ runtime.
+"""
+
+import math
+import weakref
+from typing import Any, ClassVar
+
+import torch
+import torch.distributed as dist
+from all2all_cpp import All2All as All2AllCpp
+from torch.library import custom_op
+
+# Output shapes are derived symbolically from the input tensor's shape under
+# the assumption of *uniform* sharding (caller pads up-front so
+# `total_tokens % world_size == 0`). `world_size` travels through the op as an int
+# so it enters the traced graph: the fake needs it for the output shape, and since
+# it is constant per process the Dynamo guard it installs never triggers a
+# within-run recompile -- it only keys the compile cache by world size, so a graph
+# compiled under one GPU count is never replayed under another. Per-rank token
+# counts, which do vary per call, stay in the C++ runtime state (`set_rank_tokens`)
+# and never travel through the op.
+
+
+@custom_op("ltx_kernels::send_recv_heads", mutates_args=(), device_types="cuda")
+def _send_recv_heads_op(
+    x: torch.Tensor,
+    comm_id: int,
+    world_size: int,  # noqa: ARG001
+    copy_out: bool,
+) -> torch.Tensor:
+    # copy_out=False returns a view into the IPC buffer (zero-copy). Safe under
+    # cudagraph_trees because the IPC buffer is cudaMalloc'd inside All2All, not
+    # in the static graph pool, so it isn't subject to graph-pool aliasing.
+    return All2All._runtime_registry[comm_id].send_recv_heads(x, copy_out)
+
+
+@_send_recv_heads_op.register_fake
+def _send_recv_heads_fake(
+    x: torch.Tensor,
+    comm_id: int,  # noqa: ARG001
+    world_size: int,
+    copy_out: bool,  # noqa: ARG001
+) -> torch.Tensor:
+    return x.new_empty((x.shape[0], x.shape[1] * world_size, x.shape[2] // world_size, x.shape[3]))
+
+
+@custom_op("ltx_kernels::gather_heads", mutates_args=(), device_types="cuda")
+def _gather_heads_op(
+    x: torch.Tensor,
+    comm_id: int,
+    world_size: int,  # noqa: ARG001
+    copy_out: bool,
+) -> torch.Tensor:
+    return All2All._runtime_registry[comm_id].gather_heads(x, copy_out)
+
+
+@_gather_heads_op.register_fake
+def _gather_heads_fake(
+    x: torch.Tensor,
+    comm_id: int,  # noqa: ARG001
+    world_size: int,
+    copy_out: bool,  # noqa: ARG001
+) -> torch.Tensor:
+    return x.new_empty((x.shape[0], x.shape[1] // world_size, x.shape[2] * world_size, x.shape[3]))
+
+
+class All2All:
+    """IPC-based All2All communication for distributed head-parallel attention.
+    This class manages GPU memory buffers and IPC handles to enable efficient
+    cross-GPU communication for attention head redistribution.
+    Args:
+        rank: Local rank of this process.
+        world_size: Total number of processes in the distributed group.
+        seqlen: Maximum sequence length to allocate buffers for.
+        hidden_dim: Hidden dimension size (num_heads * head_dim).
+        num_sms: Number of SMs to use for kernel execution.
+        tensor_dtype: Data type for tensors (e.g., torch.bfloat16).
+        group: PyTorch distributed process group.
+    """
+
+    _next_id: ClassVar[int] = 0
+    _runtime_registry: ClassVar[dict[int, "All2AllCpp"]] = {}
+
+    @staticmethod
+    def _release_comm(comm_id: int, runtime: "All2AllCpp") -> None:
+        All2All._runtime_registry.pop(comm_id, None)
+        runtime.destroy()
+
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        seqlen: int,
+        hidden_dim: int,
+        num_sms: int,
+        tensor_dtype: torch.dtype,
+        group: torch.distributed.ProcessGroup | None = None,
+        timeout_seconds: float | None = None,
+    ) -> None:
+        self.rank = rank
+        self.world_size = world_size
+        self.num_sms = num_sms
+        self.tensor_dtype = tensor_dtype
+        self.buffer_size = int(seqlen * hidden_dim * tensor_dtype.itemsize)
+
+        # Initialize the C++ runtime. Omit timeout_seconds to keep the kernel's default
+        # (DEFAULT_BARRIER_TIMEOUT_SECONDS); it can still be changed later via set_timeout_seconds.
+        if timeout_seconds is None:
+            self.runtime = All2AllCpp(rank, world_size, seqlen, hidden_dim, num_sms, tensor_dtype)
+        else:
+            if not math.isfinite(timeout_seconds) or timeout_seconds < 0:
+                raise ValueError(f"all2all timeout seconds must be finite and non-negative, got {timeout_seconds}")
+            self.runtime = All2AllCpp(rank, world_size, seqlen, hidden_dim, num_sms, tensor_dtype, timeout_seconds)
+
+        # Register in the class-level table so the custom ops can find us.
+        # Auto-cleanup via weakref.finalize covers the case where destroy()
+        # isn't called explicitly — without it the registry would keep the
+        # runtime (and CUDA/IPC resources) alive for the process lifetime.
+        self._comm_id = All2All._next_id
+        All2All._next_id += 1
+        All2All._runtime_registry[self._comm_id] = self.runtime
+        self._finalizer = weakref.finalize(self, All2All._release_comm, self._comm_id, self.runtime)
+
+        # Exchange IPC handles across all ranks
+        ipc_handles: list[Any] = [None] * world_size
+        local_ipc_handle = self.runtime.get_local_ipc_handle()
+        dist.all_gather_object(ipc_handles, local_ipc_handle, group)
+
+        self.runtime.sync(ipc_handles)
+
+    def set_rank_tokens(self, rank_num_tokens: list[int]) -> None:
+        """Sets per-rank token counts on the C++ runtime."""
+        self.runtime.set_rank_tokens(rank_num_tokens)
+
+    def set_timeout_seconds(self, seconds: float) -> None:
+        """Set the all2all barrier (deadlock-detection) timeout, in seconds.
+        The C++ runtime converts to clock cycles using the device's peak SM clock. Raise it
+        during the first ``torch.compile`` forward, where one rank's recompile can delay its
+        all2all launch past the steady-state timeout; reset for steady state.
+        """
+        if not math.isfinite(seconds) or seconds < 0:
+            raise ValueError(f"all2all timeout seconds must be finite and non-negative, got {seconds}")
+        self.runtime.set_timeout_seconds(seconds)
+
+    def send_recv_heads(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Exchange attention heads across ranks (All2All pattern).
+        Args:
+            x: Input tensor of shape [batch, tokens, heads, head_dim].
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+        Returns:
+            Output tensor with redistributed heads.
+        """
+        return torch.ops.ltx_kernels.send_recv_heads(x, self._comm_id, self.world_size, copy_out)
+
+    def gather_heads(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Gather heads back to original distribution (reverse All2All).
+        Args:
+            x: Input tensor with distributed heads.
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+        Returns:
+            Output tensor with gathered heads.
+        """
+        return torch.ops.ltx_kernels.gather_heads(x, self._comm_id, self.world_size, copy_out)
+
+    def allgather(self, x: torch.Tensor, *, copy_out: bool = False) -> torch.Tensor:
+        """Allgather operation across all ranks.
+        Args:
+            x: Input tensor to gather.
+            copy_out: If True, copy result to a new tensor instead of using buffer.
+        Returns:
+            Gathered tensor from all ranks.
+        """
+        return self.runtime.allgather(x, copy_out)
+
+    def destroy(self) -> None:
+        """Release IPC handles and GPU memory buffers."""
+        self._finalizer()

packages/ltx-kernels/src/ltx_kernels/arch.py

@@ -0,0 +1,21 @@
+"""GPU architecture detection for blockwise kernel dispatch."""
+
+import torch
+
+
+def get_device_arch() -> str:
+    """Return a coarse architecture name for the current CUDA device.
+    Used to pick the FP8 GEMM kernel variant (SM89 vs SM90) at runtime.
+    """
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is not available; blockwise kernels require a CUDA device.")
+    major, minor = torch.cuda.get_device_capability(torch.cuda.current_device())
+    if major == 8 and (minor >= 0 and minor < 9):
+        return "ampere"
+    if major == 8 and minor == 9:
+        return "ada"
+    if major == 9 and minor == 0:
+        return "hopper"
+    if major in {10, 12}:
+        return "blackwell"
+    raise NotImplementedError(f"Unsupported GPU compute capability sm_{major}{minor}.")

packages/ltx-kernels/src/ltx_kernels/blockwise/__init__.py

@@ -0,0 +1,36 @@
+"""Blockwise FP8/FP6 quantization kernels.
+Importing this subpackage pulls in the compiled ``ops_cpp`` / ``blockwise_cpp``
+extensions (via :mod:`.functional` and :mod:`.linear`) and ``triton``, so it
+raises :class:`ImportError` on hosts where the kernels were not built. Callers
+that want a soft dependency should guard the import (e.g. ``pytest.importorskip``
+or the lazy gate in ``ltx_core.quantization.blockwise``).
+"""
+
+from ltx_kernels.blockwise.functional import (
+    blockwise_dequantize,
+    blockwise_quantize_adanorm_triton,
+    blockwise_quantize_rms_fma_triton,
+    fp6_blockwise_quantize_weights_torch,
+    fp6_pack_tensor,
+    fp6_unpack_tensor,
+    fp8_blockwise_quantize_weights_torch,
+    gated_attention_triton,
+    rms_norm_rope,
+    rms_norm_split_rope,
+)
+from ltx_kernels.blockwise.linear import BlockwiseFP6Linear, BlockwiseFP8Linear
+
+__all__ = [
+    "BlockwiseFP6Linear",
+    "BlockwiseFP8Linear",
+    "blockwise_dequantize",
+    "blockwise_quantize_adanorm_triton",
+    "blockwise_quantize_rms_fma_triton",
+    "fp6_blockwise_quantize_weights_torch",
+    "fp6_pack_tensor",
+    "fp6_unpack_tensor",
+    "fp8_blockwise_quantize_weights_torch",
+    "gated_attention_triton",
+    "rms_norm_rope",
+    "rms_norm_split_rope",
+]

packages/ltx-kernels/src/ltx_kernels/blockwise/functional.py

@@ -0,0 +1,234 @@
+"""Blockwise FP8/FP6 functional ops.
+The CUDA kernels live in the ``ops_cpp`` extension.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+from ops_cpp import fp6_pack, fp6_unpack
+from ops_cpp import rms_norm_rope as rms_norm_rope_cuda
+from ops_cpp import rms_norm_split_rope as rms_norm_split_rope_cuda
+
+from ltx_kernels.blockwise.triton_ops import (
+    blockwise_dequantize_triton,
+    blockwise_quantize_adanorm_triton,  # noqa: F401  (re-exported for consumers)
+    blockwise_quantize_rms_fma_triton,  # noqa: F401  (re-exported for consumers)
+    blockwise_quantize_triton,  # noqa: F401  (used by linear.BlockwiseGemmLinearFunc)
+    gated_attention_triton,  # noqa: F401  (re-exported for consumers)
+)
+
+# Quantization scale constants for different precisions
+FP8_SCALE_MAX = 448.0
+FP6_SCALE_MAX = 0.1172
+
+
+@torch.library.custom_op("q8_kernels_ops::rms_norm_rope", mutates_args=())
+def _rms_norm_rope_cuda(
+    x: torch.Tensor,
+    weights: Optional[torch.Tensor],
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    out_16bit: bool,
+) -> torch.Tensor:
+    return rms_norm_rope_cuda(x, weights, cos_freqs, sin_freqs, out_16bit)
+
+
+@torch.library.register_fake("q8_kernels_ops::rms_norm_rope")
+def _rms_norm_rope_cuda_fake(
+    x: torch.Tensor,
+    weights: Optional[torch.Tensor],
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    out_16bit: bool,
+) -> torch.Tensor:
+    out = torch.empty_like(x)
+    if out_16bit:
+        return out.to(torch.bfloat16)
+    else:
+        return out.to(torch.float8_e4m3fn)
+
+
+@torch.library.custom_op("q8_kernels_ops::rms_norm_split_rope", mutates_args=())
+def _rms_norm_split_rope_cuda(
+    x: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    weights: torch.Tensor,
+    out_fp8: bool,
+) -> torch.Tensor:
+    return rms_norm_split_rope_cuda(x, sin_freqs, cos_freqs, weights, out_fp8)
+
+
+@torch.library.register_fake("q8_kernels_ops::rms_norm_split_rope")
+def _rms_norm_split_rope_cuda_fake(
+    x: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    weights: torch.Tensor,
+    out_fp8: bool,
+) -> torch.Tensor:
+    out = torch.empty_like(x)
+    if out_fp8:
+        return out.to(torch.float8_e4m3fn)
+    else:
+        return out.to(torch.bfloat16)
+
+
+class RMSNormRope(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, weights, cos_freqs, sin_freqs, out_16bit):
+        return torch.ops.q8_kernels_ops.rms_norm_rope(
+            x, weights, cos_freqs, sin_freqs, out_16bit
+        )
+
+
+class RMSNormSplitRope(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, sin_freqs, cos_freqs, weights, out_fp8):
+        return torch.ops.q8_kernels_ops.rms_norm_split_rope(
+            x, sin_freqs, cos_freqs, weights, out_fp8
+        )
+
+
+def rms_norm_rope(
+    x: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    weights: Optional[torch.Tensor] = None,
+    out_fp8: bool = False,
+) -> torch.Tensor:
+    # weights=None flows through to the kernel's no-affine path (nullptr); see
+    # rms_norm_rope.cpp, which dispatches the norm_affine template on has_value().
+    return RMSNormRope.apply(x, weights, cos_freqs, sin_freqs, not out_fp8)
+
+
+def rms_norm_split_rope(
+    x: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+    weights: Optional[torch.Tensor] = None,
+    out_fp8: bool = False,
+) -> torch.Tensor:
+    """
+    Apply RMS normalization followed by split RoPE (Rotary Position Embedding).
+    Args:
+        x: Input tensor of shape [b, s, h] where h can be 2048, 4096, or 8192
+        cos_freqs: Cos frequencies of shape [b, n, s, d] where n=32 and n*d = h/2
+        sin_freqs: Sin frequencies of shape [b, n, s, d] where n=32
+        weights: Optional RMS norm weights of shape [h]
+        out_fp8: If True, output in float8_e4m3fn, otherwise bfloat16
+    Returns:
+        Output tensor of shape [b, s, h] with dtype based on out_fp8 parameter
+    """
+    if weights is None:
+        weights = torch.ones(x.shape[-1], dtype=x.dtype, device=x.device)
+    return RMSNormSplitRope.apply(x, sin_freqs, cos_freqs, weights, out_fp8)
+
+
+def blockwise_quantize_weights(w: torch.Tensor, block_size=128, scale_max=FP8_SCALE_MAX) -> tuple[torch.Tensor, torch.Tensor]:
+    w = w.view(w.shape[0]//block_size, block_size, w.shape[1]//block_size, block_size).transpose(1, 2).contiguous()
+    w_absmax = w.float().abs().view(w.shape[0], w.shape[1], block_size*block_size).max(dim=-1, keepdim=False).values
+    w_scales = scale_max/w_absmax
+    w_quant = (w.float()*w_scales[:, :, None, None].float()).to(torch.float8_e4m3fn)
+    w_quant = w_quant.transpose(1, 2).contiguous()
+    w_quant = w_quant.view(w.shape[0]*block_size, -1)
+    return w_quant, (1/w_scales).contiguous()
+
+
+def blockwise_quantize_torch(x: torch.Tensor, block_size=128, scale_max=FP8_SCALE_MAX) -> tuple[torch.Tensor, torch.Tensor]:
+    b, n, h = x.shape
+    x = x.view(-1, x.shape[-1]//block_size, block_size)
+    x_absmax = x.float().abs().max(dim=-1, keepdim=True).values
+    x_scales = scale_max/x_absmax
+    x_scaled = (x * x_scales).to(torch.float8_e4m3fn)
+    return x_scaled.view(b, n, h), 1/x_scales.view(b*n, h//block_size).t().contiguous().float()
+
+
+# Precision-specific convenience functions
+
+def fp8_blockwise_quantize_weights_torch(x: torch.Tensor, block_size=128) -> tuple[torch.Tensor, torch.Tensor]:
+    return blockwise_quantize_weights(x, block_size, FP8_SCALE_MAX)
+
+
+def fp6_blockwise_quantize_weights_torch(x: torch.Tensor, block_size=128) -> tuple[torch.Tensor, torch.Tensor]:
+    return blockwise_quantize_weights(x, block_size, FP6_SCALE_MAX)
+
+
+def blockwise_dequantize(x: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    x_fp8, scales = x
+    return blockwise_dequantize_triton(x_fp8, scales)
+
+
+# FP6 Pack/Unpack operations
+@torch.library.custom_op("q8_kernels_ops::fp6_pack", mutates_args=())
+def _fp6_pack(x: torch.Tensor) -> torch.Tensor:
+    """Pack 8-bit tensor to 6-bit by dropping e_1 and e_2 bits.
+    Args:
+        x: Input tensor of shape [m, n] with dtype uint8 or float8_e4m3fn
+    Returns:
+        Packed tensor of shape [m, n*3/4] with dtype uint8
+    """
+    return fp6_pack(x)
+
+
+@torch.library.register_fake("q8_kernels_ops::fp6_pack")
+def _fp6_pack_fake(x: torch.Tensor) -> torch.Tensor:
+    m, n = x.shape
+    n_packed = (n * 3) // 4
+    return torch.empty(m, n_packed, dtype=torch.uint8, device=x.device)
+
+
+@torch.library.custom_op("q8_kernels_ops::fp6_unpack", mutates_args=())
+def _fp6_unpack(x: torch.Tensor, original_n: int) -> torch.Tensor:
+    """Unpack 6-bit tensor back to 8-bit (with e_1 and e_2 set to 0).
+    Args:
+        x: Packed tensor of shape [m, n_packed] with dtype uint8
+        original_n: Original dimension size (unpacked)
+    Returns:
+        Unpacked tensor of shape [m, original_n] with dtype uint8
+    """
+    return fp6_unpack(x, original_n)
+
+
+@torch.library.register_fake("q8_kernels_ops::fp6_unpack")
+def _fp6_unpack_fake(x: torch.Tensor, original_n: int) -> torch.Tensor:
+    m = x.shape[0]
+    return torch.empty(m, original_n, dtype=torch.uint8, device=x.device)
+
+
+# Convenience wrapper functions
+def fp6_pack_tensor(x: torch.Tensor) -> torch.Tensor:
+    """Pack 8-bit tensor to 6-bit by dropping e_1 and e_2 bits.
+    This function packs FP8 weights into a more memory-efficient FP6 format
+    by dropping the two highest exponent bits (e_1 and e_2). This achieves
+    25% memory reduction (4 bytes -> 3 bytes per 4 elements).
+    Args:
+        x: Input tensor of shape [m, n] with dtype uint8 or float8_e4m3fn.
+           n must be divisible by 8.
+    Returns:
+        Packed tensor of shape [m, n*3/4] with dtype uint8
+    Example:
+        >>> w_fp8 = torch.randn(1024, 4096, dtype=torch.float8_e4m3fn, device='cuda')
+        >>> w_fp8_uint8 = w_fp8.view(torch.uint8)
+        >>> w_packed = fp6_pack_tensor(w_fp8_uint8)
+        >>> print(w_packed.shape)  # [1024, 3072]
+    """
+    return torch.ops.q8_kernels_ops.fp6_pack(x)
+
+
+def fp6_unpack_tensor(x: torch.Tensor, original_n: int) -> torch.Tensor:
+    """Unpack 6-bit tensor back to 8-bit (with e_1 and e_2 set to 0).
+    This function unpacks FP6 weights back to FP8 format. Note that the
+    two dropped exponent bits (e_1 and e_2) are restored as 0.
+    Args:
+        x: Packed tensor of shape [m, n_packed] with dtype uint8
+        original_n: Original dimension size before packing (must be divisible by 8)
+    Returns:
+        Unpacked tensor of shape [m, original_n] with dtype uint8
+    Example:
+        >>> w_packed = torch.randint(0, 256, (1024, 3072), dtype=torch.uint8, device='cuda')
+        >>> w_unpacked = fp6_unpack_tensor(w_packed, 4096)
+        >>> w_fp8 = w_unpacked.view(torch.float8_e4m3fn)
+        >>> print(w_fp8.shape)  # [1024, 4096]
+    """
+    return torch.ops.q8_kernels_ops.fp6_unpack(x, original_n)

packages/ltx-kernels/src/ltx_kernels/blockwise/linear.py

@@ -0,0 +1,279 @@
+"""Blockwise FP8/FP6 quantized linear layers and the FP8 GEMM custom op.
+The FP8 GEMM kernels live in the ``blockwise_cpp`` extension.
+"""
+
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+from ops_cpp import fp6_pack, fp6_unpack
+
+from ltx_kernels.arch import get_device_arch
+from ltx_kernels.blockwise.functional import (
+    blockwise_quantize_triton,
+    blockwise_quantize_weights,
+    fp6_blockwise_quantize_weights_torch,
+)
+
+# Lazily-initialized, architecture-specific FP8 GEMM callable. Importing this
+# module must not require a CUDA device or the compiled extension; the kernel is
+# only resolved on first use (matching the lazy-import gate in ltx-core's
+# blockwise.__init__).
+_fp8_gemm = None
+
+
+def get_fp8_gemm_nt():
+    """Runtime GPU architecture dispatch for FP8 GEMM."""
+    arch = get_device_arch()
+    if arch in ["ada", "blackwell"]:
+        # Ada/Blackwell use SM89 kernels (GeForce path)
+        from blockwise_cpp import fp8_gemm_nt_sm89
+
+        def _func(a, b, d, bias=None, c=None, num_sms=132, use_fast_accum=True):
+            return fp8_gemm_nt_sm89(a, b, d, bias=bias, use_fast_accum=use_fast_accum)
+
+        return _func
+    elif arch == "hopper":
+        # Hopper uses SM90 kernels (H100 path)
+        from blockwise_cpp import fp8_gemm_nt_sm90
+
+        def _func(a, b, d, bias=None, c=None, num_sms=132, use_fast_accum=True):
+            return fp8_gemm_nt_sm90(a, b, d, bias=bias, c=c, num_sms=num_sms)
+
+        return _func
+    else:
+        raise RuntimeError(f"Unsupported GPU architecture: {arch}")
+
+
+def _fp8_gemm_dispatch(a, b, d, bias=None, c=None, num_sms=132, use_fast_accum=True):
+    global _fp8_gemm
+    if _fp8_gemm is None:
+        _fp8_gemm = get_fp8_gemm_nt()
+    return _fp8_gemm(a, b, d, bias=bias, c=c, num_sms=num_sms, use_fast_accum=use_fast_accum)
+
+
+@torch.library.custom_op("blockwise::fp8_gemm", mutates_args=())
+def blockwise_fp8_gemm(
+    a: List[torch.Tensor],
+    b: List[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    use_fast_accum: bool,
+) -> torch.Tensor:
+    d = torch.empty(a[0].shape[0], b[0].shape[0], dtype=torch.bfloat16, device=b[0].device)
+    _fp8_gemm_dispatch(a, b, d, bias=bias, use_fast_accum=use_fast_accum)
+    return d
+
+
+@blockwise_fp8_gemm.register_fake
+def blockwise_fp8_gemm_fake(
+    a: List[torch.Tensor],
+    b: List[torch.Tensor],
+    bias: Optional[torch.Tensor],
+    use_fast_accum: bool,
+) -> torch.Tensor:
+    o = torch.empty(a[0].shape[0], b[0].shape[0], dtype=torch.bfloat16, device=b[0].device)
+    return o
+
+
+def is_16bit(x) -> bool:
+    return x.dtype == torch.float16 or x.dtype == torch.bfloat16
+
+
+class BlockwiseGemmLinearFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        a: tuple[torch.Tensor, torch.Tensor],
+        w: tuple[torch.Tensor],
+        bias: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+
+        is_16bit_a = is_16bit(a[0])
+        if is_16bit_a:
+            a = blockwise_quantize_triton(a[0])
+
+        b, n, h = a[0].shape
+        out_h, _ = w[0].shape
+        new_a = (a[0].view(-1, h), a[1])
+        d = blockwise_fp8_gemm(new_a, w, bias, use_fast_accum=True)
+        return d.view(b, n, out_h)
+
+
+def blockwise_linear_func(a, b, bias=None):
+    return BlockwiseGemmLinearFunc.apply(a, b, bias)
+
+
+class BlockwiseLinear(nn.Module):
+    """Base class for blockwise quantized linear layers."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,  # for compatibility with nn.Linear
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+
+        if bias:
+            # FP32 to match the SM89/SM90 GEMM bias contract (the kernels take float*).
+            self.bias = nn.Parameter(
+                torch.empty(out_features, device=device, dtype=torch.float32),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    @property
+    def fp8weight(self):
+        """Should be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement the weight property")
+
+    def forward(self, x):
+        if not isinstance(x, tuple):
+            x = (x, None)
+        return blockwise_linear_func(x, (self.fp8weight, self.weight_scale), self.bias)
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights=True):
+        """Should be implemented by subclasses."""
+        raise NotImplementedError("Subclasses must implement from_linear method")
+
+
+class BlockwiseFP8Linear(BlockwiseLinear):
+    """Blockwise quantized linear layer using FP8 weights."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__(in_features, out_features, bias, device, dtype)
+
+        assert in_features % 128 == 0, f"in_features must be divisible by 128, got {in_features}"
+        assert out_features % 128 == 0, f"out_features must be divisible by 128, got {out_features}"
+
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features, in_features, device=device, dtype=torch.float8_e4m3fn
+            ),
+            requires_grad=False,
+        )
+        self.weight_scale = nn.Parameter(
+            torch.ones(out_features//128, in_features//128, device=device, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+    @property
+    def fp8weight(self):
+        """Return FP8 weight directly."""
+        return self.weight
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights=True):
+        layer = cls(
+            linear.in_features,
+            linear.out_features,
+            linear.bias is not None,
+            device=linear.weight.device,
+        )
+
+        if transform_weights:
+            w_fp8, w_scales = blockwise_quantize_weights(linear.weight.data.cuda())
+        else:
+            w_fp8 = torch.empty(linear.out_features, linear.in_features, dtype=torch.float8_e4m3fn, device=linear.weight.device)
+            w_scales = torch.ones(
+                w_fp8.shape[0]//128,
+                w_fp8.shape[1]//128,
+                device=w_fp8.device,
+                dtype=torch.float32,
+            )
+
+        layer.weight.data = w_fp8
+        layer.weight_scale.data = w_scales
+
+        if linear.bias is not None:
+            layer.bias.data = linear.bias.data.to(torch.float32)
+
+        return layer
+
+
+class BlockwiseFP6Linear(BlockwiseLinear):
+    """Blockwise quantized linear layer using FP6 weights with 6-bit packing."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        assert in_features % 4 == 0, f"in_features must be divisible by 4, got {in_features}"
+        assert in_features % 8 == 0, f"in_features must be divisible by 8 for fp6_pack, got {in_features}"
+        assert in_features % 128 == 0, f"in_features must be divisible by 128, got {in_features}"
+        assert out_features % 128 == 0, f"out_features must be divisible by 128, got {out_features}"
+
+        super().__init__(in_features, out_features, bias, device, dtype)
+
+        # Store packed weight: [out_features, (in_features // 4)*3]
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features, (in_features // 4)*3, device=device, dtype=torch.uint8
+            ),
+            requires_grad=False,
+        )
+
+        # Scales are based on unpacked dimensions
+        self.weight_scale = nn.Parameter(
+            torch.ones(out_features//128, in_features//128, device=device, dtype=torch.float32),
+            requires_grad=False,
+        )
+
+    @property
+    def fp8weight(self):
+        """Unpack the 6-bit weights to 8-bit format on-the-fly."""
+        # Unpack from [out_features, (in_features // 4)*3] to [out_features, in_features]
+        unpacked = fp6_unpack(self.weight, self.in_features)
+        # Convert uint8 back to float8_e4m3fn view
+        return unpacked.view(torch.float8_e4m3fn)
+
+    @classmethod
+    def from_linear(cls, linear: nn.Linear, transform_weights=True):
+        layer = cls(
+            linear.in_features,
+            linear.out_features,
+            linear.bias is not None,
+            device=linear.weight.device,
+        )
+
+        if transform_weights:
+            # Quantize to FP6 (actually FP8 with restricted range)
+            w_fp6, w_scales = fp6_blockwise_quantize_weights_torch(linear.weight.data.cuda())
+
+            # Pack the FP6 weights from [out, in] uint8 to [out, in*3/4] uint8
+            w_fp6_uint8 = w_fp6.view(torch.uint8)
+            w_packed = fp6_pack(w_fp6_uint8)
+
+            layer.weight.data = w_packed
+            layer.weight_scale.data = w_scales
+        else:
+            # If not transforming, assume weights are already in packed format
+            layer.weight.data = torch.empty(linear.out_features, (linear.in_features // 4) * 3, dtype=torch.uint8, device=linear.weight.device)
+            layer.weight_scale.data = torch.ones(
+                linear.weight.shape[0]//128,
+                linear.in_features//128,
+                device=linear.weight.device,
+                dtype=torch.float32,
+            )
+
+        if linear.bias is not None:
+            layer.bias.data = linear.bias.data.to(torch.float32)
+
+        return layer

packages/ltx-kernels/src/ltx_kernels/blockwise/triton_ops.py

@@ -0,0 +1,453 @@
+from typing import *
+
+import torch
+from triton import jit
+from triton import language as tl
+from triton import next_power_of_2
+from torch.library import triton_op, wrap_triton
+
+def get_tma_aligned_size(n, element_size):
+    num_elems_tma = 16 // element_size
+    return ((n + num_elems_tma - 1) // num_elems_tma) * num_elems_tma 
+
+@jit
+def _quantize(x, scale_max: tl.constexpr, NUM_BLOCKS: tl.constexpr, BLOCK_SIZE: tl.constexpr):
+    x = tl.reshape(x, (NUM_BLOCKS, BLOCK_SIZE))
+    x_abs = tl.abs(x)
+    x_abs = tl.broadcast_to(x_abs, (NUM_BLOCKS, BLOCK_SIZE))
+    x_absmax = tl.max(x_abs, axis=1)[:, None]
+    x_scales = scale_max / x_absmax
+    x_quant = (x_scales * x).to(tl.float8e4nv)
+    x_out_scales = 1.0/x_scales
+    return x_quant, x_out_scales
+
+@jit
+def _kernel(X, OUT, SCALES, HDIM, BLOCK_SIZE: tl.constexpr):
+    row_idx = tl.program_id(0)
+    x_ptr = X + row_idx * HDIM
+    out_ptr = OUT + row_idx * HDIM
+
+    h_offset = tl.arange(0, BLOCK_SIZE)
+
+    x = tl.load(x_ptr + h_offset, mask=h_offset < HDIM).to(tl.float32)
+    x_scale = 127.0 / tl.max(tl.abs(x))
+    x_scaled = x * x_scale
+    x_scaled += (0.5 * tl.where(x_scaled >= 0, 1, -1)).to(tl.int8)
+
+    tl.store(out_ptr + h_offset, x_scaled, mask=h_offset < HDIM)
+    tl.store(SCALES + row_idx, 1 / x_scale)
+
+def run_quantize_kernel(x: torch.Tensor, out_dtype: Optional[torch.dtype] = None):
+    x_shape_orig = x.shape
+    x = x.view(-1, x_shape_orig[-1])
+    out = torch.empty(x_shape_orig, dtype=torch.int8, device=x.device)
+    scales = torch.empty(x.shape[0], dtype=torch.float, device=x.device)
+
+    BLOCK_SIZE = next_power_of_2(x_shape_orig[-1])
+    grid = (x.shape[0],)
+    _kernel[grid](x, out, scales, x_shape_orig[-1], BLOCK_SIZE, num_warps=4)
+
+    return out.view(x_shape_orig), scales.view(x_shape_orig[:-1])
+
+@jit
+def _block_quant_norm_kernel(
+    X,
+    Norm_Scale,
+    Norm_Shift,
+    Out_scales,
+    X_out,
+    norm_scale_batch_stride: int,
+    norm_shift_batch_stride: int,
+    norm_scale_token_stride: int,
+    norm_shift_token_stride: int,
+    seqlen: int,
+    H: tl.constexpr,
+    BROADCAST_SEQLEN: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TMA_ALIGNED_MN: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    batch_id =  token_idx // seqlen
+
+    x_ptr = X + token_idx*H + tl.arange(0, H)
+    if BROADCAST_SEQLEN:
+        norm_scale_ptr = Norm_Scale + batch_id*norm_scale_batch_stride + tl.arange(0, H)
+        norm_shift_ptr = Norm_Shift + batch_id*norm_shift_batch_stride + tl.arange(0, H)
+    else:
+        norm_scale_ptr = Norm_Scale + batch_id*norm_scale_batch_stride + (token_idx % seqlen) * norm_scale_token_stride + tl.arange(0, H)
+        norm_shift_ptr = Norm_Shift + batch_id*norm_shift_batch_stride + (token_idx % seqlen) * norm_shift_token_stride + tl.arange(0, H)
+
+    norm_scales = tl.load(norm_scale_ptr)
+    norm_shift = tl.load(norm_shift_ptr)
+    x = tl.load(x_ptr)
+
+    x_sqr = x*x
+    x_norm = tl.sum(x_sqr) / H
+    x_norm = tl.rsqrt(x_norm + 0.00001)
+    x = (x * x_norm).to(tl.bfloat16)
+    x = x * (1.0+norm_scales) + norm_shift
+
+    o_quant, o_scales = _quantize(x, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+    x_out_ptr = X_out + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    tl.store(x_out_ptr, o_quant)
+    out_scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN*tl.arange(0, NUM_BLOCKS)[:, None]
+    tl.store(out_scales_ptr, o_scales)
+    
+@jit
+def _gelu(x):
+    return x * tl.sigmoid(1.702*x)
+
+@jit
+def _block_quant_kernel(
+    X,
+    Out_scales,
+    X_out,
+    seqlen: int,
+    H: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TMA_ALIGNED_MN: tl.constexpr,
+    USE_GELU: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    batch_id, seq_id = tl.program_id(0), tl.program_id(1)
+    token_idx = batch_id*seqlen + seq_id
+    x_ptr = X + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    x = tl.load(x_ptr).to(tl.float32)
+    if USE_GELU:
+        x = _gelu(x)
+
+    o_quant, o_scales = _quantize(x, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+    x_out_ptr = X_out + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    tl.store(x_out_ptr, o_quant)
+    out_scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN*tl.arange(0, NUM_BLOCKS)[:, None]
+    tl.store(out_scales_ptr, o_scales)
+
+@triton_op("blockwise::quantize", mutates_args=())
+def _quant_blockwise_tma_aligned_func(x: torch.Tensor, use_gelu: bool) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, s, h = x.shape
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+    tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+    scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    scale_max: float = 448.0
+    wrap_triton(_block_quant_kernel)[(b, s)](
+        x,
+        scales,
+        out,
+        H=h,
+        seqlen=s,
+        BLOCK_SIZE=128,
+        NUM_BLOCKS=num_blocks,
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        USE_GELU=use_gelu,
+        SCALE_MAX=scale_max
+    )
+
+    return out.view(b, s, h), scales
+
+def run_quant_blockwise_tma_aligned(x):
+    return _quant_blockwise_tma_aligned_func(x, False)
+
+def run_quant_blockwise_gelu_tma_aligned(x):
+    return _quant_blockwise_tma_aligned_func(x, True)
+
+@triton_op("blockwise::adanorm", mutates_args=())
+def run_quant_blockwise_norm_tma_aligned(x: torch.Tensor, w: Optional[torch.Tensor], norm_scale: torch.Tensor, norm_shift: torch.Tensor, out_dtype: torch.dtype, hd_scale: Optional[float]) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, s, h = x.shape
+    is_broadcast = (norm_scale.shape[1] == 1)
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+    tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+    scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    norm_scale_batch_stride = norm_scale.stride(0)
+    norm_shift_batch_stride = norm_shift.stride(0)
+
+    scale_max: float = 448.0
+    wrap_triton(_block_quant_norm_kernel)[(num_rows, 1, 1)](
+        x,
+        norm_scale,
+        norm_shift,
+        scales,
+        out,
+        norm_scale_batch_stride=norm_scale_batch_stride,
+        norm_shift_batch_stride=norm_shift_batch_stride,
+        norm_scale_token_stride=norm_scale.stride(1),
+        norm_shift_token_stride=norm_shift.stride(1),
+        H=h,
+        seqlen=s,
+        BROADCAST_SEQLEN=is_broadcast,
+        BLOCK_SIZE=128,
+        NUM_BLOCKS=num_blocks,
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        SCALE_MAX=scale_max
+    )
+
+    return out.view(b, s, h), scales
+
+@jit
+def _quant_rms_sum_mult_kernel(
+    X,
+    Y,
+    Z,
+    Out,
+    Out_scales,
+    seqlen: int,
+    z_batch_stride: int,
+    z_token_stride: int,
+    TMA_ALIGNED_MN: int,
+    H: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    IS_Z_BROADCAST: tl.constexpr,
+    QUANTIZE: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    batch_id =  token_idx // seqlen
+
+    x_ptr = X + token_idx*H + tl.arange(0, H)
+    y_ptr = Y + token_idx*H + tl.arange(0, H)
+    if IS_Z_BROADCAST:
+        z_ptr = Z + batch_id*z_batch_stride + tl.arange(0, H)
+    else:
+        z_ptr = Z + batch_id*z_batch_stride + (token_idx % seqlen) * z_token_stride + tl.arange(0, H)
+    x = tl.load(x_ptr)
+    y = tl.load(y_ptr)
+    z = tl.load(z_ptr)
+
+    o = x + y*z
+
+    tl.store(x_ptr, o)
+
+    o_sqr = o * o
+    o_norm = tl.sum(o_sqr, axis=0) / H
+    o_inv = tl.rsqrt(o_norm)
+    o *= o_inv
+
+    if QUANTIZE:
+        o_quant, o_scales = _quantize(o, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+        x_out_ptr = Out + token_idx*H + BLOCK_SIZE*tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+        tl.store(x_out_ptr, o_quant)
+        out_scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN*tl.arange(0, NUM_BLOCKS)[:, None]
+        tl.store(out_scales_ptr, o_scales)
+    else:
+        x_out_ptr = Out + token_idx*H + tl.arange(0, H)
+        tl.store(x_out_ptr, o)
+
+@triton_op("blockwise::quant_rms_fma", mutates_args=("x", ))
+def run_quant_blockwise_rms_fma(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, s, h = x.shape
+    scale_max: float = 448.0
+    is_z_broadcast = (z.shape[1] == 1)
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+    tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+    scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    z_batch_stride = z.stride(0)
+
+    wrap_triton(_quant_rms_sum_mult_kernel)[(num_rows, 1, 1)](
+        x,
+        y,
+        z,
+        out,
+        scales,
+        seqlen=s,
+        z_batch_stride=z_batch_stride,
+        z_token_stride=z.stride(1),
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        H=h,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+        IS_Z_BROADCAST=is_z_broadcast,
+        QUANTIZE=True,
+        SCALE_MAX=scale_max
+    )
+
+    return out.view(b, s, h), scales
+
+
+@triton_op("blockwise::rms_fma", mutates_args=("x", ))
+def run_rms_fma(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+    b, s, h = x.shape
+    is_z_broadcast = (z.shape[1] == 1)
+    num_rows = b*s
+    block_size = 128
+    num_blocks = h // block_size
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.bfloat16)
+    z_batch_stride = z.stride(0)
+    
+    wrap_triton(_quant_rms_sum_mult_kernel)[(num_rows, 1, 1)](
+        x,
+        y,
+        z,
+        out,
+        None,
+        seqlen=s,
+        z_batch_stride=z_batch_stride,
+        z_token_stride=z.stride(1),
+        TMA_ALIGNED_MN=0,
+        H=h,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+        IS_Z_BROADCAST=is_z_broadcast,
+        QUANTIZE=False,
+        SCALE_MAX=448.0
+    )
+
+    return out.view(b, s, h)
+
+
+@jit
+def _gated_attention_kernel(
+    X,
+    Gate_Logits,
+    Out,
+    Out_scales,
+    H: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    DIM_HEAD: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+    TMA_ALIGNED_MN: tl.constexpr,
+    QUANTIZE: tl.constexpr,
+    SCALE_MAX: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+
+    # Load gate logits for all heads: shape [NUM_HEADS]
+    gate_logits = tl.load(Gate_Logits + token_idx * NUM_HEADS + tl.arange(0, NUM_HEADS)).to(tl.float32)
+    # 2*sigmoid so that zero-init gives identity (2 * 0.5 = 1.0)
+    gates = 2.0 * tl.sigmoid(gate_logits)                           # [NUM_HEADS]
+    gates = tl.broadcast_to(gates[:, None], (NUM_HEADS, DIM_HEAD))  # [NUM_HEADS, DIM_HEAD]
+
+    # Load x viewed as [NUM_HEADS, DIM_HEAD]
+    offsets = tl.arange(0, NUM_HEADS)[:, None] * DIM_HEAD + tl.arange(0, DIM_HEAD)[None, :]
+    x = tl.load(X + token_idx * H + offsets).to(tl.float32)
+
+    # Apply per-head gating
+    gated = x * gates  # [NUM_HEADS, DIM_HEAD]
+
+    if QUANTIZE:
+        o_quant, o_scales = _quantize(gated, SCALE_MAX, NUM_BLOCKS, BLOCK_SIZE)
+        out_ptr = Out + token_idx * H + BLOCK_SIZE * tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+        tl.store(out_ptr, o_quant)
+        scales_ptr = Out_scales + token_idx + TMA_ALIGNED_MN * tl.arange(0, NUM_BLOCKS)[:, None]
+        tl.store(scales_ptr, o_scales)
+    else:
+        out_ptr = Out + token_idx * H + offsets
+        tl.store(out_ptr, gated)
+
+
+@triton_op("blockwise::gated_attention", mutates_args=())
+def run_gated_attention(x: torch.Tensor, gate_logits: torch.Tensor, quantize: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+    b, t, h = x.shape
+    num_heads = gate_logits.shape[-1]
+    dim_head = h // num_heads
+    block_size = 128
+    num_blocks = h // block_size
+    num_rows = b * t
+    scale_max: float = 448.0
+
+    if quantize:
+        out = torch.empty((num_rows, h), device=x.device, dtype=torch.float8_e4m3fn)
+        tma_aligned_mn = get_tma_aligned_size(num_rows, 4)
+        scales = torch.empty_strided((num_rows, num_blocks), (1, tma_aligned_mn), dtype=torch.float, device=x.device)
+    else:
+        out = torch.empty((num_rows, h), device=x.device, dtype=torch.bfloat16)
+        tma_aligned_mn = 0
+        scales = None
+
+    wrap_triton(_gated_attention_kernel)[(num_rows,)](
+        x,
+        gate_logits,
+        out,
+        scales,
+        H=h,
+        NUM_HEADS=num_heads,
+        DIM_HEAD=dim_head,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+        TMA_ALIGNED_MN=tma_aligned_mn,
+        QUANTIZE=quantize,
+        SCALE_MAX=scale_max,
+    )
+
+    return out.view(b, t, h), scales
+
+
+@jit
+def _blockwise_dequantize_kernel(
+    X,
+    Scales,
+    Out,
+    scales_row_stride: int,
+    scales_col_stride: int,
+    H: tl.constexpr,
+    NUM_BLOCKS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+
+    # Load per-block scales [NUM_BLOCKS] respecting the actual tensor strides
+    scales = tl.load(
+        Scales + token_idx * scales_row_stride + scales_col_stride * tl.arange(0, NUM_BLOCKS)
+    ).to(tl.float32)
+    scales = tl.broadcast_to(scales[:, None], (NUM_BLOCKS, BLOCK_SIZE))
+
+    # Load fp8 input viewed as [NUM_BLOCKS, BLOCK_SIZE]
+    x_offsets = BLOCK_SIZE * tl.arange(0, NUM_BLOCKS)[:, None] + tl.arange(0, BLOCK_SIZE)[None, :]
+    x = tl.load(X + token_idx * H + x_offsets).to(tl.float32)
+
+    # Dequantize and store as bf16
+    tl.store(Out + token_idx * H + x_offsets, (x * scales).to(tl.bfloat16))
+
+
+@triton_op("blockwise::dequantize", mutates_args=())
+def run_blockwise_dequantize(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+    b, s, h = x.shape
+    block_size = 128
+    num_blocks = h // block_size
+    num_rows = b * s
+    out = torch.empty((num_rows, h), device=x.device, dtype=torch.bfloat16)
+
+    wrap_triton(_blockwise_dequantize_kernel)[(num_rows,)](
+        x,
+        scales,
+        out,
+        scales_row_stride=scales.stride(0),
+        scales_col_stride=scales.stride(1),
+        H=h,
+        NUM_BLOCKS=num_blocks,
+        BLOCK_SIZE=128,
+    )
+
+    return out.view(b, s, h)
+
+
+rowwise_int_quantize_triton = run_quantize_kernel
+blockwise_quantize_adanorm_triton = run_quant_blockwise_norm_tma_aligned
+blockwise_quantize_triton = run_quant_blockwise_tma_aligned
+blockwise_quantize_gelu_triton = run_quant_blockwise_gelu_tma_aligned
+blockwise_quantize_rms_fma_triton = run_quant_blockwise_rms_fma
+rms_fma_triton = run_rms_fma
+gated_attention_triton = run_gated_attention
+blockwise_dequantize_triton = run_blockwise_dequantize
+
+# # Precision-specific convenience functions for Triton kernels
+# def fp8_blockwise_quantize_triton(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+#     return run_quant_blockwise_tma_aligned(x, scale_max=FP8_SCALE_MAX)
+
+# def fp7_blockwise_quantize_triton(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+#     return run_quant_blockwise_tma_aligned(x, scale_max=FP7_SCALE_MAX)
+
+# def fp6_blockwise_quantize_triton(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+#     return run_quant_blockwise_tma_aligned(x, scale_max=FP6_SCALE_MAX)