Kernels:

kernels-community
/

sonic-moe

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Aug 4

Commit

7151f25

verified ·

1 Parent(s): e24502c

Uploaded using `kernel-builder`.

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

build/torch-cuda/_ops.py +8 -2
build/torch-cuda/_ops_compat.py +0 -10
build/torch-cuda/functional/__init__.py +36 -37
build/torch-cuda/functional/backward.py +1 -54
build/torch-cuda/functional/forward.py +1 -50
build/torch-cuda/functional/reduction_over_k_gather.py +6 -6
build/torch-cuda/functional/tile_scheduler.py +0 -91
build/torch-cuda/functional/triton_kernels/__init__.py +4 -1
build/torch-cuda/metadata.json +117 -2
build/torch-cuda/metadata.json.sigstore +1 -0
build/torch-cuda/quack/__init__.py +12 -3
build/torch-cuda/quack/_compile_worker.py +0 -102
build/torch-cuda/quack/_ops_compat.py +9 -3
build/torch-cuda/quack/activation.py +301 -72
build/torch-cuda/quack/autotuner.py +176 -170
build/torch-cuda/quack/bench/__init__.py +0 -0
build/torch-cuda/quack/bench/bench_utils.py +202 -0
build/torch-cuda/quack/blockscaled/__init__.py +31 -0
build/torch-cuda/quack/{mx_utils.py → blockscaled/quantize.py} +0 -0
build/torch-cuda/quack/{blockscaled_gemm_utils.py → blockscaled/utils.py} +185 -79
build/torch-cuda/quack/broadcast_utils.py +3 -4
build/torch-cuda/quack/cache/__init__.py +74 -0
build/torch-cuda/quack/cache/_pool_preload.py +50 -0
build/torch-cuda/quack/cache/async_compile.py +413 -0
build/torch-cuda/quack/{cache_utils.py → cache/jit.py} +194 -57
build/torch-cuda/quack/compile_utils.py +17 -2
build/torch-cuda/quack/complex.py +292 -0
build/torch-cuda/quack/copy_utils.py +613 -87
build/torch-cuda/quack/cross_entropy.py +191 -141
build/torch-cuda/quack/cute_dsl_utils.py +63 -3
build/torch-cuda/quack/dsl/__init__.py +17 -0
build/torch-cuda/quack/{cute_dsl_ptxas.py → dsl/cute_dsl_ptxas.py} +58 -14
build/torch-cuda/quack/dsl/cute_tensor.py +167 -0
build/torch-cuda/quack/dsl/cute_tensor_indexing.py +139 -0
build/torch-cuda/quack/dsl/smem_struct.py +99 -0
build/torch-cuda/quack/dsl/torch_library_op.py +65 -0
build/torch-cuda/quack/epi_composable.py +161 -54
build/torch-cuda/quack/epi_ops.py +640 -187
build/torch-cuda/quack/epi_utils.py +12 -7
build/torch-cuda/quack/fast_math.py +67 -21
build/torch-cuda/quack/gemm.py +72 -25
build/torch-cuda/quack/gemm_act.py +223 -113
build/torch-cuda/quack/gemm_base.py +731 -0
build/torch-cuda/quack/gemm_blockscaled_interface.py +0 -326
build/torch-cuda/quack/gemm_config.py +40 -3
build/torch-cuda/quack/gemm_dact.py +71 -93
build/torch-cuda/quack/gemm_default_epi.py +18 -26
build/torch-cuda/quack/gemm_interface.py +595 -240
build/torch-cuda/quack/gemm_norm_act.py +46 -40
build/torch-cuda/quack/gemm_sm100.py +411 -359

build/torch-cuda/_ops.py CHANGED Viewed

@@ -4,7 +4,13 @@ def get_backend() -> str:
     """Detect the backend by inspecting torch."""
     import torch
-    if hasattr(torch, "neuron"):
         # Needs to be sorted before specific Torch builds, since Neuron
         # extension can be loaded into e.g. CUDA Torch builds.
         return "neuron"
@@ -22,7 +28,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "sonic_moe"
-    unique_id = "86f75d9"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

     """Detect the backend by inspecting torch."""
     import torch
+    if hasattr(torch.backends, "tpu"):
+        # torch_tpu sets torch.backends.tpu when it is imported (via
+        # torch's device-backend autoload), regardless of whether TPU
+        # hardware is present — analogous to torch.version.cuda being
+        # set on CUDA builds without a GPU.
+        return "tpu"
+    elif hasattr(torch, "neuron"):
         # Needs to be sorted before specific Torch builds, since Neuron
         # extension can be loaded into e.g. CUDA Torch builds.
         return "neuron"
 def _find_ops_name() -> str:
     kernel_name = "sonic_moe"
+    unique_id = "83d1d6e"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-cuda/_ops_compat.py DELETED Viewed

@@ -1,10 +0,0 @@
-"""Compatibility helpers for op namespacing in source and built layouts."""
-try:
-    from ._ops import add_op_namespace_prefix as _generated_add_op_namespace_prefix
-except ImportError:
-    def _generated_add_op_namespace_prefix(name: str) -> str:
-        return name if "::" in name else f"sonicmoe::{name}"
-def add_op_namespace_prefix(name: str) -> str:
-    return _generated_add_op_namespace_prefix(name)

build/torch-cuda/functional/__init__.py CHANGED Viewed

@@ -11,13 +11,11 @@ from ..quack.gemm_interface import gemm, gemm_dgated, gemm_gated
 from ..enums import ActivationType, is_glu
 from .backward import (
     _down_projection_backward_act,
-    _down_projection_backward_weight,
     _token_broadcast_backward,
     _topk_softmax_bwd,
     _up_projection_backward_act,
-    _up_projection_backward_weight,
 )
-from .forward import _down_projection_forward, _router_forward, _topk_softmax_fwd, _up_projection_forward
 from .triton_kernels import TC_topk_router_metadata_triton, general_routing_router_metadata_triton
@@ -107,17 +105,21 @@ class _UpProjection(torch.autograd.Function):
             else None
         )
-        _up_projection_forward(
-            x=x,
-            w1=w1,
-            h=h,
-            a=a,
-            b1=b1,
-            expert_frequency_offset=expert_frequency_offset,
-            x_gather_idx=x_gather_idx,
-            activation_type=activation_type.value,
-            is_inference_mode_enabled=is_inference_mode_enabled,
-            concat_layout=concat_layout,
         )
         ctx.T = T
@@ -182,14 +184,15 @@ class _UpProjection(torch.autograd.Function):
             concat_layout=concat_layout,
         )
-        _up_projection_backward_weight(
-            x=x,
-            dw1=dw1,
-            dh=dh,
-            expert_frequency_offset=expert_frequency_offset,
-            x_gather_idx=x_gather_idx,
-            is_glu_activation=is_glu_activation,
-            concat_layout=concat_layout,
         )
         dx_reduced = torch.empty(T, H, dtype=dh.dtype, device=dh.device)
@@ -231,13 +234,7 @@ class _DownProjection(torch.autograd.Function):
         y = torch.empty(TK, H, dtype=a.dtype, device=a.device)
-        _down_projection_forward(
-            w2=w2,
-            a=a,
-            y=y,
-            b2=b2,
-            expert_frequency_offset=expert_frequency_offset,
-        )
         o = torch.empty(T, H, device=a.device, dtype=a.dtype)
         topk_scores = topk_scores.view(-1)
@@ -266,6 +263,7 @@ class _DownProjection(torch.autograd.Function):
             expert_frequency_offset,
             x_gather_idx,
             s_scatter_idx,
         )
         return o
@@ -285,6 +283,7 @@ class _DownProjection(torch.autograd.Function):
             expert_frequency_offset,
             x_gather_idx,
             s_scatter_idx,
         ) = ctx.saved_tensors
         dw2 = torch.empty_like(w2)
@@ -313,12 +312,14 @@ class _DownProjection(torch.autograd.Function):
             activation_type=activation_type.value,
         )
-        _down_projection_backward_weight(
-            dout=dout,
-            a_prime=a_prime,
-            dw2=dw2,
-            expert_frequency_offset=expert_frequency_offset,
-            x_gather_idx=x_gather_idx,
         )
         # TC top-K routing
@@ -369,7 +370,6 @@ def moe_TC_softmax_topk_layer(
     if type(activation_type) == str:
         activation_type = ActivationType(activation_type)
-    assert not torch.compiler.is_compiling()
     assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet"
     a, h = _UpProjection.apply(
@@ -467,7 +467,6 @@ def moe_general_routing_inputs(
         num_activated_expert_per_token_offset,
     )
-    assert not torch.compiler.is_compiling()
     assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet"
     a, h = _UpProjection.apply(

 from ..enums import ActivationType, is_glu
 from .backward import (
     _down_projection_backward_act,
     _token_broadcast_backward,
     _topk_softmax_bwd,
     _up_projection_backward_act,
 )
+from .forward import _router_forward, _topk_softmax_fwd
 from .triton_kernels import TC_topk_router_metadata_triton, general_routing_router_metadata_triton
             else None
         )
+        assert activation_type.value in (
+            "swiglu",
+            "geglu",
+        ), f"QuACK gemm_gated only supports glu activations, got {activation_type.value}"
+        gemm_gated(
+            x,
+            w1.permute(2, 1, 0),
+            activation=activation_type.value,
+            cu_seqlens_m=expert_frequency_offset,
+            A_idx=x_gather_idx,
+            preact_out=h,
+            postact_out=a,
+            store_preact=(not is_inference_mode_enabled),
+            bias=b1,
+            concat_layout=(("B", "bias") if b1 is not None else ("B",)) if concat_layout else None,
         )
         ctx.T = T
             concat_layout=concat_layout,
         )
+        gemm(
+            x.T,
+            dh,
+            out=dw1.permute(2, 1, 0),
+            cu_seqlens_k=expert_frequency_offset,
+            A_idx=x_gather_idx,
+            batch_idx_permute=None,
+            dynamic_scheduler=False,
+            concat_layout=(("out",) if concat_layout else None),
         )
         dx_reduced = torch.empty(T, H, dtype=dh.dtype, device=dh.device)
         y = torch.empty(TK, H, dtype=a.dtype, device=a.device)
+        gemm(a, w2.permute(2, 1, 0), out=y, cu_seqlens_m=expert_frequency_offset, bias=b2)
         o = torch.empty(T, H, device=a.device, dtype=a.dtype)
         topk_scores = topk_scores.view(-1)
             expert_frequency_offset,
             x_gather_idx,
             s_scatter_idx,
+            s_reverse_scatter_idx,
         )
         return o
             expert_frequency_offset,
             x_gather_idx,
             s_scatter_idx,
+            s_reverse_scatter_idx,
         ) = ctx.saved_tensors
         dw2 = torch.empty_like(w2)
             activation_type=activation_type.value,
         )
+        gemm(
+            dout.T,
+            a_prime,
+            out=dw2.permute(2, 0, 1),
+            cu_seqlens_k=expert_frequency_offset,
+            A_idx=x_gather_idx,
+            batch_idx_permute=None,
+            dynamic_scheduler=False,
         )
         # TC top-K routing
     if type(activation_type) == str:
         activation_type = ActivationType(activation_type)
     assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet"
     a, h = _UpProjection.apply(
         num_activated_expert_per_token_offset,
     )
     assert is_glu(activation_type), "QuACK GEMM does not support non GLU activation yet"
     a, h = _UpProjection.apply(

build/torch-cuda/functional/backward.py CHANGED Viewed

@@ -11,7 +11,7 @@ import triton
 import triton.language as tl
 from ..quack.gemm_interface import gemm, gemm_dgated
-from .._ops_compat import add_op_namespace_prefix
 from ..utils import get_powers_of_2
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
@@ -208,35 +208,6 @@ def _up_projection_backward_act(
 _up_projection_backward_act.compile_cache = {}
-@torch.library.custom_op(add_op_namespace_prefix("_up_projection_backward_weight"), mutates_args={"dw1"})
-def _up_projection_backward_weight(
-    x: torch.Tensor,
-    dw1: torch.Tensor,
-    dh: torch.Tensor,
-    expert_frequency_offset: torch.Tensor,
-    x_gather_idx: torch.Tensor,
-    is_glu_activation: bool,
-    concat_layout: bool = False,
-) -> None:
-    I, H, E = dw1.size()
-    if is_glu_activation:
-        I //= 2
-    gemm(
-        x.T,
-        dh,
-        out=dw1.permute(2, 1, 0),
-        cu_seqlens_k=expert_frequency_offset,
-        A_idx=x_gather_idx,
-        batch_idx_permute=None,
-        dynamic_scheduler=False,
-        concat_layout=(("out",) if concat_layout else None),
-    )
-_up_projection_backward_weight.compile_cache = {}
 @torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_act"), mutates_args={"dh", "ds", "db2", "a_prime"})
 def _down_projection_backward_act(
     dout: torch.Tensor,
@@ -272,8 +243,6 @@ def _down_projection_backward_act(
         A_idx=x_gather_idx,
         dynamic_scheduler=False,
     )
-    ds[s_scatter_idx] = ds_scattered
     if db2 is None:
         ds[s_scatter_idx] = ds_scattered
     else:
@@ -314,28 +283,6 @@ def _down_projection_backward_act(
 _down_projection_backward_act.compile_cache = {}
-@torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_weight"), mutates_args={"dw2"})
-def _down_projection_backward_weight(
-    dout: torch.Tensor,
-    a_prime: torch.Tensor,
-    dw2: torch.Tensor,
-    expert_frequency_offset: torch.Tensor,
-    x_gather_idx: torch.Tensor,
-) -> None:
-    gemm(
-        dout.T,
-        a_prime,
-        out=dw2.permute(2, 0, 1),
-        cu_seqlens_k=expert_frequency_offset,
-        A_idx=x_gather_idx,
-        batch_idx_permute=None,
-        dynamic_scheduler=False,
-    )
-_down_projection_backward_weight.compile_cache = {}
 @torch.library.custom_op(add_op_namespace_prefix("_token_broadcast_backward"), mutates_args={"dx_reduced"})
 def _token_broadcast_backward(
     dx_reduced: torch.Tensor,

 import triton.language as tl
 from ..quack.gemm_interface import gemm, gemm_dgated
+from .._ops import add_op_namespace_prefix
 from ..utils import get_powers_of_2
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 _up_projection_backward_act.compile_cache = {}
 @torch.library.custom_op(add_op_namespace_prefix("_down_projection_backward_act"), mutates_args={"dh", "ds", "db2", "a_prime"})
 def _down_projection_backward_act(
     dout: torch.Tensor,
         A_idx=x_gather_idx,
         dynamic_scheduler=False,
     )
     if db2 is None:
         ds[s_scatter_idx] = ds_scattered
     else:
 _down_projection_backward_act.compile_cache = {}
 @torch.library.custom_op(add_op_namespace_prefix("_token_broadcast_backward"), mutates_args={"dx_reduced"})
 def _token_broadcast_backward(
     dx_reduced: torch.Tensor,

build/torch-cuda/functional/forward.py CHANGED Viewed

@@ -9,9 +9,8 @@ import triton
 import triton.language as tl
 from cutlass.cute.runtime import from_dlpack
 from ..quack.cute_dsl_utils import torch2cute_dtype_map
-from ..quack.gemm_interface import gemm, gemm_gated
-from .._ops_compat import add_op_namespace_prefix
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 from .topk import Softmax_Over_TopK, TopK_Over_Softmax
@@ -62,54 +61,6 @@ def _topk_fwd(
 _topk_fwd.compile_cache = {}
-@torch.library.custom_op(add_op_namespace_prefix("_up_projection_forward"), mutates_args={"h", "a"})
-def _up_projection_forward(
-    x: torch.Tensor,
-    w1: torch.Tensor,
-    h: torch.Tensor,
-    a: torch.Tensor,
-    b1: torch.Tensor | None,
-    expert_frequency_offset: torch.Tensor,
-    x_gather_idx: torch.Tensor,
-    activation_type: str,
-    is_inference_mode_enabled: bool = False,
-    concat_layout: bool = False,
-) -> None:
-    assert activation_type in (
-        "swiglu",
-        "geglu",
-    ), f"QuACK gemm_gated only supports glu activations, got {activation_type}"
-    gemm_gated(
-        x,
-        w1.permute(2, 1, 0),
-        activation=activation_type,
-        cu_seqlens_m=expert_frequency_offset,
-        A_idx=x_gather_idx,
-        preact_out=h,
-        postact_out=a,
-        store_preact=(not is_inference_mode_enabled),
-        bias=b1,
-        concat_layout=(("B", "bias") if b1 is not None else ("B",)) if concat_layout else None,
-    )
-_up_projection_forward.compile_cache = {}
-@torch.library.custom_op(add_op_namespace_prefix("_down_projection_forward"), mutates_args={"y"})
-def _down_projection_forward(
-    w2: torch.Tensor,
-    a: torch.Tensor,
-    y: torch.Tensor,
-    b2: torch.Tensor | None,
-    expert_frequency_offset: torch.Tensor,
-) -> None:
-    gemm(a, w2.permute(2, 1, 0), out=y, cu_seqlens_m=expert_frequency_offset, bias=b2)
-_down_projection_forward.compile_cache = {}
 @torch.library.custom_op(add_op_namespace_prefix("_router_forward"), mutates_args={"o"})
 def _router_forward(
     y: torch.Tensor,

 import triton.language as tl
 from cutlass.cute.runtime import from_dlpack
 from ..quack.cute_dsl_utils import torch2cute_dtype_map
+from .._ops import add_op_namespace_prefix
 from .reduction_over_k_gather import token_gather_and_sum_varlen_K_triton
 from .topk import Softmax_Over_TopK, TopK_Over_Softmax
 _topk_fwd.compile_cache = {}
 @torch.library.custom_op(add_op_namespace_prefix("_router_forward"), mutates_args={"o"})
 def _router_forward(
     y: torch.Tensor,

build/torch-cuda/functional/reduction_over_k_gather.py CHANGED Viewed

@@ -71,12 +71,12 @@ def token_gather_sum_kernel(
 ):
     # 1D tiling over T only
     pid_t = tl.program_id(axis=0)
-    t_idx = pid_t.to(tl.uint32)
     # Load segment starts and ends for this token
     if is_varlen_K:
-        Ms = tl.load(M_offset_ptr + t_idx).to(tl.uint32)
-        Me = tl.load(M_offset_ptr + t_idx + 1).to(tl.uint32)
         K_this_token = Me - Ms  # actual K for this token
     else:
         Ms = MAX_K * t_idx
@@ -84,7 +84,7 @@ def token_gather_sum_kernel(
     # Outer loop over H tiles
     for h_tile in tl.static_range(triton.cdiv(H, BLOCK_H)):
-        h_idx = (h_tile * BLOCK_H + tl.arange(0, BLOCK_H)).to(tl.uint32)  # [BLOCK_H]
         m_h = h_idx < H
         # Initialize accumulator for this H tile
@@ -94,7 +94,7 @@ def token_gather_sum_kernel(
         for k_tile in tl.range(tl.cdiv(K_this_token, BLOCK_K)):
             k_offset = k_tile * BLOCK_K
-            k_idx = (k_offset + tl.arange(0, BLOCK_K)).to(tl.uint32)  # [BLOCK_K]
             # Mask for valid K indices
             m_k = k_idx < K_this_token  # [BLOCK_K]
@@ -103,7 +103,7 @@ def token_gather_sum_kernel(
             m_abs = Ms + k_idx  # [BLOCK_K]
             # Gather permuted indices
-            perm_idx = tl.load(M_perm_ptr + m_abs, mask=m_k, other=0).to(tl.uint32)  # [BLOCK_K]
             # Load x values: [BLOCK_K, BLOCK_H]
             x_ptrs = x_ptr + perm_idx[:, None] * stride_xM + h_idx[None, :] * stride_xH

 ):
     # 1D tiling over T only
     pid_t = tl.program_id(axis=0)
+    t_idx = pid_t.to(tl.int64)
     # Load segment starts and ends for this token
     if is_varlen_K:
+        Ms = tl.load(M_offset_ptr + t_idx).to(tl.int64)
+        Me = tl.load(M_offset_ptr + t_idx + 1).to(tl.int64)
         K_this_token = Me - Ms  # actual K for this token
     else:
         Ms = MAX_K * t_idx
     # Outer loop over H tiles
     for h_tile in tl.static_range(triton.cdiv(H, BLOCK_H)):
+        h_idx = (h_tile * BLOCK_H + tl.arange(0, BLOCK_H)).to(tl.int64)  # [BLOCK_H]
         m_h = h_idx < H
         # Initialize accumulator for this H tile
         for k_tile in tl.range(tl.cdiv(K_this_token, BLOCK_K)):
             k_offset = k_tile * BLOCK_K
+            k_idx = (k_offset + tl.arange(0, BLOCK_K)).to(tl.int64)  # [BLOCK_K]
             # Mask for valid K indices
             m_k = k_idx < K_this_token  # [BLOCK_K]
             m_abs = Ms + k_idx  # [BLOCK_K]
             # Gather permuted indices
+            perm_idx = tl.load(M_perm_ptr + m_abs, mask=m_k, other=0).to(tl.int64)  # [BLOCK_K]
             # Load x values: [BLOCK_K, BLOCK_H]
             x_ptrs = x_ptr + perm_idx[:, None] * stride_xM + h_idx[None, :] * stride_xH

build/torch-cuda/functional/tile_scheduler.py DELETED Viewed

@@ -1,91 +0,0 @@
-# ********************************************************************************
-# Copyright (c) 2025, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao
-# ********************************************************************************
-from __future__ import annotations
-import cutlass
-import cutlass.cute as cute
-from cutlass import Boolean, Int32, const_expr
-from ..quack.pipeline import PipelineStateWAdvance
-from ..quack.tile_scheduler import TileScheduler, VarlenMTileScheduler
-class SonicMoETileScheduler(TileScheduler):
-    @staticmethod
-    @cute.jit
-    def create(
-        params: TileScheduler.Params,
-        tile_count: cute.Tensor | None = None,
-        scheduler_pipeline: cutlass.pipeline.PipelineAsync | None = None,
-        is_scheduler_warp: bool | Boolean = False,
-        *,
-        loc=None,
-        ip=None,
-    ) -> SonicMoETileScheduler:
-        """is_scheduler_warp should only be true for one warp in the whole cluster"""
-        stages = 0
-        if const_expr(not params.is_persistent):
-            cidx, cidy, _ = cute.arch.cluster_idx()
-            cdimx, _, _ = cute.arch.cluster_dim()
-            cluster_id = cidx + cidy * cdimx
-            current_work_linear_idx = Int32(cluster_id)
-        else:
-            _, _, bidz = cute.arch.block_idx()
-            current_work_linear_idx = Int32(bidz)
-            if const_expr(params.tile_count_semaphore is not None):
-                assert tile_count is not None
-                assert scheduler_pipeline is not None
-                stages = const_expr(cute.size(tile_count))
-        return SonicMoETileScheduler(
-            current_work_linear_idx,
-            Int32(0),  # num_tiles_executed
-            tile_count,
-            scheduler_pipeline,
-            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)),
-            params,
-            loc=loc,
-            ip=ip,
-        )
-    def prefetch_next_work(self, *, advance_count: int = 1, loc=None, ip=None):
-        old_current_work_linear_idx = self._current_work_linear_idx
-        if const_expr(self.params.is_persistent):
-            num_persistent_clusters = cute.arch.grid_dim()[2]
-            self._current_work_linear_idx += advance_count * Int32(num_persistent_clusters)
-        future_tile_coord_mnkl = self.get_current_work()
-        self._current_work_linear_idx = old_current_work_linear_idx
-        return future_tile_coord_mnkl
-class SonicMoEVarlenMTileScheduler(VarlenMTileScheduler, SonicMoETileScheduler):
-    @staticmethod
-    @cute.jit
-    def create(
-        params: VarlenMTileScheduler.Params,
-        tile_count: cute.Tensor | None = None,
-        scheduler_pipeline: cutlass.pipeline.PipelineAsync | None = None,
-        is_scheduler_warp: bool | Boolean = False,
-        *,
-        loc=None,
-        ip=None,
-    ) -> SonicMoEVarlenMTileScheduler:
-        stages = 0
-        _, _, bidz = cute.arch.block_idx()
-        current_work_linear_idx = Int32(bidz)
-        if const_expr(params.tile_count_semaphore is not None):
-            assert tile_count is not None
-            assert scheduler_pipeline is not None
-            stages = const_expr(cute.size(tile_count))
-        return SonicMoEVarlenMTileScheduler(
-            current_work_linear_idx,
-            Int32(0),  # num_tiles_executed
-            Int32(0),  # current_batch_idx
-            Int32(0),  # num_work_idx_before_cur_batch
-            tile_count,
-            scheduler_pipeline,
-            PipelineStateWAdvance(stages, Int32(0), Int32(0), Int32(1 if is_scheduler_warp else 0)),
-            params,
-            loc=loc,
-            ip=ip,
-        )

build/torch-cuda/functional/triton_kernels/__init__.py CHANGED Viewed

@@ -1,10 +1,13 @@
 import math
 import torch
 import triton
 import triton.language as tl
-from ..._ops_compat import add_op_namespace_prefix
 from .bitmatrix import _bitmatrix_metadata_compute_stage1, _bitmatrix_metadata_compute_stage2, _keyed_add

+# ********************************************************************************
+# Copyright (c) 2026, Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao
+# ********************************************************************************
 import math
 import torch
+from ..._ops import add_op_namespace_prefix
 import triton
 import triton.language as tl
 from .bitmatrix import _bitmatrix_metadata_compute_stage1, _bitmatrix_metadata_compute_stage2, _keyed_add

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,7 +1,7 @@
 {
   "name": "sonic-moe",
-  "id": "_sonic_moe_cuda_86f75d9",
-  "version": 1,
   "license": "Apache-2.0",
   "python-depends": [
     "tvm-ffi",
@@ -9,5 +9,120 @@
   ],
   "backend": {
     "type": "cuda"
   }
 }

 {
   "name": "sonic-moe",
+  "id": "_sonic_moe_cuda_83d1d6e",
+  "version": 2,
   "license": "Apache-2.0",
   "python-depends": [
     "tvm-ffi",
   ],
   "backend": {
     "type": "cuda"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "DDepLW9l03NRbdlAuR2+UkWgJuQN/8ymgKZTMmdV0Dc=",
+      "_ops.py": "9EVL/sWTDKjpk2E9q9TV3k0BIZR8kzK88mQRF7lwkm8=",
+      "enums.py": "eeBuPJFCiiGCsL7OmYWCMfdvj4iQE4irrXun8niAUgg=",
+      "functional/__init__.py": "FDQaNUpQGY+3YeKXK9PNmtwKjXtn16GxsGdyFiQul08=",
+      "functional/backward.py": "W10eywFlVh1mMO+9oktuSWK7HllWZxzaByzQfXt1jM4=",
+      "functional/forward.py": "Tq1hGu14TvCXMptVrIZJQLgnAp0wSVzATEv9T5nc35w=",
+      "functional/reduction_over_k_gather.py": "fBYTaHbHTXpl7i2dtJ1InNIh0PIpxfE+6ImvPoNcOaU=",
+      "functional/topk.py": "Gsq7DJpeAuL0co/Zm9tSpyP0t9tSEwxSlRyjrTyCVhg=",
+      "functional/triton_kernels/__init__.py": "19C1f3DyiRGJfW3TcfHZH9x4q2vhqzAj2YkK/if9TUU=",
+      "functional/triton_kernels/bitmatrix.py": "kJ+VIr07n0e48+XQHSgQbLnmcRgFZ91M3NIUxNN4hwI=",
+      "jit.py": "pko7Lkttkab+0YDUkXbAAN3I/yMvTnw7d2jJE6zml+g=",
+      "moe.py": "7UtJ7Oe+401gE+pn0axai3MWrVipfJl9pJc72mSwB/g=",
+      "quack/__init__.py": "BsCyHg2lDdBg8lhuijV2CjN2CFJnQQuovcg5v4gM+Xo=",
+      "quack/_ops_compat.py": "Df2expY3Aqaob3ZuykC0qsSWm4DuUUSPUvI9J9gY9Rg=",
+      "quack/activation.py": "YIsLHpQauAokG2r/V9ToZnFzkCDll3JKU4IhTCsK8Hw=",
+      "quack/autotuner.py": "nbdngdlWYF19B4R6z6N7fgnDoO40MFghpeHlcKe2Yv8=",
+      "quack/bench/__init__.py": "47DEQpj8HBSa+/TImW+5JCeuQeRkm5NMpJWZG3hSuFU=",
+      "quack/bench/bench_utils.py": "OZfHYnYliR7+o1paEdakdQ8YS22ySnOwPjwxuwfXpxE=",
+      "quack/blockscaled/__init__.py": "8aZJFIdaeT2ND6x77yrR/kAFVlSAkjx/nEwgPXrSv5s=",
+      "quack/blockscaled/quantize.py": "EPiyWvbHg1PC0UI5IRX3CqHTNk6HnjzhkUY+vkwgNSU=",
+      "quack/blockscaled/utils.py": "IE/QeL4IurRCCLLF/m/PKrZjuJ/Tp1j5jcgZggxy2eM=",
+      "quack/broadcast_utils.py": "HVUVszpqpbhv3ZrfCMSQVtqmjfRJJljYH+3J0WuSaMs=",
+      "quack/cache/__init__.py": "9qjakMy4yFmpI7XdQz4NDUyqfCti6xi0tkZFPnp9KVU=",
+      "quack/cache/_pool_preload.py": "qGKHxfFndDBCYyHUjSRDC7A+conMOQTqsjPfX5OCZc0=",
+      "quack/cache/async_compile.py": "BzJjQ5YswF4YAreXFDvbI05pQHzdtE5zxSdN4M7pOp4=",
+      "quack/cache/jit.py": "vXgLNNeAMrEY8a1fpigiiCUWZGCTbrtxMYWFrHIffkE=",
+      "quack/compile_utils.py": "d8AzPa2ONwKF3DhsB1E5orFzrNmJJinrYLJ0kstFq6k=",
+      "quack/complex.py": "WmYNde9V3ZDLaOt7ttCnjH4Dxahan4+PvLfoHHHU+jo=",
+      "quack/copy_utils.py": "0e1AtA1GZjFdc9Ts0mNT+gCftK0PVbDDkzOkXiBgSHE=",
+      "quack/cross_entropy.py": "OL8SJ435IOxTxL7L/C6vl/+3+BQWbIgMUt86rM8hJ84=",
+      "quack/cute_dsl_utils.py": "9lUc4T4L+Nm05V8doN+iiUqBFa2yBUSjYaVUseXCfH4=",
+      "quack/dsl/__init__.py": "v0w1PKdAXyF8yOzc/OwvR5+4M5rcdni62pCtVWpj2bI=",
+      "quack/dsl/cute_dsl_ptxas.py": "jbam6fUwzm4X5CZc29LevlWXBAzbedWu3j6hK/2lEmA=",
+      "quack/dsl/cute_tensor.py": "l1Z8XNzBSQuHwRKR6gX5+Se1DR39NoVR/JEV2IJbveo=",
+      "quack/dsl/cute_tensor_indexing.py": "1QnPrnekkECXkgHC5NWY6JxyDmTTf5kgClCIYwxMfG0=",
+      "quack/dsl/smem_struct.py": "piMjPM+wy8w/Ig6IZI/AOoVpeKsKSEes7xPxY2eiWCI=",
+      "quack/dsl/torch_library_op.py": "QfwtFO4P/3FSivWuY1LH2dL2xiead+m2VpK3iBKjKp0=",
+      "quack/epi_composable.py": "iCA8mYu9CLXNVvWmdP7847Pk+51tUYv6U6lTb5ptgfQ=",
+      "quack/epi_ops.py": "5+MCkOUqeLxFrzevU5OH4KXgxJJTKC+10jUHkxiDplA=",
+      "quack/epi_utils.py": "pdXiffZgdAbSA2WM5jkOIWvrPTBGqsvIkKrIffiAXFM=",
+      "quack/fast_math.py": "Nvhr8+bgjPAS+e/XKzlxO0Xzqa3gi8vRxV3gl82ECJI=",
+      "quack/gemm.py": "mAVLwuD+tC7O+NGhLzGaYQhsoWm4lo8h7TWEsqyJosA=",
+      "quack/gemm_act.py": "OIIB130+VHkMvyfP/kXCXzeNr7BkK68pIhxY2YwTqXM=",
+      "quack/gemm_base.py": "9RUR26bnwPd5LkM0PDo5201x0QdxgLO6VRZNCxBuDR8=",
+      "quack/gemm_config.py": "NinZfsUbbcaFyMR+yzuBtSfF2Q6TJPd3AoxAr89I5AQ=",
+      "quack/gemm_dact.py": "HV3dHIx4Er+QogUnp/lvBbkOAsXm2+z2YRiuOHyO5gQ=",
+      "quack/gemm_default_epi.py": "YFh3eGsC5wPVHiw//IoNQKiDH0u1ioT/5QPwtasRats=",
+      "quack/gemm_interface.py": "XTqyTBBZxOFEnE9/DCJjAsUQ3BBICY9nFgOSGjc8OZ4=",
+      "quack/gemm_norm_act.py": "u5CVGXIN4g1Z6pPD4fHKQkjDl1V/hdH1S3eWBlq8Hf0=",
+      "quack/gemm_sm100.py": "6vehfja8xPV1A78HBC6kUjHfOkUGGVtTiTLW7YVlr48=",
+      "quack/gemm_sm120.py": "aCy6GOMLgyvuFq+L4mRPXMXuKVcobswO6FSwLWhfx+k=",
+      "quack/gemm_sm80.py": "mGyJSOf7HSTJbOQawSJiu53hfKNg7HvSPzTVDLSwB4c=",
+      "quack/gemm_sm90.py": "38ONe3cM8g8hAEx4JomzbbkJ5lTo92UpbLNoUMn9EQE=",
+      "quack/gemm_sq_reduce.py": "N6JefW8NwxiMjFs6JssVu97Ysn9n4nErQkW2J6VMZHI=",
+      "quack/gemm_symmetric.py": "PRVtXM+nCya7nyiDH5ozvndGnfLX5TkM9kXYnSuJ5sg=",
+      "quack/gemm_tvm_ffi_utils.py": "76C1usihWxUD0lRi82AUO8OVjWVW5ypxIRelaFSAzTo=",
+      "quack/jax_utils.py": "z9CKuM2HbYXkgib8+mVNEvI5WF2HxBl8cgTCmiosTus=",
+      "quack/layout_utils.py": "I6MaaKCC67T+Yh/d6QdU4xbiqfPQ0ZJRMJ7/eUGFlFg=",
+      "quack/linear.py": "DK6ZcxbVB5PNnvRhUilds7BbRRdqovwNXJTPUOfUo2A=",
+      "quack/linear_cross_entropy.py": "UmkMdBB+8Tp4C4ikwjM/AGG/Ql+lWoqgmbADEHZ4I6o=",
+      "quack/mlp.py": "LKrK0VX3MToBM98oELVUMQJ3Zmgt93QNyZ0u7xdaFNE=",
+      "quack/nvmmh_heuristic.py": "mxImbqtdhY4NxYS6kZRL41zgG8oKa94uo5yLUtbLs2E=",
+      "quack/pipeline.py": "veqXF6OrbggTrcjhw4E/vDc7ufja27tifRDSuok4XHM=",
+      "quack/reduce.py": "Ul6PlgNnYQvm4dq948stCPrBdKA0Ct+jjgSeZFbxdNk=",
+      "quack/reduction_base.py": "cJc9fa+DsOFigD7d7bgNWR03gXkyXNzoSXwjPJxKTl0=",
+      "quack/rms_final_reduce.py": "9X2rSFH25BjNdLFDaJKmZeq/i1/gnmYhoYtCyPmAhrQ=",
+      "quack/rmsnorm.py": "yFyoI/p3ZyZqvOaDgN/+a3PbK1kJZvnYxjgLxAorgN4=",
+      "quack/rmsnorm_config.py": "HDOdXv5Ej2egQVQDVk4k1ynXUE0D0HpPaN9O7Qd9EK8=",
+      "quack/rotary.py": "A0FaRr7LKdWw+kjeHEg79I4eYb9+MbGKSMPuawocGBw=",
+      "quack/rounding.py": "jI6fP/qCjSMh9s8PtjPDYfTPs200C3xx9zKll1hlodU=",
+      "quack/sm100_utils.py": "+AkAPWzg+2cXZMwTijIVTvWWwZmk5m3bAsOJQTffzDo=",
+      "quack/sm80_utils.py": "0yOgRslAIq8B7MS8tIkgZuilOMIkYnJDUl7c2gjdkzg=",
+      "quack/sm90_utils.py": "gX9BE07yDSBMnrFdIJY0EXvyQuutTiXlCT8j5mHDLh0=",
+      "quack/softmax.py": "YQ4R6wVI2J/+FsjuTVGQ2Jj4shl4KU/YrHgFZmQ99do=",
+      "quack/softmax_jax.py": "VZ1phS40aMQDyx4QgsrWPxbkbdLnv3jNEc/+j2VqNN0=",
+      "quack/sort/bitonic_sort.py": "XjGKnZ51NI2N2BoHBgxz2Ne2OvYfnuOTgiOSnyjN08Q=",
+      "quack/sort/generate_sorting_networks.py": "al5BepGSLhn/bBPg64p2KpW7itx9gaCIGxhpqIceKB8=",
+      "quack/sort/sorting_networks.py": "okUmplKNQMAu8/lHClRiQf31X3tIe7MR07KybQJfH8Y=",
+      "quack/sort/utils.py": "lifuGThMk4D1VxfAUYDxrrUzHwxLGgLdoinpE8UVN3g=",
+      "quack/spec/__init__.py": "gZsivx8yPk0Tk7GqzHpOtrjN2b3dSB0/MNKVt0/e8qs=",
+      "quack/spec/mma.py": "3SSQiMkkncb3vuaoc9Dkr687vIBxEf7icnsdWJv5q58=",
+      "quack/spec/smem.py": "McZhRq5NJlHEFye+TQZpvDvnarRgLFKh2Hk6JslyOUA=",
+      "quack/spec/tensor_spec.py": "2d30W/TI7kwvc7tTw6/q1JU35/k0tJGD81JAXeX/P2A=",
+      "quack/spec/tma.py": "sz94fm7k50nGQ2jr7ubCzeZBX/YWZRmIerwp8P4+EFs=",
+      "quack/spec/tmem.py": "7YVy9F+Xn3iYEx2aDKiVW1xSIO8GMw+nyrwIdFgnrxs=",
+      "quack/sync/__init__.py": "5dkW0RJn9GodDYJeXxyQJfNaMBkXczp5BjkXeQK4/yg=",
+      "quack/sync/barrier.py": "sTHk/GcIC+OcgF6SK+YxSo7R4mdCB1n/6i42fdx+iVg=",
+      "quack/tensormap_manager.py": "twLLdCS6s8+BCYIh4XYdPtE3XX4OAASHklnsGssznPg=",
+      "quack/testing/__init__.py": "F4fOLg+7e5MRZUU+2kKdpR+XykrUhy3hrkrXke139+0=",
+      "quack/testing/pytest_plugin.py": "drWmLh35ffAa9J2VjGOiyhfWYBOcyr54HRNb1UkGTHo=",
+      "quack/testing/trace.py": "/pVNHOyiv1mgnvGKWy7RmreheJzUZx229+fMJoQZJo0=",
+      "quack/tile_scheduler.py": "VftFntcYMwF3t6op2H4arOSjmcQzvSqBp5YXjDSS5a0=",
+      "quack/topk.py": "md6UD7+W1OlZ3K0UDwhfT5fIT3kBrL3ccPYDuEcIedM=",
+      "quack/trace.py": "nmHaefNLu7k9MjI+LPYmGe69WU1UumN6D7IDAtX9sms=",
+      "quack/transform/__init__.py": "JhfsXbnxVdojWTz6VbGHCTkCANeR1b9Js2wDyldnO+c=",
+      "quack/transform/hadamard.py": "fehxPxaWlhWr1LHgBJijuFHQ2Q9houh/ZLEeO1WOvtQ=",
+      "quack/utils.py": "/FsIscjaxmOSj/g6frh4gLgB6iuIZ6e3EeGNunKWBuw=",
+      "quack/varlen_utils.py": "UtdJmeCWbAMeJhTYNF78cjN9CH3fEpNQvicoF1YhGL8=",
+      "utils.py": "mggbSpSYFQIJTyYgccngut/SeK2AqJfSXL1/CMZ8VyM="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "a6564d1f481adcbd3273099c0f4432e7b833e846",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "83d1d6e670b94603c6bdda6f039416ce988d7683",
+      "dirty": false
+    }
   }
 }

build/torch-cuda/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSjCCBtCgAwIBAgIUFiUTErfHXgcT0avZDUrtqVpjPVMwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwODA0MDgyNzA5WhcNMjYwODA0MDgzNzA5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE8E1Q2yBXTAmi4adPxCkrbQcPf17/xBZ3tfE6qg9RHWEdmi7WM5D7S43lqmnObiB4CkefF+OIxcSlHV8UJ2opOKOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU36UElV8bsJe0ibWvj7caucWF9xgwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoODNkMWQ2ZTY3MGI5NDYwM2M2YmRkYTZmMDM5NDE2Y2U5ODhkNzY4MzATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoODNkMWQ2ZTY3MGI5NDYwM2M2YmRkYTZmMDM5NDE2Y2U5ODhkNzY4MzAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoODNkMWQ2ZTY3MGI5NDYwM2M2YmRkYTZmMDM5NDE2Y2U5ODhkNzY4MzAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDgzZDFkNmU2NzBiOTQ2MDNjNmJkZGE2ZjAzOTQxNmNlOTg4ZDc2ODMwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzA4OTE2ODY1NTMvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABn8viPTkAAAQDAEYwRAIgfLsNELCVlyuHUG5H2cJ2naaV5F0SYQZv0YdXyC9K6asCIGpaO2l3hSUXIiLRBq6dsIXZHZyZqZ2/6nNPdklmZMlzMAoGCCqGSM49BAMDA2gAMGUCMHeyUvclqJRFc6yuhmZtURySzkJ6RpnXasTQ5hXcx4H/3BVH1hHtFeXi4epa0GNcLgIxAPspbTGLLNMVY+7Rhv7SaSsX5AGzaWVABXmQPDy7S7T4IJArASfsqBm3NmO0SR0vtA=="}, "tlogEntries":[{"logIndex":"2339551903", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1785832029", "inclusionPromise":{"signedEntryTimestamp":"MEYCIQDitpUL1h/KusQ9OMEtNPCsHo084FfnQmZK0WwFpKnoZAIhAPTvy8x20w6lcWucTqcZVLYevLSY3i2S/rLsZ34hFk44"}, "inclusionProof":{"logIndex":"2217647641", "rootHash":"nQRMpxeZMxZKoBM0Aid5ajFNCBrokJCPxAZe9n5BiME=", "treeSize":"2217647645", "hashes":["48NZ8qo8UwRJEAZ3CLfS1KfxJZxkxn+iOuNafjEb1/M=", "u2Qapxkuy0/IKrnn3gJDePaVd9YyBiTAyigk5uwxM8Y=", "yzLVy2ZAAv+TndVc45A90uQ1k7RULtWa9nqp7rNIt9U=", "ZeI1SGr9GP/SskvVce8yHk+6DRrxLrWqJpQnRQRaKpc=", "nQWu7U6/pF5XpZlDhkQlNjCedXmCTB3GYqdStVFhcGo=", "WeMerEQh8gnrs+xRVnIEQXfykDf4KiXgRHX8oVHQYIs=", "2uIbGD5iEQJHWVjTX0311tWS7CFZVaKdy3hZGbyhI8s=", "DLtcvkp4dRTrpmSEDyxo9Y5x1mCGHjNNX5a+Pu2gajA=", "7VIK52c6IvZNnE+Au/wuGXo7lX61HQyDCibo+Dfcvs8=", "jDU3DoSktlu0G+qce/KFrCSxbMWDf2MyMt8j1NVTV/g=", "p0bIaoDFhBhr45Lurz3ap+CC2ZtzymaJJb6Ij0aF6Xg=", "7xWj0qvDr9hJ4vooBulXteOWWqidK4mqZKahpqVHaoE=", "r3u4sx0tkA1pkDYoU+5zLDv1az999M9yZP3h7rXTwYE=", "8cBkID7yRXftjQqclK9Jj0hUDPpp7HMqoKXxzdkwOzg=", "i5Zl8FZrDwxCDv2e2DNO2M8JvpR/c11ElvCZS53/teA=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2217647645\nnQRMpxeZMxZKoBM0Aid5ajFNCBrokJCPxAZe9n5BiME=\n\n— rekor.sigstore.dev wNI9ajBFAiEAm8BuxbwwhrjVkIerX6mU3KfouLswOpqLlyxRei9yV58CIEW4vOQQYXUx761ALJD54e5Q5qfizDBuosCfSV/l8+Jc\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI0ZWYzZTQwM2RjMGJlM2Q1N2YzNWYyYjNkNzIwNTMwNDZkZmZlYjczZWQ2YmE3ODE4NGMwYjA0Njc5MzllNzU2In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJSElhM3JzWDI4a1ZOKzFPVDQyUDhLNk82anY3Wk9iSW9neXg1M2t5a0xFWkFpRUF5K1QxY0VyTjJ3M2xhRWpwaXRZYi9JNi9KMEhlOFNiQVN6UjBVVFdEYk84PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRha05EUW5SRFowRjNTVUpCWjBsVlJtbFZWRVZ5WmtoWVoyTlVNR0YyV2tSVmNuUnhWbkJxVUZaTmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlFUVRCTlJHZDVUbnBCTlZkb1kwNU5hbGwzVDBSQk1FMUVaM3BPZWtFMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVU0UlRGUk1ubENXRlJCYldrMFlXUlFlRU5yY21KUlkxQm1NVGN2ZUVKYU0zUm1SVFlLY1djNVVraFhSV1J0YVRkWFRUVkVOMU0wTTJ4eGJXNVBZbWxDTkVOclpXWkdLMDlKZUdOVGJFaFdPRlZLTW05d1QwdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlV6TmxWRkNteFdPR0p6U21Vd2FXSlhkbW8zWTJGMVkxZEdPWGhuZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVBSRTVyVFZkUk1scFVXVE5OUjBrMVRrUlpkMDB5VFRKWmJWSnJDbGxVV20xTlJFMDFUa1JGTWxreVZUVlBSR2hyVG5wWk5FMTZRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDlFVG10TlYxRXlXbFJaTTAxSFNUVk9SRmwzVFRKTk1sbHRVbXRaVkZwdFRVUk5OVTVFUlRJS1dUSlZOVTlFYUd0T2VsazBUWHBCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDlFVG10TlYxRXlXbFJaTTAxSFNUVk9SRmwzVFRKTk1sbHRVbXNLV1ZSYWJVMUVUVFZPUkVVeVdUSlZOVTlFYUd0T2VsazBUWHBCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUm5lbHBFUm1zS1RtMVZNazU2UW1sUFZGRXlUVVJPYWs1dFNtdGFSMFV5V21wQmVrOVVVWGhPYlU1c1QxUm5ORnBFWXpKUFJFMTNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWtFMFQxUkZNazlFV1RGT1ZFMTJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNDRkbWxRVkd0QlFVRlJSQXBCUlZsM1VrRkpaMlpNYzA1RlRFTldiSGwxU0ZWSE5VZ3lZMG95Ym1GaFZqVkdNRk5aVVZwMk1GbGtXSGxET1VzMllYTkRTVWR3WVU4eWJETm9VMVZZQ2tscFRGSkNjVFprYzBsWVdraGFlVnB4V2pJdk5tNU9VR1JyYkcxYVRXeDZUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlaMEZOUjFWRFRVaGxlVlYyWTJ3S2NVcFNSbU0yZVhWb2JWcDBWVko1VTNwclNqWlNjRzVZWVhOVVVUVm9XR040TkVndk0wSldTREZvU0hSR1pWaHBOR1Z3WVRCSFRtTk1aMGw0UVZCemNBcGlWRWRNVEU1TlZsa3JOMUpvZGpkVFlWTnpXRFZCUjNwaFYxWkJRbGh0VVZCRWVUZFROMVEwU1VwQmNrRlRabk54UW0welRtMVBNRk5TTUhaMFFUMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgx31mO/vQ8UTM6GrWl7ysWYCg9VM1q5BDPZqvFVfIaYMCFQD1qIDvOHE1Cgo864r6GhDv1PSMORgPMjAyNjA4MDQwODI3MDlaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDgwNDA4MjcwOVowLwYJKoZIhvcNAQkEMSIEIHmfCoWI9C2TufxbIehLRJjjeOkcBfMZEcgFW3CNowVHMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMHTZ3i8JqZ77m5Vccy3iUVKabZVryj/JRhzGaPfzud9V06kWLkG/Y8pmaHEOZXR4tAIxAOK0xDAb+/Kj7Vn16OaMNEQGjedN9J3fDf2Nyf+/TvnAlwVVpbX++c717VMpDOG6pw=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"TvPkA9wL49V/NfKz1yBTBG3/63Pta6eBhMCwRnk551Y="}, "signature":"MEUCIHIa3rsX28kVN+1OT42P8K6O6jv7ZObIogyx53kykLEZAiEAy+T1cErN2w3laEjpitYb/I6/J0He8SbASzR0UTWDbO8="}}

build/torch-cuda/quack/__init__.py CHANGED Viewed

@@ -1,8 +1,17 @@
-__version__ = "0.3.11"
 import os
 if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None:
-    from . import cute_dsl_ptxas  # noqa: F401
-    cute_dsl_ptxas.patch()

+__version__ = "0.6.1"
 import os
+from . import dsl as _quack_dsl  # noqa: F401
 if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None:
+    from .dsl import cute_dsl_ptxas as _cute_dsl_ptxas
+    # Patch before importing any modules that instantiate CuTeDSL. The patch
+    # forces PTX dumping so the CUDA library loader can replace CUTLASS DSL's
+    # embedded ptxas-library cubin with one assembled by system ptxas.
+    _cute_dsl_ptxas.patch()
+# Pythonic CuTe tensor indexing (`:` / `...` sugar) is installed as a side effect
+# of importing `quack.dsl`, which imports `quack.dsl.cute_tensor_indexing` and
+# monkey-patches CuTe's tensor classes process-wide.

build/torch-cuda/quack/_compile_worker.py DELETED Viewed

@@ -1,102 +0,0 @@
-# Copyright (c) 2025, Tri Dao.
-# Persistent subprocess worker for parallel autotuning pre-compilation.
-# Receives length-prefixed pickled tasks on stdin, creates FakeTensors
-# matching the parent's tensor metadata, and compiles with COMPILE_ONLY=True.
-# Stays alive to process multiple configs (amortizes import overhead).
-import importlib
-import pickle
-import struct
-import sys
-import torch
-from torch._subclasses.fake_tensor import FakeTensorMode
-from . import cache_utils
-cache_utils.COMPILE_ONLY = True
-_dtype_map = {
-    "torch.float16": torch.float16,
-    "torch.bfloat16": torch.bfloat16,
-    "torch.float32": torch.float32,
-    "torch.float64": torch.float64,
-    "torch.int32": torch.int32,
-    "torch.int64": torch.int64,
-    "torch.int8": torch.int8,
-    "torch.uint8": torch.uint8,
-    "torch.bool": torch.bool,
-}
-def _make_fake_tensor(meta):
-    shape = meta["shape"]
-    stride = meta["stride"]
-    dtype = _dtype_map[meta["dtype"]]
-    return torch.empty_strided(shape, stride, dtype=dtype, device="cuda")
-def _recv(stream):
-    """Read a length-prefixed pickled message. Returns None on EOF."""
-    header = stream.read(4)
-    if len(header) < 4:
-        return None
-    length = struct.unpack("<I", header)[0]
-    if length == 0:
-        return None
-    data = stream.read(length)
-    return pickle.loads(data)
-def _send(stream, msg):
-    """Write a length-prefixed pickled message."""
-    data = pickle.dumps(msg)
-    stream.write(struct.pack("<I", len(data)))
-    stream.write(data)
-    stream.flush()
-def main():
-    stdin = sys.stdin.buffer
-    stdout = sys.stdout.buffer
-    # Signal ready
-    _send(stdout, "READY")
-    fn_cache = {}
-    while True:
-        payload = _recv(stdin)
-        if payload is None:
-            break
-        fn_module = payload["fn_module"]
-        fn_qualname = payload["fn_qualname"]
-        fn_key = (fn_module, fn_qualname)
-        if fn_key not in fn_cache:
-            mod = importlib.import_module(fn_module)
-            obj = mod
-            for part in fn_qualname.split("."):
-                obj = getattr(obj, part)
-            fn_cache[fn_key] = getattr(obj, "fn", obj)
-        fn = fn_cache[fn_key]
-        tensor_meta = payload["tensor_meta"]
-        kwargs = payload["kwargs"]
-        config_kwargs = payload["config_kwargs"]
-        with FakeTensorMode():
-            fake_args = []
-            for meta in tensor_meta:
-                if isinstance(meta, dict) and "shape" in meta:
-                    fake_args.append(_make_fake_tensor(meta))
-                else:
-                    fake_args.append(meta)
-            try:
-                fn(*fake_args, **kwargs, **config_kwargs)
-                _send(stdout, "OK")
-            except Exception as e:
-                _send(stdout, f"ERR:{e}")
-if __name__ == "__main__":
-    main()

build/torch-cuda/quack/_ops_compat.py CHANGED Viewed

@@ -1,4 +1,10 @@
-from .._ops_compat import add_op_namespace_prefix
-def add_quack_op_namespace_prefix(name: str) -> str:
-    return add_op_namespace_prefix(f"quack__{name}")

+from .._ops import add_op_namespace_prefix as _add_op_namespace_prefix
+# For quack we need to prefix the function name because some names
+# overlap between quack and sonic-moe itself. Name the function the
+# same as the function it is wrapping for the prefix check to be
+# happy.
+def add_op_namespace_prefix(name: str) -> str:
+    return _add_op_namespace_prefix(f"quack__{name}")

build/torch-cuda/quack/activation.py CHANGED Viewed

@@ -4,10 +4,12 @@ import math
 from typing import Tuple
 from functools import partial
 import cutlass.cute as cute
 from cutlass import Float32, Boolean, const_expr
-from cutlass.cutlass_dsl import T, dsl_user_op
-from cutlass._mlir.dialects import llvm, nvvm
 F32_or_F32x2 = Float32 | Tuple[Float32, Float32]
@@ -21,30 +23,61 @@ sub_packed_f32x2 = partial(
 @dsl_user_op
-def tanh(a: float | Float32, *, loc=None, ip=None) -> Float32:
-    return Float32(
-        llvm.inline_asm(
-            T.f32(),
-            [Float32(a).ir_value(loc=loc, ip=ip)],
-            "tanh.approx.f32 $0, $1;",
-            "=f,f",
-            has_side_effects=False,
-            is_align_stack=False,
         )
-    )
 @dsl_user_op
-def sigmoid(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
         # return 0.5 + 0.5 * cute.math.tanh(0.5 * x, fastmath=True)
         return 0.5 + 0.5 * tanh(0.5 * x)
     else:
         x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
-        tanh_x_half = (tanh(x_half[0]), tanh(x_half[1]))
         return cute.arch.fma_packed_f32x2(tanh_x_half, (0.5, 0.5), (0.5, 0.5))
 @dsl_user_op
 def dsigmoid_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> Float32:
     # return dout * out * (1.0 - out)
@@ -54,9 +87,9 @@ def dsigmoid_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> F
 @dsl_user_op
 def relu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
-        return cute.arch.fmax(x, Float32(0.0))
     else:
-        return cute.arch.fmax(x[0], Float32(0.0)), cute.arch.fmax(x[1], Float32(0.0))
 @dsl_user_op
@@ -66,7 +99,7 @@ def drelu(
 ) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
     if const_expr(not isinstance(x, tuple)):
         x_pos = Boolean(x > 0)
-        return dout if x_pos else Float32(0.0), cute.arch.fmax(x, Float32(0.0))
     else:
         x0_pos = Boolean(x[0] > 0)
         x1_pos = Boolean(x[1] > 0)
@@ -77,9 +110,9 @@ def drelu(
 @dsl_user_op
 def relu_sq(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
-        return cute.arch.fmax(x, Float32(0.0)) * x
     else:
-        relu_x = (cute.arch.fmax(x[0], Float32(0.0)), cute.arch.fmax(x[1], Float32(0.0)))
         return cute.arch.mul_packed_f32x2(relu_x, x)
@@ -117,19 +150,17 @@ def gelu_tanh_approx(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     sqrt_2_over_pi = math.sqrt(2 / math.pi)  # ~0.797885
     sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi  # ~0.0356774
     if const_expr(not isinstance(x, tuple)):
-        return 0.5 * (
-            x
-            # Currently cute.math.tanh(x, fastmath=True) generates very slow code
-            # * (1 + cute.math.tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x)), fastmath=True))
-            * (1.0 + tanh(x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * (x * x))))
-        )
     else:
         x_sq = cute.arch.mul_packed_f32x2(x, x)
         x_sq_scaled = cute.arch.fma_packed_f32x2(
             x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi)
         )
         z = cute.arch.mul_packed_f32x2(x, x_sq_scaled)
-        tanh_z = (tanh(z[0]), tanh(z[1]))
         x_tanh_z = cute.arch.fma_packed_f32x2(tanh_z, x, x)
         return cute.arch.mul_packed_f32x2((0.5, 0.5), x_tanh_z)
@@ -162,11 +193,16 @@ def dgelu_tanh_approx(
         # Compute gradient
         # sech^2(z) = 1 - tanh^2(z)
-        sech2_z = 1 - tanh_z * tanh_z
         # dz/dx = c1 + 3 * c2 * x^2
         dz_dx = sqrt_2_over_pi + sqrt_2_over_pi_coeff_3 * x_sq
         # d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
-        dgelu = half_tanh_z_plus_one + x * (0.5 * (sech2_z * dz_dx))
         dx = dout * dgelu
         return dx, gelu_out
@@ -177,7 +213,7 @@ def dgelu_tanh_approx(
             x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi)
         )
         z = cute.arch.mul_packed_f32x2(x, x_sq_scaled)
-        tanh_z = (tanh(z[0]), tanh(z[1]))
         half_tanh_z_plus_one = cute.arch.fma_packed_f32x2(tanh_z, (0.5, 0.5), (0.5, 0.5))
         gelu_out = cute.arch.mul_packed_f32x2(x, half_tanh_z_plus_one)
@@ -236,7 +272,18 @@ def dsoftplus_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) ->
 @dsl_user_op
-def silu(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) -> F32_or_F32x2:
     """
     silu(x) = x * sigmoid(x) = x * (1 + tanh(x / 2)) / 2 = (0.5 * x) * tanh(0.5 * x) + (0.5 * x)
     This compiles down to 3 SASS instructions: FMUL to get 0.5 * x, MUFU.TANH, and FFMA.
@@ -247,10 +294,91 @@ def silu(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) ->
         return x_half * tanh(x_half) + x_half
     else:
         x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) if const_expr(not already_halved) else x
-        tanh_x_half = (tanh(x_half[0]), tanh(x_half[1]))
         return cute.arch.fma_packed_f32x2(x_half, tanh_x_half, x_half)
 @dsl_user_op
 def swiglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
@@ -259,13 +387,20 @@ def swiglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32
         return cute.arch.mul_packed_f32x2(silu(x), y)
 @dsl_user_op
 def dswiglu(
     x: F32_or_F32x2,
     y: F32_or_F32x2,
     dout: F32_or_F32x2,
     *,
-    already_halved: bool = False,
     loc=None,
     ip=None,
 ) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
@@ -280,15 +415,8 @@ def dswiglu(
     to use FFMA instead of FADD and FMUL).
     """
     if const_expr(not isinstance(x, tuple)):
-        # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(0.5 * x))
-        # FMUL, MUFU.TANH, then FFMA
-        if const_expr(not already_halved):
-            sigmoid_x = sigmoid(x)
-            silu_x = x * sigmoid_x  # FMUL
-        else:
-            tanh_x = tanh(x)  # MUFU.TANH
-            sigmoid_x = 0.5 * tanh_x + 0.5  # FFMA
-            silu_x = x * tanh_x + x  # FFMA
         silu_x_dout = silu_x * dout  # FMUL
         #   d_silu(x) * dout
         # = sigmoid_x * (1 + x * (1 - sigmoid_x)) * dout
@@ -296,19 +424,65 @@ def dswiglu(
         # = (sigmoid_x + silu_x * (1 - sigmoid_x)) * dout
         # = (sigmoid_x + silu_x - silu_x * sigmoid_x) * dout
         # = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout
-        d_silu_x_dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x_dout  # FFMA, FFMA
         dx = d_silu_x_dout * y  # FMUL
         dy = silu_x_dout
         swiglu_out = silu_x * y  # FMUL
-        # Overall it's 1 MUFU.TANH, 5 FMUL, 3 FFMA
         return dx, dy, swiglu_out
     else:
         # Compute sigmoid(x) and silu(x)
         if const_expr(not already_halved):
-            sigmoid_x = sigmoid(x)
             silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x)
         else:
-            tanh_x = (tanh(x[0]), tanh(x[1]))
             sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x, (0.5, 0.5), (0.5, 0.5))
             silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x)
         silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
@@ -332,18 +506,31 @@ def swiglu_oai(
     """The swiglu variant used in gpt-oss, which has a scaling factor on x and bias of 1 to y.
     https://github.com/openai/gpt-oss/blob/7be9334950053a888e24887a57dac797a17d6e00/gpt_oss/torch/model.py#L249
     x * sigmoid(alpha * x) * (y + 1)
-    Compile down to FMUL, FMUL, TANH, FFMA, FFMA
     """
-    # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2))
     if const_expr(not isinstance(x, tuple)):
         x_half = 0.5 * x
-        # silu_x = x_half * cute.math.tanh(alpha * x_half, fastmath=True) + x_half
         silu_x = x_half * tanh(alpha * x_half) + x_half
         return silu_x * y + silu_x
     else:
         x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
         alpha_x_half = cute.arch.mul_packed_f32x2((alpha, alpha), x_half)
-        tanh_alpha_x_half = (tanh(alpha_x_half[0]), tanh(alpha_x_half[1]))
         silu_x = cute.arch.fma_packed_f32x2(x_half, tanh_alpha_x_half, x_half)
         return cute.arch.fma_packed_f32x2(silu_x, y, silu_x)
@@ -361,28 +548,62 @@ def dswiglu_oai(
     d/dx[x * sigmoid(alpha * x)] = sigmoid(alpha * x) + alpha * x * sigmoid(alpha * x) * (1 - sigmoid(alpha * x))
     """
     if const_expr(not isinstance(x, tuple)):
-        # Compute sigmoid(alpha * x) using tanh: sigmoid(z) = 0.5 * (1 + tanh(z/2))
-        alpha_x_half = (0.5 * alpha) * x  # FMUL
-        # MUFU.TANH, then FFMA
-        # sigmoid_alpha_x = 0.5 + 0.5 * cute.math.tanh(alpha_x_half, fastmath=True)
         sigmoid_alpha_x = 0.5 + 0.5 * tanh(alpha_x_half)
-        silu_x = x * sigmoid_alpha_x  # FMUL
-        silu_x_dout = silu_x * dout  # FMUL
-        # FFMA, FFMA, FMUL
-        d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout
-        dx = d_silu_x_dout * y + d_silu_x_dout  # FFMA, instead of multiply by y + 1
         dy = silu_x_dout
-        swiglu_out = silu_x * y + silu_x  # FFMA, instead of multiply by y + 1
-        # Overall it's 1 MUFU.TANH, 4 FMUL, 5 FFMA
         return dx, dy, swiglu_out
     else:
-        # Compute sigmoid(alpha * x)
         alpha_x_half = cute.arch.mul_packed_f32x2(((0.5 * alpha), (0.5 * alpha)), x)
-        tanh_alpha_x_half = (tanh(alpha_x_half[0]), tanh(alpha_x_half[1]))
         sigmoid_alpha_x = cute.arch.fma_packed_f32x2(tanh_alpha_x_half, (0.5, 0.5), (0.5, 0.5))
         silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x)
         silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
-        # d_silu_x_dout = (sigmoid_alpha_x + alpha * (silu_x - silu_x * sigmoid_alpha_x)) * dout
         silu_x_minus_product = cute.arch.fma_packed_f32x2(
             silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x
         )
@@ -400,10 +621,9 @@ def dswiglu_oai(
 def glu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     """GLU: Gated Linear Unit
     glu(x, y) = sigmoid(x) * y
-    Using tanh to compute sigmoid: sigmoid(x) = 0.5 * (1 + tanh(x/2))
     """
     if const_expr(not isinstance(x, tuple)):
-        sigmoid_x = sigmoid(x)  # FMUL, MUFU.TANH, then FFMA
         return sigmoid_x * y  # FMUL
     else:
         sigmoid_x = sigmoid(x)
@@ -423,8 +643,7 @@ def dglu(
     - glu_out = sigmoid(x) * y
     """
     if const_expr(not isinstance(x, tuple)):
-        # Compute sigmoid(x) using tanh: sigmoid(x) = 0.5 * (1 + tanh(x/2))
-        sigmoid_x = sigmoid(x)  # FMUL, MUFU.TANH, then FFMA
         sigmoid_x_dout = sigmoid_x * dout  # FMUL
         glu_out = sigmoid_x * y  # FMUL
         # dx = y * sigmoid(x) * (1 - sigmoid(x)) * dout
@@ -452,9 +671,9 @@ def reglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x
     reglu(x, y) = relu(x) * y = max(x, 0) * y
     """
     if const_expr(not isinstance(x, tuple)):
-        return cute.arch.fmax(x, Float32(0.0)) * y
     else:
-        relu_x = relu(x)
         return cute.arch.mul_packed_f32x2(relu_x, y)
@@ -473,15 +692,16 @@ def dreglu(
     """
     if const_expr(not isinstance(x, tuple)):
         x_pos = Boolean(x > 0)
-        relu_x = cute.arch.fmax(x, Float32(0.0))
-        dx = (dout * y) if x_pos else Float32(0.0)
         dy = dout * relu_x
         reglu_out = relu_x * y
         return dx, dy, reglu_out
     else:
         x0_pos = Boolean(x[0] > 0)
         x1_pos = Boolean(x[1] > 0)
-        relu_x = relu(x)
         dout_y = cute.arch.mul_packed_f32x2(dout, y)
         dx = ((dout_y[0] if x0_pos else Float32(0.0)), (dout_y[1] if x1_pos else Float32(0.0)))
         dy = cute.arch.mul_packed_f32x2(dout, relu_x)
@@ -538,21 +758,28 @@ def dgeglu(
 act_fn_map = {
     None: None,
     "silu": silu,
     "relu": relu,
     "relu_sq": relu_sq,
     "gelu_tanh_approx": gelu_tanh_approx,
 }
 dact_fn_map = {
     None: None,
     "relu": drelu,
     "relu_sq": drelu_sq,
     "gelu_tanh_approx": dgelu_tanh_approx,
 }
 gate_fn_map = {
     "swiglu": swiglu,
     "swiglu_oai": swiglu_oai,
     "reglu": reglu,
     "geglu": geglu,
     "glu": glu,
@@ -560,7 +787,9 @@ gate_fn_map = {
 dgate_fn_map = {
     "swiglu": dswiglu,
     "swiglu_oai": dswiglu_oai,
     "reglu": dreglu,
     "geglu": dgeglu,
     "glu": dglu,

 from typing import Tuple
 from functools import partial
+import cutlass
 import cutlass.cute as cute
 from cutlass import Float32, Boolean, const_expr
+from cutlass.cutlass_dsl import dsl_user_op
+from cutlass._mlir.dialects import nvvm
+from cutlass._mlir_helpers import math as mlir_math
 F32_or_F32x2 = Float32 | Tuple[Float32, Float32]
 @dsl_user_op
+def tanh(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
+    if const_expr(not isinstance(x, tuple)):
+        return cute.math.tanh(x, fastmath=True, loc=loc, ip=ip)
+    else:
+        return (
+            cute.math.tanh(x[0], fastmath=True, loc=loc, ip=ip),
+            cute.math.tanh(x[1], fastmath=True, loc=loc, ip=ip),
         )
 @dsl_user_op
+def dtanh(
+    x: F32_or_F32x2, dout: F32_or_F32x2, *, loc=None, ip=None
+) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
+    if const_expr(not isinstance(x, tuple)):
+        tanh_x = tanh(x, loc=loc, ip=ip)
+        dx = dout * (1.0 - tanh_x * tanh_x)
+        return dx, tanh_x
+    else:
+        tanh_x = tanh(x, loc=loc, ip=ip)
+        sech2_x = cute.arch.fma_packed_f32x2(tanh_x, (-tanh_x[0], -tanh_x[1]), (1.0, 1.0))
+        dx = cute.arch.mul_packed_f32x2(dout, sech2_x)
+        return dx, tanh_x
+@dsl_user_op
+def sigmoid_tanh(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
         # return 0.5 + 0.5 * cute.math.tanh(0.5 * x, fastmath=True)
         return 0.5 + 0.5 * tanh(0.5 * x)
     else:
         x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
+        tanh_x_half = tanh(x_half)
         return cute.arch.fma_packed_f32x2(tanh_x_half, (0.5, 0.5), (0.5, 0.5))
+@dsl_user_op
+def sigmoid(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
+    log2_e = math.log2(math.e)
+    if const_expr(not isinstance(x, tuple)):
+        exp_neg_x = cute.math.exp2(x * (-log2_e), fastmath=True, loc=loc, ip=ip)
+        return mlir_math.rcp(exp_neg_x + 1.0, approx=True, ftz=True, loc=loc, ip=ip)
+    else:
+        neg_x = cute.arch.mul_packed_f32x2(x, (-log2_e, -log2_e))
+        exp_neg_x = (
+            cute.math.exp2(neg_x[0], fastmath=True, loc=loc, ip=ip),
+            cute.math.exp2(neg_x[1], fastmath=True, loc=loc, ip=ip),
+        )
+        denom = cute.arch.add_packed_f32x2(exp_neg_x, (1.0, 1.0))
+        return (
+            mlir_math.rcp(denom[0], approx=True, ftz=True, loc=loc, ip=ip),
+            mlir_math.rcp(denom[1], approx=True, ftz=True, loc=loc, ip=ip),
+        )
 @dsl_user_op
 def dsigmoid_from_output(out: Float32, dout: Float32, *, loc=None, ip=None) -> Float32:
     # return dout * out * (1.0 - out)
 @dsl_user_op
 def relu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
+        return cutlass.max(x, Float32(0.0), loc=loc, ip=ip)
     else:
+        return relu(x[0], loc=loc, ip=ip), relu(x[1], loc=loc, ip=ip)
 @dsl_user_op
 ) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
     if const_expr(not isinstance(x, tuple)):
         x_pos = Boolean(x > 0)
+        return dout if x_pos else Float32(0.0), relu(x, loc=loc, ip=ip)
     else:
         x0_pos = Boolean(x[0] > 0)
         x1_pos = Boolean(x[1] > 0)
 @dsl_user_op
 def relu_sq(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
+        return relu(x, loc=loc, ip=ip) * x
     else:
+        relu_x = relu(x, loc=loc, ip=ip)
         return cute.arch.mul_packed_f32x2(relu_x, x)
     sqrt_2_over_pi = math.sqrt(2 / math.pi)  # ~0.797885
     sqrt_2_over_pi_coeff = 0.044715 * sqrt_2_over_pi  # ~0.0356774
     if const_expr(not isinstance(x, tuple)):
+        x_sq = x * x
+        z = x * (sqrt_2_over_pi + sqrt_2_over_pi_coeff * x_sq)
+        tanh_z = tanh(z)
+        return 0.5 * (x * tanh_z + x)
     else:
         x_sq = cute.arch.mul_packed_f32x2(x, x)
         x_sq_scaled = cute.arch.fma_packed_f32x2(
             x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi)
         )
         z = cute.arch.mul_packed_f32x2(x, x_sq_scaled)
+        tanh_z = tanh(z)
         x_tanh_z = cute.arch.fma_packed_f32x2(tanh_z, x, x)
         return cute.arch.mul_packed_f32x2((0.5, 0.5), x_tanh_z)
         # Compute gradient
         # sech^2(z) = 1 - tanh^2(z)
+        # Keep this as a multiply-add expression so vectorize=True lowers to
+        # FFMA2 like the explicit F32x2 path; `1.0 - tanh_z * tanh_z` costs
+        # an extra FADD2/FMUL2 pair per vector.
+        sech2_z = tanh_z * (-tanh_z) + 1.0
         # dz/dx = c1 + 3 * c2 * x^2
         dz_dx = sqrt_2_over_pi + sqrt_2_over_pi_coeff_3 * x_sq
         # d/dx[gelu(x)] = 0.5 * (1 + tanh(z)) + 0.5 * x * sech^2(z) * dz/dx
+        sech2_dz_dx = sech2_z * dz_dx
+        x_sech2_dz_dx = x * sech2_dz_dx
+        dgelu = x_sech2_dz_dx * 0.5 + half_tanh_z_plus_one
         dx = dout * dgelu
         return dx, gelu_out
             x_sq, (sqrt_2_over_pi_coeff, sqrt_2_over_pi_coeff), (sqrt_2_over_pi, sqrt_2_over_pi)
         )
         z = cute.arch.mul_packed_f32x2(x, x_sq_scaled)
+        tanh_z = tanh(z)
         half_tanh_z_plus_one = cute.arch.fma_packed_f32x2(tanh_z, (0.5, 0.5), (0.5, 0.5))
         gelu_out = cute.arch.mul_packed_f32x2(x, half_tanh_z_plus_one)
 @dsl_user_op
+def silu(x: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
+    """
+    silu(x) = x * sigmoid(x) = x * rcp(1 + exp(-x)).
+    """
+    if const_expr(not isinstance(x, tuple)):
+        return x * sigmoid(x, loc=loc, ip=ip)
+    else:
+        return cute.arch.mul_packed_f32x2(x, sigmoid(x, loc=loc, ip=ip))
+@dsl_user_op
+def silu_tanh(x: F32_or_F32x2, *, already_halved: bool = False, loc=None, ip=None) -> F32_or_F32x2:
     """
     silu(x) = x * sigmoid(x) = x * (1 + tanh(x / 2)) / 2 = (0.5 * x) * tanh(0.5 * x) + (0.5 * x)
     This compiles down to 3 SASS instructions: FMUL to get 0.5 * x, MUFU.TANH, and FFMA.
         return x_half * tanh(x_half) + x_half
     else:
         x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x) if const_expr(not already_halved) else x
+        tanh_x_half = tanh(x_half)
         return cute.arch.fma_packed_f32x2(x_half, tanh_x_half, x_half)
+@dsl_user_op
+def dsilu(
+    x: F32_or_F32x2,
+    dout: F32_or_F32x2,
+    *,
+    loc=None,
+    ip=None,
+) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
+    """
+    SiLU backward pass: computes d_silu(x) * dout and recomputes silu(x).
+    d_silu(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))).
+    """
+    if const_expr(not isinstance(x, tuple)):
+        sigmoid_x = sigmoid(x, loc=loc, ip=ip)
+        silu_x = x * sigmoid_x
+        # This form vectorizes cleanly with cutlass.range(..., vectorize=True):
+        # FADD2 (1 - sigmoid_x), FFMA2 (silu_x * tmp + sigmoid_x), FMUL2 (* dout).
+        d_silu_x_dout = (sigmoid_x + silu_x * (1.0 - sigmoid_x)) * dout
+        return d_silu_x_dout, silu_x
+    else:
+        sigmoid_x = sigmoid(x)
+        silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x)
+        sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2(
+            sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x
+        )
+        sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x = cute.arch.add_packed_f32x2(
+            sigmoid_x_minus_silu_x_sigmoid_x, silu_x
+        )
+        d_silu_x_dout = cute.arch.mul_packed_f32x2(
+            sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x, dout
+        )
+        return d_silu_x_dout, silu_x
+@dsl_user_op
+def dsilu_tanh(
+    x: F32_or_F32x2,
+    dout: F32_or_F32x2,
+    *,
+    already_halved: bool = False,
+    loc=None,
+    ip=None,
+) -> Tuple[F32_or_F32x2, F32_or_F32x2]:
+    """
+    SiLU backward using sigmoid(x) = 0.5 * (1 + tanh(0.5 * x)).
+    """
+    if const_expr(not isinstance(x, tuple)):
+        if const_expr(not already_halved):
+            x_half = 0.5 * x
+            tanh_x_half = tanh(x_half)
+            sigmoid_x = 0.5 * tanh_x_half + 0.5
+            silu_x = x_half * tanh_x_half + x_half
+        else:
+            tanh_x = tanh(x)
+            sigmoid_x = 0.5 * tanh_x + 0.5
+            silu_x = x * tanh_x + x
+        d_silu_x_dout = (sigmoid_x + silu_x * (1.0 - sigmoid_x)) * dout
+        return d_silu_x_dout, silu_x
+    else:
+        if const_expr(not already_halved):
+            x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
+            tanh_x_half = tanh(x_half)
+            sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x_half, (0.5, 0.5), (0.5, 0.5))
+            silu_x = cute.arch.fma_packed_f32x2(x_half, tanh_x_half, x_half)
+        else:
+            tanh_x = tanh(x)
+            sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x, (0.5, 0.5), (0.5, 0.5))
+            silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x)
+        sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2(
+            sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x
+        )
+        sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x = cute.arch.add_packed_f32x2(
+            sigmoid_x_minus_silu_x_sigmoid_x, silu_x
+        )
+        d_silu_x_dout = cute.arch.mul_packed_f32x2(
+            sigmoid_x_minus_silu_x_sigmoid_x_plus_silu_x, dout
+        )
+        return d_silu_x_dout, silu_x
 @dsl_user_op
 def swiglu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     if const_expr(not isinstance(x, tuple)):
         return cute.arch.mul_packed_f32x2(silu(x), y)
+@dsl_user_op
+def swiglu_tanh(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
+    if const_expr(not isinstance(x, tuple)):
+        return silu_tanh(x) * y
+    else:
+        return cute.arch.mul_packed_f32x2(silu_tanh(x), y)
 @dsl_user_op
 def dswiglu(
     x: F32_or_F32x2,
     y: F32_or_F32x2,
     dout: F32_or_F32x2,
     *,
     loc=None,
     ip=None,
 ) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
     to use FFMA instead of FADD and FMUL).
     """
     if const_expr(not isinstance(x, tuple)):
+        sigmoid_x = sigmoid(x)
+        silu_x = x * sigmoid_x  # FMUL
         silu_x_dout = silu_x * dout  # FMUL
         #   d_silu(x) * dout
         # = sigmoid_x * (1 + x * (1 - sigmoid_x)) * dout
         # = (sigmoid_x + silu_x * (1 - sigmoid_x)) * dout
         # = (sigmoid_x + silu_x - silu_x * sigmoid_x) * dout
         # = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout
+        # This form lets ptxas recover the same two packed FFMA instructions
+        # as the explicit F32x2 path while reusing silu_x_dout for dy.
+        d_silu_x_dout = (sigmoid_x + (-silu_x) * sigmoid_x) * dout + silu_x_dout
         dx = d_silu_x_dout * y  # FMUL
         dy = silu_x_dout
         swiglu_out = silu_x * y  # FMUL
         return dx, dy, swiglu_out
     else:
         # Compute sigmoid(x) and silu(x)
+        sigmoid_x = sigmoid(x)
+        silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x)
+        silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
+        # d_silu(x) * dout = (sigmoid_x - silu_x * sigmoid_x) * dout + silu_x * dout
+        sigmoid_x_minus_silu_x_sigmoid_x = cute.arch.fma_packed_f32x2(
+            sigmoid_x, (-silu_x[0], -silu_x[1]), sigmoid_x
+        )
+        d_silu_x_dout = cute.arch.fma_packed_f32x2(
+            sigmoid_x_minus_silu_x_sigmoid_x, dout, silu_x_dout
+        )
+        dx = cute.arch.mul_packed_f32x2(d_silu_x_dout, y)
+        dy = silu_x_dout
+        swiglu_out = cute.arch.mul_packed_f32x2(silu_x, y)
+        return dx, dy, swiglu_out
+@dsl_user_op
+def dswiglu_tanh(
+    x: F32_or_F32x2,
+    y: F32_or_F32x2,
+    dout: F32_or_F32x2,
+    *,
+    already_halved: bool = False,
+    loc=None,
+    ip=None,
+) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
+    """
+    SwiGLU backward using sigmoid(x) = 0.5 * (1 + tanh(0.5 * x)).
+    """
+    if const_expr(not isinstance(x, tuple)):
+        if const_expr(not already_halved):
+            sigmoid_x = sigmoid_tanh(x)
+            silu_x = x * sigmoid_x  # FMUL
+        else:
+            tanh_x = tanh(x)
+            sigmoid_x = 0.5 * tanh_x + 0.5
+            silu_x = x * tanh_x + x
+        silu_x_dout = silu_x * dout
+        d_silu_x_dout = (sigmoid_x + (-silu_x) * sigmoid_x) * dout + silu_x_dout
+        dx = d_silu_x_dout * y
+        dy = silu_x_dout
+        swiglu_out = silu_x * y
+        # Overall it's 1 MUFU.TANH, 5 FMUL, 3 FFMA
+        return dx, dy, swiglu_out
+    else:
         if const_expr(not already_halved):
+            sigmoid_x = sigmoid_tanh(x)
             silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_x)
         else:
+            tanh_x = tanh(x)
             sigmoid_x = cute.arch.fma_packed_f32x2(tanh_x, (0.5, 0.5), (0.5, 0.5))
             silu_x = cute.arch.fma_packed_f32x2(x, tanh_x, x)
         silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
     """The swiglu variant used in gpt-oss, which has a scaling factor on x and bias of 1 to y.
     https://github.com/openai/gpt-oss/blob/7be9334950053a888e24887a57dac797a17d6e00/gpt_oss/torch/model.py#L249
     x * sigmoid(alpha * x) * (y + 1)
     """
+    if const_expr(not isinstance(x, tuple)):
+        sigmoid_alpha_x = sigmoid(alpha * x)
+        silu_x = x * sigmoid_alpha_x
+        return silu_x * y + silu_x
+    else:
+        alpha_x = cute.arch.mul_packed_f32x2((alpha, alpha), x)
+        sigmoid_alpha_x = sigmoid(alpha_x)
+        silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x)
+        return cute.arch.fma_packed_f32x2(silu_x, y, silu_x)
+@dsl_user_op
+def swiglu_oai_tanh(
+    x: F32_or_F32x2, y: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None
+) -> F32_or_F32x2:
+    """Tanh-based swiglu_oai kept for SASS/accuracy comparison."""
     if const_expr(not isinstance(x, tuple)):
         x_half = 0.5 * x
         silu_x = x_half * tanh(alpha * x_half) + x_half
         return silu_x * y + silu_x
     else:
         x_half = cute.arch.mul_packed_f32x2((0.5, 0.5), x)
         alpha_x_half = cute.arch.mul_packed_f32x2((alpha, alpha), x_half)
+        tanh_alpha_x_half = tanh(alpha_x_half)
         silu_x = cute.arch.fma_packed_f32x2(x_half, tanh_alpha_x_half, x_half)
         return cute.arch.fma_packed_f32x2(silu_x, y, silu_x)
     d/dx[x * sigmoid(alpha * x)] = sigmoid(alpha * x) + alpha * x * sigmoid(alpha * x) * (1 - sigmoid(alpha * x))
     """
     if const_expr(not isinstance(x, tuple)):
+        sigmoid_alpha_x = sigmoid(alpha * x)
+        silu_x = x * sigmoid_alpha_x
+        silu_x_dout = silu_x * dout
+        # Keep this as two multiply-add expressions. With vectorize=True this
+        # matches the explicit F32x2 path; spelling it as (1 - sigmoid) costs
+        # an extra FADD2/FMUL2 pair per vector.
+        silu_x_minus_product = silu_x * (-sigmoid_alpha_x) + silu_x
+        sigmoid_plus_alpha_diff = silu_x_minus_product * alpha + sigmoid_alpha_x
+        d_silu_x_dout = sigmoid_plus_alpha_diff * dout
+        dx = d_silu_x_dout * y + d_silu_x_dout
+        dy = silu_x_dout
+        swiglu_out = silu_x * y + silu_x
+        return dx, dy, swiglu_out
+    else:
+        alpha_x = cute.arch.mul_packed_f32x2((alpha, alpha), x)
+        sigmoid_alpha_x = sigmoid(alpha_x)
+        silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x)
+        silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
+        silu_x_minus_product = cute.arch.fma_packed_f32x2(
+            silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x
+        )
+        sigmoid_plus_alpha_diff = cute.arch.fma_packed_f32x2(
+            (alpha, alpha), silu_x_minus_product, sigmoid_alpha_x
+        )
+        d_silu_x_dout = cute.arch.mul_packed_f32x2(sigmoid_plus_alpha_diff, dout)
+        dx = cute.arch.fma_packed_f32x2(d_silu_x_dout, y, d_silu_x_dout)
+        dy = silu_x_dout
+        swiglu_out = cute.arch.fma_packed_f32x2(silu_x, y, silu_x)
+        return dx, dy, swiglu_out
+@dsl_user_op
+def dswiglu_oai_tanh(
+    x: F32_or_F32x2, y: F32_or_F32x2, dout: F32_or_F32x2, alpha: float = 1.702, *, loc=None, ip=None
+) -> Tuple[F32_or_F32x2, F32_or_F32x2, F32_or_F32x2]:
+    """Tanh-based dswiglu_oai kept for SASS/accuracy comparison."""
+    if const_expr(not isinstance(x, tuple)):
+        alpha_x_half = (0.5 * alpha) * x
         sigmoid_alpha_x = 0.5 + 0.5 * tanh(alpha_x_half)
+        silu_x = x * sigmoid_alpha_x
+        silu_x_dout = silu_x * dout
+        # Same spelling as dswiglu_oai: this preserves the packed FFMA2 chain
+        # under cutlass.range(..., vectorize=True).
+        silu_x_minus_product = silu_x * (-sigmoid_alpha_x) + silu_x
+        sigmoid_plus_alpha_diff = silu_x_minus_product * alpha + sigmoid_alpha_x
+        d_silu_x_dout = sigmoid_plus_alpha_diff * dout
+        dx = d_silu_x_dout * y + d_silu_x_dout
         dy = silu_x_dout
+        swiglu_out = silu_x * y + silu_x
         return dx, dy, swiglu_out
     else:
         alpha_x_half = cute.arch.mul_packed_f32x2(((0.5 * alpha), (0.5 * alpha)), x)
+        tanh_alpha_x_half = tanh(alpha_x_half)
         sigmoid_alpha_x = cute.arch.fma_packed_f32x2(tanh_alpha_x_half, (0.5, 0.5), (0.5, 0.5))
         silu_x = cute.arch.mul_packed_f32x2(x, sigmoid_alpha_x)
         silu_x_dout = cute.arch.mul_packed_f32x2(silu_x, dout)
         silu_x_minus_product = cute.arch.fma_packed_f32x2(
             silu_x, (-sigmoid_alpha_x[0], -sigmoid_alpha_x[1]), silu_x
         )
 def glu(x: F32_or_F32x2, y: F32_or_F32x2, *, loc=None, ip=None) -> F32_or_F32x2:
     """GLU: Gated Linear Unit
     glu(x, y) = sigmoid(x) * y
     """
     if const_expr(not isinstance(x, tuple)):
+        sigmoid_x = sigmoid(x)
         return sigmoid_x * y  # FMUL
     else:
         sigmoid_x = sigmoid(x)
     - glu_out = sigmoid(x) * y
     """
     if const_expr(not isinstance(x, tuple)):
+        sigmoid_x = sigmoid(x)
         sigmoid_x_dout = sigmoid_x * dout  # FMUL
         glu_out = sigmoid_x * y  # FMUL
         # dx = y * sigmoid(x) * (1 - sigmoid(x)) * dout
     reglu(x, y) = relu(x) * y = max(x, 0) * y
     """
     if const_expr(not isinstance(x, tuple)):
+        return relu(x, loc=loc, ip=ip) * y
     else:
+        relu_x = relu(x, loc=loc, ip=ip)
         return cute.arch.mul_packed_f32x2(relu_x, y)
     """
     if const_expr(not isinstance(x, tuple)):
         x_pos = Boolean(x > 0)
+        relu_x = relu(x, loc=loc, ip=ip)
+        dout_y = dout * y
+        dx = dout_y if x_pos else Float32(0.0)
         dy = dout * relu_x
         reglu_out = relu_x * y
         return dx, dy, reglu_out
     else:
         x0_pos = Boolean(x[0] > 0)
         x1_pos = Boolean(x[1] > 0)
+        relu_x = relu(x, loc=loc, ip=ip)
         dout_y = cute.arch.mul_packed_f32x2(dout, y)
         dx = ((dout_y[0] if x0_pos else Float32(0.0)), (dout_y[1] if x1_pos else Float32(0.0)))
         dy = cute.arch.mul_packed_f32x2(dout, relu_x)
 act_fn_map = {
     None: None,
     "silu": silu,
+    "silu-tanh": silu_tanh,
     "relu": relu,
     "relu_sq": relu_sq,
     "gelu_tanh_approx": gelu_tanh_approx,
+    "tanh": tanh,
 }
 dact_fn_map = {
     None: None,
+    "silu": dsilu,
+    "silu-tanh": dsilu_tanh,
     "relu": drelu,
     "relu_sq": drelu_sq,
     "gelu_tanh_approx": dgelu_tanh_approx,
+    "tanh": dtanh,
 }
 gate_fn_map = {
     "swiglu": swiglu,
+    "swiglu-tanh": swiglu_tanh,
     "swiglu_oai": swiglu_oai,
+    "swiglu_oai-tanh": swiglu_oai_tanh,
     "reglu": reglu,
     "geglu": geglu,
     "glu": glu,
 dgate_fn_map = {
     "swiglu": dswiglu,
+    "swiglu-tanh": dswiglu_tanh,
     "swiglu_oai": dswiglu_oai,
+    "swiglu_oai-tanh": dswiglu_oai_tanh,
     "reglu": dreglu,
     "geglu": dgeglu,
     "glu": dglu,

build/torch-cuda/quack/autotuner.py CHANGED Viewed

@@ -4,6 +4,7 @@ from __future__ import annotations
 import builtins
 import os
 import time
 import inspect
 import base64
@@ -12,6 +13,11 @@ import json
 from pathlib import Path
 from functools import cached_property, partial
 from typing import Dict, Tuple, List, Optional, Any
 import torch
 from torch import Tensor
@@ -25,29 +31,6 @@ PACKAGE_NAME = "quack"
 VERSION = __version__
-def _get_current_cuda_device() -> str | None:
-    """Return the physical CUDA device identifier for the current process.
-    Maps the logical ``torch.cuda.current_device()`` index through
-    ``CUDA_VISIBLE_DEVICES`` (if set) so the result is valid as a
-    standalone ``CUDA_VISIBLE_DEVICES`` value (handles integer IDs,
-    GPU UUIDs, and MIG IDs).
-    Returns ``None`` if CUDA is not initialized or the device cannot
-    be determined.
-    """
-    if not (torch.cuda.is_available() and torch.cuda.is_initialized()):
-        return None
-    logical_device = torch.cuda.current_device()
-    parent_visible = os.environ.get("CUDA_VISIBLE_DEVICES")
-    if parent_visible is not None:
-        visible_devices = [d.strip() for d in parent_visible.split(",")]
-        if logical_device < len(visible_devices):
-            return visible_devices[logical_device]
-        return None
-    return str(logical_device)
 def get_home_dir():
     return os.getenv(f"{PACKAGE_NAME.upper()}_HOME", Path.home())
@@ -75,6 +58,14 @@ def _base32(key):
     return base64.b32encode(bytes.fromhex(key)).decode("utf-8").rstrip("=")
 def _gpu_warmup(duration_ms=200):
     """Saturate the GPU to reach thermal steady-state before benchmarking.
@@ -91,6 +82,21 @@ def _gpu_warmup(duration_ms=200):
         torch.cuda.synchronize()
 class Autotuner:
     def __init__(
         self,
@@ -163,146 +169,6 @@ class Autotuner:
             return partial(triton.testing.do_bench, warmup=5, rep=25)
         return self._do_bench
-    def _precompile(self, *args, configs, **kwargs):
-        """Pre-compile all configs in parallel subprocesses to populate .o cache.
-        cute.compile() is not thread-safe (MLIR thread-local state) and fork after
-        CUDA init causes segfaults. So we spawn persistent subprocess workers: each
-        has its own CUDA context, creates FakeTensors matching the parent's tensor
-        metadata, and compiles with COMPILE_ONLY=True. Workers stay alive to amortize
-        import overhead across multiple configs. The parent then loads instantly from
-        the .o cache during benchmarking.
-        """
-        from .cache_utils import CACHE_ENABLED
-        if not CACHE_ENABLED:
-            return
-        max_workers = min(len(configs), int(os.getenv("QUACK_COMPILE_WORKERS", "8")))
-        if max_workers <= 1:
-            return
-        # Quick check: compile first config in-process. If it loads from .o cache
-        # (<0.5s), the rest are likely cached too — skip spawning workers.
-        t_check = time.time()
-        try:
-            current = dict(kwargs, **configs[0].all_kwargs())
-            self.fn(*args, **current)
-        except Exception:
-            pass
-        if time.time() - t_check < 0.5:
-            return
-        verbose = os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1"
-        if verbose:
-            print(f"Pre-compiling {len(configs)} configs with {max_workers} workers")
-        t0 = time.time()
-        import pickle
-        import struct
-        import subprocess
-        import sys
-        def _send(stream, msg):
-            data = pickle.dumps(msg)
-            stream.write(struct.pack("<I", len(data)))
-            stream.write(data)
-            stream.flush()
-        def _recv(stream):
-            header = stream.read(4)
-            if len(header) < 4:
-                return None
-            length = struct.unpack("<I", header)[0]
-            return pickle.loads(stream.read(length)) if length else None
-        # Serialize tensor metadata
-        tensor_meta = []
-        for arg in args:
-            if isinstance(arg, Tensor):
-                tensor_meta.append(
-                    {
-                        "shape": list(arg.shape),
-                        "stride": list(arg.stride()),
-                        "dtype": str(arg.dtype),
-                    }
-                )
-            else:
-                tensor_meta.append(arg)
-        fn_module = self.fn.__module__
-        fn_qualname = self.fn.__qualname__
-        # Restrict worker subprocesses to the parent's current CUDA device.
-        # Without this, all workers default to cuda:0 and their CUDA context
-        # initialization can OOM when many ranks share a node.
-        worker_env = os.environ.copy()
-        current_device = _get_current_cuda_device()
-        if current_device is not None:
-            worker_env["CUDA_VISIBLE_DEVICES"] = current_device
-        # Launch persistent worker pool. When vendored under sonic_moe (loaded
-        # via kernels.get_kernel), the quack package isn't importable as a
-        # top-level module, so invoke the worker via its fully-qualified dotted
-        # path and inject PYTHONPATH so the subprocess can import it.
-        worker_module = __package__ + "._compile_worker" if __package__ else "quack._compile_worker"
-        if __package__:
-            import importlib.util
-            spec = importlib.util.find_spec(__package__.split(".")[0])
-            if spec is not None and spec.submodule_search_locations:
-                pkg_parent = os.path.dirname(list(spec.submodule_search_locations)[0])
-                existing_pp = worker_env.get("PYTHONPATH", "")
-                worker_env["PYTHONPATH"] = (
-                    f"{pkg_parent}{os.pathsep}{existing_pp}" if existing_pp else pkg_parent
-                )
-        workers = []
-        for _ in range(max_workers):
-            p = subprocess.Popen(
-                [sys.executable, "-m", worker_module],
-                stdin=subprocess.PIPE,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.DEVNULL if not verbose else None,
-                env=worker_env,
-            )
-            ready = _recv(p.stdout)
-            if ready != "READY":
-                p.kill()
-                continue
-            workers.append(p)
-        if not workers:
-            return
-        # Round-robin dispatch configs to workers
-        pending = [0] * len(workers)
-        for i, config in enumerate(configs):
-            w = workers[i % len(workers)]
-            _send(
-                w.stdin,
-                {
-                    "fn_module": fn_module,
-                    "fn_qualname": fn_qualname,
-                    "tensor_meta": tensor_meta,
-                    "kwargs": kwargs,
-                    "config_kwargs": config.all_kwargs(),
-                },
-            )
-            pending[i % len(workers)] += 1
-        # Collect all results
-        for wi, w in enumerate(workers):
-            for _ in range(pending[wi]):
-                _recv(w.stdout)
-        # Shutdown workers (close stdin → worker exits)
-        for w in workers:
-            w.stdin.close()
-            w.wait()
-        if verbose:
-            print(f"Pre-compilation done in {time.time() - t0:.1f}s")
     def _bench(self, *args, config, **meta):
         verbose = os.environ.get(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1"
         if verbose:
@@ -320,6 +186,45 @@ class Autotuner:
         current = dict(meta, **config.all_kwargs())
         full_nargs = {**self.nargs, **current}
         def kernel_call():
             if self.pre_hook is not None:
                 self.pre_hook(full_nargs)
@@ -406,17 +311,118 @@ class Autotuner:
                 @torch.compiler.disable  # Don't want any tracing here
                 def benchmark():
-                    self._precompile(*args, configs=pruned_configs, **kwargs)
-                    _gpu_warmup()
                     bench_start = time.time()
-                    timings = {
-                        config: self._bench(*args, config=config, **kwargs)
-                        for config in pruned_configs
-                    }
                     bench_end = time.time()
-                    if os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1":
                         for config, time_ in timings.items():
                             print(f"[{config}] -> {time_[0]:.3f}ms")
                     self.bench_time = bench_end - bench_start
                     self.cache[key] = builtins.min(timings, key=timings.get)
                     self.configs_timings = timings

 import builtins
 import os
+import sys
 import time
 import inspect
 import base64
 from pathlib import Path
 from functools import cached_property, partial
 from typing import Dict, Tuple, List, Optional, Any
+from .bench.bench_utils import (
+    _bench_cuda_graph_l2_rotate,
+    _clone_l2_rotate_inputs,
+    _pick_l2_rotate_count,
+)
 import torch
 from torch import Tensor
 VERSION = __version__
 def get_home_dir():
     return os.getenv(f"{PACKAGE_NAME.upper()}_HOME", Path.home())
     return base64.b32encode(bytes.fromhex(key)).decode("utf-8").rstrip("=")
+#: How long a deferred config may wait on its pool compile before the bench
+#: loop stops trusting the pool and benches it with the pool suppressed
+#: (in-process compile). Guards against a wedged worker / a foreign flock
+#: holder that never produces the .o; without it a permanently-"pending"
+#: sha would rotate forever. Tests override this.
+_POOL_WEDGE_TIMEOUT_S = 300.0
 def _gpu_warmup(duration_ms=200):
     """Saturate the GPU to reach thermal steady-state before benchmarking.
         torch.cuda.synchronize()
+# ---------------------------------------------------------------------------
+# Candidate-config compilation
+#
+# There is no separate precompile phase: the bench loop in ``benchmark()``
+# (inside ``Autotuner.__call__``) runs under ``pool_scope()`` from
+# quack.cache.async_compile. A config whose kernel misses the .o cache
+# raises ``CompilePending`` from jit_cache after shipping the pickled
+# ``_compile_*`` key to a CPU worker; the loop rotates that config to the
+# back and benches whichever config is ready. Total wall stays
+# max(parallel_compile, serial_bench), key discovery uses the real tensors
+# in-process, and workers never launch kernels (they call the tensor-free
+# ``_compile_*`` functions directly).
+# ---------------------------------------------------------------------------
 class Autotuner:
     def __init__(
         self,
             return partial(triton.testing.do_bench, warmup=5, rep=25)
         return self._do_bench
     def _bench(self, *args, config, **meta):
         verbose = os.environ.get(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1"
         if verbose:
         current = dict(meta, **config.all_kwargs())
         full_nargs = {**self.nargs, **current}
+        # Default path: L2-cold CUDA-graph round-robin bench. ``__call__``
+        # sets ``self._l2_cold_arg_sets`` / ``self._l2_cold_kwarg_sets`` to
+        # pre-cloned (args, kwargs) sets once per shape (reused across all
+        # configs). Round-robin over fresh sets keeps the kernel measured
+        # under the cache-cold conditions that match production access
+        # patterns, so the autotuner picks configs that win at the same
+        # workload the user actually runs.
+        l2_cold_arg_sets = getattr(self, "_l2_cold_arg_sets", None)
+        l2_cold_kwarg_sets = getattr(self, "_l2_cold_kwarg_sets", None)
+        has_hooks = self.pre_hook is not None or self.post_hook is not None
+        use_l2_cold = (
+            self._do_bench is None
+            and l2_cold_arg_sets is not None
+            and l2_cold_kwarg_sets is not None
+            and not has_hooks
+        )
+        if use_l2_cold:
+            try:
+                return _bench_cuda_graph_l2_rotate(
+                    self.fn,
+                    l2_cold_arg_sets,
+                    l2_cold_kwarg_sets,
+                    extra_kwargs=config.all_kwargs(),
+                    quantiles=(0.5, 0.2, 0.8),
+                )
+            except (RuntimeError, MemoryError) as e:
+                # Narrow catch: only swallow GPU-side failures (smem
+                # overflow, kernel launch errors, OOM). Programming errors
+                # (TypeError, AssertionError, ValueError from conflict check
+                # above) propagate so the user sees them.
+                if verbose:
+                    print(f"Autotuning failed with {type(e).__name__}: {e}")
+                return [float("inf"), float("inf"), float("inf")]
+        # Legacy path: triton.testing.do_bench or user-supplied do_bench.
+        # Used when (a) a custom do_bench was passed via the decorator's
+        # ``do_bench=`` arg, or (b) pre/post hooks are configured (the
+        # clone/restore inside hooks doesn't work under CUDA graph capture).
         def kernel_call():
             if self.pre_hook is not None:
                 self.pre_hook(full_nargs)
                 @torch.compiler.disable  # Don't want any tracing here
                 def benchmark():
+                    # Compile/bench overlap via the async compile pool
+                    # (quack.cache.async_compile): the bench loop runs inside
+                    # pool_scope(). A config whose kernel isn't compiled yet
+                    # raises CompilePending from jit_cache (after shipping the
+                    # key to a CPU worker); the loop rotates it to the back
+                    # and benches whichever config is ready, retrying once its
+                    # .o lands. Discovery happens in-process with the real
+                    # tensors (no fake-tensor reconstruction), and the pool
+                    # workers replay the pickled _compile_* key directly --
+                    # which never launches, by construction.
+                    #
+                    # CompilePending can only fire OUTSIDE CUDA graph capture:
+                    # the L2-cold bench does priming launches before capture,
+                    # and the legacy do_bench path warms up first, so a cold
+                    # key raises at the first plain launch.
+                    from collections import deque
+                    from .cache.async_compile import (
+                        CompilePending,
+                        pool_scope,
+                        suppress_pool,
+                    )
                     bench_start = time.time()
+                    verbose = os.getenv(f"{PACKAGE_NAME.upper()}_PRINT_AUTOTUNING", None) == "1"
+                    has_hooks = self.pre_hook is not None or self.post_hook is not None
+                    timings = {}
+                    _MAX_ATTEMPTS = 20
+                    try:
+                        _gpu_warmup()
+                        # Pre-allocate cloned (args, kwargs) sets once per
+                        # shape; the same sets are reused across all configs
+                        # to avoid ~400x re-cloning. Skipped when hooks are
+                        # present or a custom do_bench was supplied (legacy
+                        # fallback in _bench).
+                        if self._do_bench is None and not has_hooks:
+                            try:
+                                n_buffers = _pick_l2_rotate_count(args, kwargs)
+                                arg_sets, kwarg_sets = _clone_l2_rotate_inputs(
+                                    args, kwargs, n_buffers
+                                )
+                                self._l2_cold_arg_sets = arg_sets
+                                self._l2_cold_kwarg_sets = kwarg_sets
+                            except (RuntimeError, MemoryError):
+                                # Cloning failed (likely OOM at extreme N);
+                                # legacy do_bench path will be used by _bench.
+                                self._l2_cold_arg_sets = None
+                                self._l2_cold_kwarg_sets = None
+                        else:
+                            self._l2_cold_arg_sets = None
+                            self._l2_cold_kwarg_sets = None
+                        with pool_scope() as pool:
+                            queue = deque(pruned_configs)
+                            awaiting = {}  # id(config) -> sha
+                            attempts = {}  # id(config) -> int
+                            deadline = {}  # id(config) -> wedge deadline
+                            spins = 0
+                            while queue:
+                                config = queue.popleft()
+                                sha = awaiting.get(id(config))
+                                wedged = sha is not None and time.monotonic() > deadline[id(config)]
+                                if sha is not None and not wedged:
+                                    state, _ = pool.poll(sha)
+                                    if state == "pending":
+                                        queue.append(config)
+                                        spins += 1
+                                        if spins >= len(queue):
+                                            time.sleep(0.05)
+                                            spins = 0
+                                        continue
+                                spins = 0
+                                n = attempts.get(id(config), 0) + 1
+                                attempts[id(config)] = n
+                                try:
+                                    if wedged or n > _MAX_ATTEMPTS:
+                                        # Wedged pool: compile in-process so
+                                        # the sweep always terminates.
+                                        with suppress_pool():
+                                            timings[config] = self._bench(
+                                                *args, config=config, **kwargs
+                                            )
+                                    else:
+                                        timings[config] = self._bench(
+                                            *args, config=config, **kwargs
+                                        )
+                                except CompilePending as e:
+                                    awaiting[id(config)] = e.sha
+                                    deadline.setdefault(
+                                        id(config),
+                                        time.monotonic() + _POOL_WEDGE_TIMEOUT_S,
+                                    )
+                                    queue.append(config)
+                    finally:
+                        # Free L2-cold sets before persisting the cache so the
+                        # user's subsequent .fn(...) call has full HBM.
+                        self._l2_cold_arg_sets = None
+                        self._l2_cold_kwarg_sets = None
                     bench_end = time.time()
+                    if verbose:
                         for config, time_ in timings.items():
                             print(f"[{config}] -> {time_[0]:.3f}ms")
+                    # Surface bench failures (configs returning inf timings)
+                    # so smem-overflow / launch errors aren't silently masked.
+                    n_failed = sum(1 for t in timings.values() if t[0] == float("inf"))
+                    if n_failed:
+                        print(
+                            f"quack autotune: {n_failed}/{len(timings)} configs "
+                            f"failed for {self.fn.__name__}{key}; "
+                            f"set {PACKAGE_NAME.upper()}_PRINT_AUTOTUNING=1 for details",
+                            file=sys.stderr,
+                        )
                     self.bench_time = bench_end - bench_start
                     self.cache[key] = builtins.min(timings, key=timings.get)
                     self.configs_timings = timings

build/torch-cuda/quack/bench/__init__.py ADDED Viewed

File without changes

build/torch-cuda/quack/bench/bench_utils.py ADDED Viewed

	@@ -0,0 +1,202 @@

+"""Shared helpers for triton perf_report-based benchmarks."""
+from __future__ import annotations
+import os
+from typing import TYPE_CHECKING
+import torch
+from torch import Tensor
+if TYPE_CHECKING:
+    import pandas as pd
+    from triton.testing import Benchmark
+def run_and_print(mark, save_path=None):
+    """Run a triton ``Mark`` (from ``perf_report``) and print/save results.
+    Each runner is expected to return a ``dict[str, Any]`` mapping stat name to
+    value, e.g. ``{"ms": 0.123, "GB/s": 1234}``. All providers in a benchmark
+    must return the same set of keys. Values are written through unchanged --
+    rounding/formatting is the caller's responsibility.
+    Output columns are ``x_names + [f"{line_name} ({stat})" for ...]``.
+    """
+    benchmarks = mark.benchmarks if isinstance(mark.benchmarks, list) else [mark.benchmarks]
+    for bench in benchmarks:
+        df = _run_one(mark.fn, bench)
+        print(bench.plot_name + ":")
+        print(df.to_string())
+        if save_path:
+            os.makedirs(save_path, exist_ok=True)
+            df.to_csv(os.path.join(save_path, f"{bench.plot_name}.csv"), index=False)
+def _run_one(fn, bench: Benchmark) -> pd.DataFrame:
+    try:
+        import pandas as pd
+    except ImportError as e:
+        raise ImportError(
+            "pandas is required to format benchmark results. "
+            "Install it with `pip install pandas` or `pip install -e '.[bench]'`."
+        ) from e
+    x_names = list(bench.x_names)
+    rows = []
+    stat_keys = None  # locked in from the first runner result
+    for x in bench.x_vals:
+        if not isinstance(x, (list, tuple)):
+            x = [x] * len(x_names)
+        x_args = dict(zip(x_names, x))
+        row = list(x)
+        for line_val in bench.line_vals:
+            stats = fn(**x_args, **{bench.line_arg: line_val}, **bench.args)
+            if not isinstance(stats, dict):
+                raise TypeError(f"runner must return dict[str, Any], got {type(stats).__name__}")
+            if stat_keys is None:
+                stat_keys = list(stats.keys())
+            elif list(stats.keys()) != stat_keys:
+                raise ValueError(f"runner returned keys {list(stats.keys())}, expected {stat_keys}")
+            row.extend(stats[k] for k in stat_keys)
+        rows.append(row)
+    cols = list(x_names) + [
+        f"{name} ({stat})" for name in bench.line_names for stat in (stat_keys or [])
+    ]
+    return pd.DataFrame(rows, columns=cols)
+def _bench_cuda_graph_l2_rotate(
+    fn,
+    arg_sets,
+    kwarg_sets,
+    extra_kwargs,
+    warmup_target_ms: float = 200.0,
+    n_timed_calls: int = 200,
+    quantiles=None,
+):
+    """L2-cold single-replay CUDA-graph benchmark.
+    Warmup is time-based: probe a single kernel launch to estimate the
+    per-call cost, then iterate round-robin over the pre-cloned
+    ``(arg_sets[i], kwarg_sets[i])`` pairs as a plain Python loop for
+    ``warmup_target_ms`` of wall-clock GPU work. Heavy pipelined configs
+    (TMA + smem_stages=3) need enough warmup to drain the pipeline-fill
+    phase or the timed window catches them artificially fast - a fixed
+    count of warmup launches underwarms heavy configs while overpaying
+    for cheap ones.
+    The timed window is a single ``graph.replay()`` of a captured CUDA
+    graph whose body records ``n_timed_calls`` round-robin invocations -
+    no Python loop, no per-launch CPU overhead - so the measurement is
+    just GPU work / total recorded calls.
+    Round-robin over fresh tensor sets defeats the L2-resident caching
+    that inflates short-kernel timing under ``triton.testing.do_bench``
+    (which calls the kernel on the same single tensor each iteration).
+    For cache-cold production workloads the round-robin number predicts
+    real-world latency; the L2-hot number favours wider layouts / deeper
+    smem stages that don't actually win once data has to come from HBM.
+    ``fn`` is called as ``fn(*arg_sets[i], **kwarg_sets[i], **extra_kwargs)``
+    once per recorded launch. ``extra_kwargs`` holds the per-config kwargs
+    (e.g. ``{"config": <RmsNormBwdConfig>}``) which don't need cloning;
+    keys in ``extra_kwargs`` must not overlap with ``kwarg_sets[i]``.
+    Returns ms/call, or a ``len(quantiles)``-list replicating that value
+    when ``quantiles`` is provided (for API parity with
+    ``triton.testing.do_bench``).
+    """
+    n_sets = len(arg_sets)
+    # Round timed-call count to a multiple of n_sets for even L2 turnover.
+    rounds_timed = max(1, n_timed_calls // n_sets)
+    total_timed_calls = rounds_timed * n_sets
+    # A few priming launches so the probe doesn't catch first-launch
+    # driver / kernel-load overhead.
+    for _ in range(3):
+        fn(*arg_sets[0], **kwarg_sets[0], **extra_kwargs)
+    torch.cuda.synchronize()
+    # Probe a single launch to estimate per-call ms; size the warmup loop
+    # to hit ``warmup_target_ms`` of GPU work regardless of kernel cost.
+    probe_start = torch.cuda.Event(enable_timing=True)
+    probe_end = torch.cuda.Event(enable_timing=True)
+    probe_start.record()
+    fn(*arg_sets[0], **kwarg_sets[0], **extra_kwargs)
+    probe_end.record()
+    torch.cuda.synchronize()
+    est_ms = max(probe_start.elapsed_time(probe_end), 1e-3)
+    n_warmup_calls = max(50, int(warmup_target_ms / est_ms))
+    # Warmup: plain Python loop over rotating sets, no graph capture.
+    for i in range(n_warmup_calls):
+        idx = i % n_sets
+        fn(*arg_sets[idx], **kwarg_sets[idx], **extra_kwargs)
+    torch.cuda.synchronize()
+    # Capture timed graph: a single replay covers all timed kernel launches.
+    timed_graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(timed_graph):
+        for _ in range(rounds_timed):
+            for i in range(n_sets):
+                fn(*arg_sets[i], **kwarg_sets[i], **extra_kwargs)
+    torch.cuda.synchronize()
+    start_evt = torch.cuda.Event(enable_timing=True)
+    end_evt = torch.cuda.Event(enable_timing=True)
+    start_evt.record()
+    timed_graph.replay()
+    end_evt.record()
+    torch.cuda.synchronize()
+    ms = start_evt.elapsed_time(end_evt) / total_timed_calls
+    if quantiles:
+        return [ms for _ in quantiles]
+    return ms
+def _clone_l2_rotate_inputs(args, kwargs, n_buffers: int):
+    """Clone tensor args AND tensor kwargs ``n_buffers`` times.
+    Returns ``(arg_sets, kwarg_sets)``: ``arg_sets[i]`` is a tuple matching
+    ``args``' positional shape with tensors cloned to fresh memory;
+    ``kwarg_sets[i]`` is a dict matching ``kwargs``' keys with tensor values
+    cloned. Non-tensor values (ints, strings, None, dataclasses, etc.) are
+    shared across all sets (no clone).
+    Both args and kwargs are cloned so that every recorded launch in the
+    L2-cold round-robin CUDA graph touches distinct GMEM addresses,
+    including for write-target kwargs like ``dw_partial`` / ``dx``.
+    """
+    arg_sets = []
+    kwarg_sets = []
+    for _ in range(n_buffers):
+        arg_sets.append(tuple(a.clone() if isinstance(a, Tensor) else a for a in args))
+        kwarg_sets.append({k: v.clone() if isinstance(v, Tensor) else v for k, v in kwargs.items()})
+    return arg_sets, kwarg_sets
+def _pick_l2_rotate_count(
+    args, kwargs, target_ratio: int = 3, min_buffers: int = 4, max_buffers: int = 16
+):
+    """Pick ``n_bufs`` so cloned input bytes per round exceed
+    ``target_ratio * L2_size`` (defeats L2 reuse), capped by HBM headroom and
+    [min_buffers, max_buffers]. Counts tensor bytes across both ``args`` and
+    ``kwargs`` so write-target kwargs (``dw_partial``, ``dx``, etc.) are
+    included in the L2-turnover calculation.
+    """
+    if not torch.cuda.is_available():
+        return min_buffers
+    tensor_bytes = sum(a.numel() * a.element_size() for a in args if isinstance(a, Tensor)) + sum(
+        v.numel() * v.element_size() for v in kwargs.values() if isinstance(v, Tensor)
+    )
+    if tensor_bytes == 0:
+        return min_buffers
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    l2_size = props.L2_cache_size
+    n_by_l2 = (target_ratio * l2_size + tensor_bytes - 1) // tensor_bytes
+    free_bytes, _ = torch.cuda.mem_get_info()
+    # Leave half of free memory headroom for the kernel's own scratch + the
+    # user's other allocations.
+    n_by_mem = max(1, int(free_bytes * 0.5) // tensor_bytes)
+    return max(min_buffers, min(max_buffers, min(n_by_l2, n_by_mem)))

build/torch-cuda/quack/blockscaled/__init__.py ADDED Viewed

	@@ -0,0 +1,31 @@

+# Copyright (c) 2026, Tri Dao.
+"""Blockscaled (MXFP8 / MXFP4 / NVFP4) GEMM support.
+- :mod:`quack.blockscaled.quantize` — pure-PyTorch quantizers (ported from
+  torchao) with torch.compile'd fast paths.
+- :mod:`quack.blockscaled.utils` — scale-factor packing/unpacking, operand
+  builders for tests/benchmarks, and the kernel-level compile path.
+The GEMM entry points live in :mod:`quack.gemm_interface` (pass ``(A, SFA)`` /
+``(B, SFB)`` tuples).
+"""
+from .quantize import (  # noqa: F401
+    nvfp4_per_tensor_scale,
+    to_blocked,
+    to_mx,
+    to_mx_compiled,
+    to_mxfp4,
+    to_mxfp4_compiled,
+    to_nvfp4,
+    to_nvfp4_compiled,
+)
+from .utils import (  # noqa: F401
+    BLOCKSCALED_FORMATS,
+    blockscaled_gemm_reference,
+    blockscaled_quantize,
+    dequant_operand,
+    pack_scale_2d_to_blocked_contig,
+    scale_blocked_for_cublas,
+    unpack_scale_blocked_to_2d,
+)

build/torch-cuda/quack/{mx_utils.py → blockscaled/quantize.py} RENAMED Viewed

File without changes

build/torch-cuda/quack/{blockscaled_gemm_utils.py → blockscaled/utils.py} RENAMED Viewed

@@ -9,16 +9,16 @@ import torch
 import cutlass
 import cutlass.cute as cute
-from .compile_utils import make_fake_tensor as fake_tensor
-from .cute_dsl_utils import get_device_capacity, get_max_active_clusters
-from .gemm_default_epi import GemmDefaultSm100
-from .gemm_tvm_ffi_utils import div_for_dtype, make_scheduler_args
-from .mx_utils import (
     to_mx_compiled,
     to_mxfp4_compiled,
     to_nvfp4_compiled,
 )
-from .varlen_utils import VarlenArguments
 TORCH_DTYPE_MAP = {
@@ -188,7 +188,7 @@ def create_blockscaled_operand_tensor(
 def _pack_blockscaled_scales(ref_blocks: torch.Tensor) -> torch.Tensor:
-    """Rearrange (mn, sf_k, l) scales into the (l, rm, rk, 512) blocked layout."""
     mn, sf_k, l = ref_blocks.shape
     rm = ceil_div(mn, 128)
     rk = ceil_div(sf_k, 4)
@@ -205,7 +205,7 @@ def _pack_blockscaled_scales(ref_blocks: torch.Tensor) -> torch.Tensor:
         k_idx[None, :, None] // 4,
         l_idx[None, None, :],
     ] = ref_blocks
-    return packed_6d.view(l, rm, rk, 512)
 def create_blockscaled_scale_tensor(
@@ -237,10 +237,10 @@ def create_blockscaled_scale_tensor(
 def pack_scale_2d_to_blocked_contig(scale_2d: torch.Tensor) -> torch.Tensor:
     """Rearrange a (l, mn, sf_k) or (mn, sf_k) e8m0 scale tensor into the
-    contiguous (l, rm, rk, 512) blocked layout shared by the quack kernel and
-    cuBLAS's block-scaling. Each 512 B inner block holds one 128 MN × 4 K
-    swizzled tile. Pads `mn` to a multiple of 128 and `sf_k` to a multiple of
-    4 with zeros."""
     if scale_2d.dim() == 2:
         scale_2d = scale_2d.unsqueeze(0)
     assert scale_2d.dim() == 3, f"expected (l, mn, sf_k), got shape {tuple(scale_2d.shape)}"
@@ -260,22 +260,92 @@ def pack_scale_2d_to_blocked_contig(scale_2d: torch.Tensor) -> torch.Tensor:
     blocks = padded.view(l, rm, 128, rk, 4).permute(0, 1, 3, 2, 4)
     # split 128 into (4 outer, 32 inner), then swap to (32, 4)
     blocks = blocks.reshape(l, rm, rk, 4, 32, 4).transpose(3, 4).contiguous()
-    return blocks.view(l, rm, rk, 512).view(orig_dtype)
 def scale_view_for_kernel(scale_contig: torch.Tensor, mn: int, sf_k: int, l: int) -> torch.Tensor:
-    """Validate a (l, rm, rk, 512) scale tensor and return it unchanged.
-    Only the innermost 512-B tile must be contiguous (stride 1, size 512);
-    outer (L, rm, rk) strides are free — the kernel reads them from the
-    passed tensor. This lets callers pass a slice/view of a larger buffer
-    with no extra copy. Works for both E8M0 (MX) and E4M3 (NVFP4)."""
     rm = ceil_div(mn, 128)
     rk = ceil_div(sf_k, 4)
-    assert scale_contig.shape == (l, rm, rk, 512), (
-        f"expected (l, rm, rk, 512) = ({l}, {rm}, {rk}, 512), got {tuple(scale_contig.shape)}"
     )
-    assert scale_contig.stride(-1) == 1, (
-        f"innermost 512-B dim must be unit-stride, got stride {scale_contig.stride(-1)}"
     )
     return scale_contig
@@ -283,9 +353,9 @@ def scale_view_for_kernel(scale_contig: torch.Tensor, mn: int, sf_k: int, l: int
 def scale_blocked_for_cublas(
     scale_contig: torch.Tensor, mn: int, sf_k: int, l_idx: int = 0
 ) -> torch.Tensor:
-    """Flatten a (l, rm, rk, 512) scale tensor to the 1D swizzled layout
     torch._scaled_mm expects. Uses a single l slice."""
-    assert scale_contig.is_contiguous() and scale_contig.dim() == 4
     return scale_contig[l_idx].reshape(-1)
@@ -328,8 +398,8 @@ def create_blockscaled_operand_quantized(
     ref:   (mn, k, l) float32 dequantized reference
     q_mkl: (mn, k, l) operand tensor in the layout the quack kernel consumes
            (float8_e4m3fn for fp8 formats; int8 with packed nibbles for fp4)
-    scale_contig: (l, rm, rk, 512) contiguous scale storage. Each 512 B
-           inner block is one 128 MN × 4 K swizzled tile. Byte layout matches
            cuBLAS `to_blocked`. Pass directly to the quack kernel, or use
            `scale_blocked_for_cublas` for cuBLAS.
     """
@@ -403,7 +473,7 @@ def create_blockscaled_varlen_m_operands(
     """Generate bf16 randn + quantize for a varlen_m blockscaled GEMM.
     Per-expert seqlens may be arbitrary (not required to be multiples of 128).
-    SF is stored in dQaccum-style padded format: each expert `i`'s scales
     occupy `ceildiv(m_i, 128) * 128` rows at offset
     `(cu_seqlens_m[i] + i * 128) // 128 * 128` in the padded scale buffer.
     The kernel decodes via `VarlenManager.offset_batch_SFA` which applies the
@@ -414,10 +484,14 @@ def create_blockscaled_varlen_m_operands(
       b_ref: (num_experts, n, k) fp32 dequantized
       qa:   (total_m, k) 2D K-major quantized operand (fp8) or (total_m, k/2) (fp4)
       qb:   (n, k, num_experts) 3D K-major quantized operand (fp8) or (n, k/2, num_experts) (fp4)
-      a_sc_contig: (1, total_padded_rm, rk, 512) — dQaccum-padded SFA.
         total_padded_rm = ((total_m + num_experts * 128) // 128).
-      b_sc_contig: (num_experts, rn, rk, 512) — regular per-expert SFB.
       cu_seqlens_m: (num_experts+1,) int32
     """
     assert k % sf_vec_size == 0
     if seqlens_m is None:
@@ -429,23 +503,31 @@ def create_blockscaled_varlen_m_operands(
     std = randn_std if randn_std is not None else k**-0.5
     sf_k = k // sf_vec_size
-    if ab_dtype == cutlass.Float8E4M3FN and sf_dtype == cutlass.Float8E8M0FNU and sf_vec_size == 32:
-        from .mx_utils import to_mx_compiled
-        to_fn = to_mx_compiled
-    else:
-        raise NotImplementedError(
-            f"varlen_m currently only supports MXFP8 (got ab={ab_dtype}, sf={sf_dtype}, vec={sf_vec_size}). "
-            "FP4 support pending."
-        )
-    # Quantize A: (total_m, k) bf16 -> (total_m, k) fp8 K-major.
     # A data itself is stored packed (no per-expert padding); only SFA is padded.
     a_hp = (torch.randn(total_m, k, dtype=torch.bfloat16, device="cuda") * std).contiguous()
-    qa, sa_2d = to_fn(a_hp, sf_vec_size)  # (total_m, k), (total_m, sf_k)
-    a_ref = qa.float() * sa_2d.float().repeat_interleave(sf_vec_size, dim=-1)
-    # Build padded SFA storage (dQaccum format). Each expert's m_i rows of
     # scales are written at padded tile offset `cu_seqlens[i] // 128 + i`.
     # Allocation: `ceildiv(total_m, 128) + (L - 1)` tiles — proven sufficient
     # in AI/varlen_blockscaled_sf_layout.md (proof 2's "tighter alternative").
@@ -462,24 +544,22 @@ def create_blockscaled_varlen_m_operands(
         offset += m_i
     a_sc_contig = pack_scale_2d_to_blocked_contig(sa_2d_padded.view(1, total_padded_m, sf_k))
-    # Quantize B: (num_experts, n, k) bf16 -> (n, k, num_experts). b_major selects
-    # k-major (stride (k, 1, n*k)) or n-major (stride (1, n, n*k)).
     assert b_major in ("k", "n"), f"b_major must be 'k' or 'n', got {b_major!r}"
     b_hp = (torch.randn(num_experts, n, k, dtype=torch.bfloat16, device="cuda") * std).contiguous()
-    qb_flat, sb_2d = to_fn(b_hp.view(num_experts * n, k), sf_vec_size)
     if b_major == "k":
         qb = (
-            qb_flat.view(num_experts, n, k).contiguous().permute(1, 2, 0)
-        )  # (n, k, l) stride (k, 1, n*k)
     else:
         qb = (
-            qb_flat.view(num_experts, n, k).transpose(1, 2).contiguous().permute(2, 1, 0)
         )  # (n, k, l) stride (1, n, n*k)
-    sb_2d = sb_2d.view(num_experts, n, sf_k)
-    b_sc_contig = pack_scale_2d_to_blocked_contig(sb_2d)
-    b_ref = qb_flat.float().view(num_experts, n, k) * sb_2d.float().repeat_interleave(
-        sf_vec_size, dim=-1
-    )
     cu_seqlens_m = torch.tensor(
         [0] + list(itertools.accumulate(seqlens_m)), dtype=torch.int32, device="cuda"
@@ -498,23 +578,34 @@ def create_blockscaled_varlen_k_operands(
     *,
     randn_std: Optional[float] = None,
     seqlens_k: Optional[list] = None,
 ):
     """Generate bf16 randn + quantize for a varlen_k blockscaled GEMM.
-    Per-expert `k_i` must be a multiple of `sf_vec_size` (quantization chunk)
-    but NOT necessarily a multiple of `sf_vec_size * 4` (= 128 for MXFP8).
-    The SF buffer uses dQaccum-style K padding: each expert `i`'s scales occupy
     `ceildiv(k_i, 128) * 128` bytes worth of K at offset
     `(cu_seqlens_k[i] + i * 128) // 128 * 128` (in source-K units). A and B
     operand data stay packed and unpadded along K — only their SF buffers pad.
     Returns (a_ref_list, b_ref_list, qa, qb, a_sc_contig, b_sc_contig, cu_seqlens_k):
       a_ref_list: list of per-expert (m, k_i) fp32 dequantized A.
       b_ref_list: list of per-expert (n, k_i) fp32 dequantized B.
       qa:  (m, total_k) K-major fp8 (stride (total_k, 1)).
       qb:  (n, total_k) K-major fp8 (stride (total_k, 1)).
-      a_sc_contig: (1, rm, total_padded_rk, 512) dQaccum-padded SFA.
-      b_sc_contig: (1, rn, total_padded_rk, 512) dQaccum-padded SFB.
       cu_seqlens_k: (num_experts+1,) int32.
     """
     if not (
@@ -530,30 +621,37 @@ def create_blockscaled_varlen_k_operands(
         f"seqlens_k length {len(seqlens_k)} != num_experts {num_experts}"
     )
     for i, k_i in enumerate(seqlens_k):
-        assert k_i % sf_vec_size == 0, (
-            f"seqlens_k[{i}]={k_i} must be divisible by sf_vec_size={sf_vec_size}"
-        )
     total_k = int(sum(seqlens_k))
     std = randn_std if randn_std is not None else (max(seqlens_k)) ** -0.5
-    sf_k_total = total_k // sf_vec_size
-    from .mx_utils import to_mx_compiled
     a_q_list, a_sc_list, a_ref_list = [], [], []
     b_q_list, b_sc_list, b_ref_list = [], [], []
     for k_i in seqlens_k:
-        # A slice: (m, k_i) bf16 -> fp8, scales (m, k_i // sf_vec_size).
-        a_hp = (torch.randn(m, k_i, dtype=torch.bfloat16, device="cuda") * std).contiguous()
-        a_q, a_sc = to_mx_compiled(a_hp, sf_vec_size)
         a_q_list.append(a_q)
         a_sc_list.append(a_sc)
-        a_ref_list.append(a_q.float() * a_sc.float().repeat_interleave(sf_vec_size, dim=-1))
-        b_hp = (torch.randn(n, k_i, dtype=torch.bfloat16, device="cuda") * std).contiguous()
-        b_q, b_sc = to_mx_compiled(b_hp, sf_vec_size)
         b_q_list.append(b_q)
         b_sc_list.append(b_sc)
-        b_ref_list.append(b_q.float() * b_sc.float().repeat_interleave(sf_vec_size, dim=-1))
     # Pack operand data along K: (m, total_k), (n, total_k). varlen_k's
     # ragged TMA descriptors are built for MN-major operands (stride 1 on
@@ -572,11 +670,15 @@ def create_blockscaled_varlen_k_operands(
     total_padded_rk = (total_k + tile - 1) // tile + (num_experts - 1)
     total_padded_k = total_padded_rk * tile
     total_padded_sf_k = total_padded_k // sf_vec_size
-    sa_2d_padded = torch.zeros(m, total_padded_sf_k, dtype=a_sc_list[0].dtype, device="cuda")
-    sb_2d_padded = torch.zeros(n, total_padded_sf_k, dtype=b_sc_list[0].dtype, device="cuda")
     k_offset = 0
     for i, k_i in enumerate(seqlens_k):
-        sf_k_i = k_i // sf_vec_size
         k_offset_padded = (k_offset // tile + i) * tile
         sf_k_offset_padded = k_offset_padded // sf_vec_size
         sa_2d_padded[:, sf_k_offset_padded : sf_k_offset_padded + sf_k_i] = a_sc_list[i]
@@ -635,16 +737,20 @@ def compile_blockscaled_gemm_tvm_ffi(
     )
     stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
-    from .gemm_tvm_ffi_utils import make_fake_varlen_args
     varlen_args_fake = make_fake_varlen_args(varlen_m, varlen_k, False, None) or VarlenArguments()
     # Fake operand tensors with sym_ints (varlen-aware shapes).
     if varlen_m:
         total_m_sym = cute.sym_int()
-        n_sym, k_sym, l_sym = cute.sym_int(), cute.sym_int(), cute.sym_int()
-        # Detect each operand's leading (stride-1) dim so m-major A / n-major B
-        # are accepted for varlen_m (MXFP8 only — fp4 is rejected upstream).
         fake_mA = fake_tensor(
             ab_dtype,
             (total_m_sym, k_sym),
@@ -709,7 +815,7 @@ def compile_blockscaled_gemm_tvm_ffi(
         varlen_args,
         stream,
     ):
-        gemm(a, b, d, None, compile_epi_args, scheduler_args, varlen_args, stream, sfa, sfb, None)
     compiled = cute.compile(
         runner,

 import cutlass
 import cutlass.cute as cute
+from ..compile_utils import make_fake_tensor as fake_tensor
+from ..cute_dsl_utils import get_device_capacity, get_max_active_clusters
+from ..gemm_default_epi import GemmDefaultSm100
+from ..gemm_tvm_ffi_utils import div_for_dtype, make_scheduler_args
+from .quantize import (
     to_mx_compiled,
     to_mxfp4_compiled,
     to_nvfp4_compiled,
 )
+from ..varlen_utils import VarlenArguments
 TORCH_DTYPE_MAP = {
 def _pack_blockscaled_scales(ref_blocks: torch.Tensor) -> torch.Tensor:
+    """Rearrange (mn, sf_k, l) scales into the (l, rm, rk, 32, 4, 4) blocked layout."""
     mn, sf_k, l = ref_blocks.shape
     rm = ceil_div(mn, 128)
     rk = ceil_div(sf_k, 4)
         k_idx[None, :, None] // 4,
         l_idx[None, None, :],
     ] = ref_blocks
+    return packed_6d
 def create_blockscaled_scale_tensor(
 def pack_scale_2d_to_blocked_contig(scale_2d: torch.Tensor) -> torch.Tensor:
     """Rearrange a (l, mn, sf_k) or (mn, sf_k) e8m0 scale tensor into the
+    contiguous (l, rm, rk, 32, 4, 4) blocked layout shared by the quack kernel
+    and cuBLAS's block-scaling. Each inner (32, 4, 4) atom (512 B) holds one
+    128 MN × 4 K swizzled tile. Pads `mn` to a multiple of 128 and `sf_k` to a
+    multiple of 4 with zeros."""
     if scale_2d.dim() == 2:
         scale_2d = scale_2d.unsqueeze(0)
     assert scale_2d.dim() == 3, f"expected (l, mn, sf_k), got shape {tuple(scale_2d.shape)}"
     blocks = padded.view(l, rm, 128, rk, 4).permute(0, 1, 3, 2, 4)
     # split 128 into (4 outer, 32 inner), then swap to (32, 4)
     blocks = blocks.reshape(l, rm, rk, 4, 32, 4).transpose(3, 4).contiguous()
+    return blocks.view(orig_dtype)
+def unpack_scale_blocked_to_2d(blocked: torch.Tensor, mn: int, sf_k: int) -> torch.Tensor:
+    """Unswizzle (l, rm, rk, 32, 4, 4) blocked scale factors to (l, mn, sf_k)."""
+    l, rm, rk = blocked.shape[:3]
+    assert tuple(blocked.shape[3:]) == (32, 4, 4)
+    orig_dtype = blocked.dtype
+    u8 = blocked.view(torch.uint8)
+    # (32=m%32, 4=m//32, 4=k%4) -> (4, 32, 4) -> (l, rm, rk, 128, 4) -> (l, mn_pad, sf_k_pad)
+    u8 = u8.transpose(3, 4).reshape(l, rm, rk, 128, 4)
+    u8 = u8.permute(0, 1, 3, 2, 4).reshape(l, rm * 128, rk * 4)
+    return u8[:, :mn, :sf_k].contiguous().view(orig_dtype)
+def dequant_operand(x: torch.Tensor) -> torch.Tensor:
+    """Dequantize an operand tensor to float32 values (without scale factors).
+    fp8 tensors convert directly; ``float4_e2m1fn_x2`` tensors unpack two codes
+    per byte (low nibble = even K, high nibble = odd K), doubling the last dim.
+    """
+    if x.dtype == torch.float4_e2m1fn_x2:
+        u8 = x.view(torch.uint8)
+        lo = _fp4_unpacked_to_value(u8 & 0x0F)
+        hi = _fp4_unpacked_to_value((u8 >> 4) & 0x0F)
+        return torch.stack([lo, hi], dim=-1).reshape(*x.shape[:-1], x.shape[-1] * 2)
+    return x.float()
+BLOCKSCALED_FORMATS = {
+    # format: (torch operand dtype, torch SF dtype, sf_vec_size)
+    "mxfp8": (torch.float8_e4m3fn, torch.float8_e8m0fnu, 32),
+    "mxfp4": (torch.float4_e2m1fn_x2, torch.float8_e8m0fnu, 32),
+    "nvfp4": (torch.float4_e2m1fn_x2, torch.float8_e4m3fn, 16),
+}
+def blockscaled_quantize(
+    x: torch.Tensor, format: str = "mxfp8", per_tensor_scale: Optional[torch.Tensor] = None
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Quantize a (M, K) or (L, M, K) bf16/fp32 tensor along K for blockscaled GEMM.
+    Returns ``(q, sf)`` ready to pass as an ``(A, SFA)`` / ``(B, SFB)`` tuple to
+    :func:`quack.gemm_interface.gemm`:
+      q:  same leading shape as ``x``; fp8 for mxfp8 (M, K), packed fp4x2 for
+          mxfp4/nvfp4 (M, K/2), K-contiguous.
+      sf: blocked scale factors, (rm, rk, 32, 4, 4) or (L, rm, rk, 32, 4, 4).
+    For nvfp4, ``per_tensor_scale`` (scalar fp32) folds the global scale; pass the
+    product of A's and B's per-tensor scales as ``alpha`` to the GEMM.
+    """
+    from .quantize import to_mx_compiled, to_mxfp4_compiled, to_nvfp4_compiled
+    assert format in BLOCKSCALED_FORMATS, f"unknown blockscaled format: {format}"
+    q_dtype, sf_dtype, sf_vec = BLOCKSCALED_FORMATS[format]
+    assert x.shape[-1] % sf_vec == 0, f"K ({x.shape[-1]}) must be divisible by {sf_vec}"
+    batched = x.ndim == 3
+    l, mn, k = x.shape if batched else (1, *x.shape)
+    x_flat = x.reshape(l * mn, k)
+    if format == "mxfp8":
+        q, sc = to_mx_compiled(x_flat, sf_vec)
+    elif format == "mxfp4":
+        q, sc = to_mxfp4_compiled(x_flat, sf_vec)
+    else:
+        q, sc, _ = to_nvfp4_compiled(x_flat, sf_vec, per_tensor_scale)
+    q = q.view(torch.uint8).view(q_dtype) if q_dtype == torch.float4_e2m1fn_x2 else q
+    q = q.reshape(*x.shape[:-1], -1)
+    sf = pack_scale_2d_to_blocked_contig(sc.view(l, mn, k // sf_vec))
+    return q, sf if batched else sf.squeeze(0)
 def scale_view_for_kernel(scale_contig: torch.Tensor, mn: int, sf_k: int, l: int) -> torch.Tensor:
+    """Validate a (l, rm, rk, 32, 4, 4) scale tensor and return it unchanged.
+    Only the innermost (32, 4, 4) atom (one 512 B tile) must be contiguous
+    (strides (16, 4, 1)); outer (L, rm, rk) strides are free — the kernel
+    reads them from the passed tensor. This lets callers pass a slice/view of
+    a larger buffer with no extra copy. Works for both E8M0 (MX) and E4M3
+    (NVFP4)."""
     rm = ceil_div(mn, 128)
     rk = ceil_div(sf_k, 4)
+    assert scale_contig.shape == (l, rm, rk, 32, 4, 4), (
+        f"expected (l, rm, rk, 32, 4, 4) = ({l}, {rm}, {rk}, 32, 4, 4), "
+        f"got {tuple(scale_contig.shape)}"
     )
+    assert scale_contig.stride()[-3:] == (16, 4, 1), (
+        f"inner (32, 4, 4) atom must be contiguous with strides (16, 4, 1), "
+        f"got {scale_contig.stride()[-3:]}"
     )
     return scale_contig
 def scale_blocked_for_cublas(
     scale_contig: torch.Tensor, mn: int, sf_k: int, l_idx: int = 0
 ) -> torch.Tensor:
+    """Flatten a (l, rm, rk, 32, 4, 4) scale tensor to the 1D swizzled layout
     torch._scaled_mm expects. Uses a single l slice."""
+    assert scale_contig.is_contiguous() and scale_contig.dim() == 6
     return scale_contig[l_idx].reshape(-1)
     ref:   (mn, k, l) float32 dequantized reference
     q_mkl: (mn, k, l) operand tensor in the layout the quack kernel consumes
            (float8_e4m3fn for fp8 formats; int8 with packed nibbles for fp4)
+    scale_contig: (l, rm, rk, 32, 4, 4) contiguous scale storage. Each inner
+           (32, 4, 4) atom (512 B) is one 128 MN × 4 K swizzled tile. Byte layout matches
            cuBLAS `to_blocked`. Pass directly to the quack kernel, or use
            `scale_blocked_for_cublas` for cuBLAS.
     """
     """Generate bf16 randn + quantize for a varlen_m blockscaled GEMM.
     Per-expert seqlens may be arbitrary (not required to be multiples of 128).
+    SF is stored with tile-aligned per-batch padding: each expert `i`'s scales
     occupy `ceildiv(m_i, 128) * 128` rows at offset
     `(cu_seqlens_m[i] + i * 128) // 128 * 128` in the padded scale buffer.
     The kernel decodes via `VarlenManager.offset_batch_SFA` which applies the
       b_ref: (num_experts, n, k) fp32 dequantized
       qa:   (total_m, k) 2D K-major quantized operand (fp8) or (total_m, k/2) (fp4)
       qb:   (n, k, num_experts) 3D K-major quantized operand (fp8) or (n, k/2, num_experts) (fp4)
+      a_sc_contig: (1, total_padded_rm, rk, 32, 4, 4) — M-padded SFA (tile-aligned per batch).
         total_padded_rm = ((total_m + num_experts * 128) // 128).
+      b_sc_contig: (num_experts, rn, rk, 32, 4, 4) — regular per-expert SFB.
       cu_seqlens_m: (num_experts+1,) int32
+    Supports MXFP8 / MXFP4 / NVFP4; fp4 formats require b_major="k" (tcgen05
+    MMA needs K-major fp4 operands). NVFP4 uses no per-tensor scale here (it
+    would just fold into alpha).
     """
     assert k % sf_vec_size == 0
     if seqlens_m is None:
     std = randn_std if randn_std is not None else k**-0.5
     sf_k = k // sf_vec_size
+    fmt = _blockscaled_format_of(ab_dtype, sf_dtype, sf_vec_size)
+    if fmt != "mxfp8":
+        assert b_major == "k", f"{fmt} requires K-major operands, got b_major={b_major!r}"
+    def quantize(x2d):
+        """(rows, k) bf16 -> (q, scale_2d, dequant_ref); q is fp8 (rows, k) or fp4x2 (rows, k/2)."""
+        if fmt == "mxfp8":
+            q, sc = to_mx_compiled(x2d, sf_vec_size)
+            vals = q.float()
+        else:
+            if fmt == "mxfp4":
+                q_packed, sc = to_mxfp4_compiled(x2d, sf_vec_size)
+            else:  # nvfp4
+                q_packed, sc, _ = to_nvfp4_compiled(x2d, sf_vec_size, None)
+            q = q_packed.view(torch.uint8).view(torch.float4_e2m1fn_x2)
+            vals = dequant_operand(q)
+        ref = vals * sc.float().repeat_interleave(sf_vec_size, dim=-1)
+        return q, sc, ref
+    # Quantize A: (total_m, k) bf16 -> (total_m, k[/2]) K-major.
     # A data itself is stored packed (no per-expert padding); only SFA is padded.
     a_hp = (torch.randn(total_m, k, dtype=torch.bfloat16, device="cuda") * std).contiguous()
+    qa, sa_2d, a_ref = quantize(a_hp)  # (total_m, k[/2]), (total_m, sf_k), (total_m, k)
+    # Build padded SFA storage (tile-aligned per-batch). Each expert's m_i rows of
     # scales are written at padded tile offset `cu_seqlens[i] // 128 + i`.
     # Allocation: `ceildiv(total_m, 128) + (L - 1)` tiles — proven sufficient
     # in AI/varlen_blockscaled_sf_layout.md (proof 2's "tighter alternative").
         offset += m_i
     a_sc_contig = pack_scale_2d_to_blocked_contig(sa_2d_padded.view(1, total_padded_m, sf_k))
+    # Quantize B: (num_experts, n, k) bf16 -> (n, k[/2], num_experts). b_major selects
+    # k-major (stride (kb, 1, n*kb)) or n-major (stride (1, n, n*k), mxfp8 only).
     assert b_major in ("k", "n"), f"b_major must be 'k' or 'n', got {b_major!r}"
     b_hp = (torch.randn(num_experts, n, k, dtype=torch.bfloat16, device="cuda") * std).contiguous()
+    qb_flat, sb_2d, b_ref_flat = quantize(b_hp.view(num_experts * n, k))
+    kb = qb_flat.shape[-1]  # k for fp8, k/2 for packed fp4
     if b_major == "k":
         qb = (
+            qb_flat.view(num_experts, n, kb).contiguous().permute(1, 2, 0)
+        )  # (n, kb, l) stride (kb, 1, n*kb)
     else:
         qb = (
+            qb_flat.view(num_experts, n, kb).transpose(1, 2).contiguous().permute(2, 1, 0)
         )  # (n, k, l) stride (1, n, n*k)
+    b_sc_contig = pack_scale_2d_to_blocked_contig(sb_2d.view(num_experts, n, sf_k))
+    b_ref = b_ref_flat.view(num_experts, n, k)
     cu_seqlens_m = torch.tensor(
         [0] + list(itertools.accumulate(seqlens_m)), dtype=torch.int32, device="cuda"
     *,
     randn_std: Optional[float] = None,
     seqlens_k: Optional[list] = None,
+    sf_pad_byte: int = 0,
 ):
     """Generate bf16 randn + quantize for a varlen_k blockscaled GEMM.
+    Per-expert `k_i` is arbitrary (any positive int): neither `sf_vec_size` nor
+    `sf_vec_size * 4` (= 128 for MXFP8) alignment is required. A non-multiple-of-32
+    `k_i` just means the expert's last scale block covers a partial chunk; the
+    kernel's ragged value TMA zero-fills beyond `cu_seqlens_k[i+1]`, so the tail
+    contributes exactly 0.
+    The SF buffer uses tile-aligned per-batch K padding: each expert `i`'s scales occupy
     `ceildiv(k_i, 128) * 128` bytes worth of K at offset
     `(cu_seqlens_k[i] + i * 128) // 128 * 128` (in source-K units). A and B
     operand data stay packed and unpadded along K — only their SF buffers pad.
+    SF pad regions inside each expert's last 512 B atom are loaded by the
+    kernel (TMA loads whole atom columns) but never consumed: the mma loop
+    skips the MMA instructions for pad k-blocks (one instruction per SF block
+    for mxfp8; see `GemmSm100.mma`), so the pad may hold arbitrary bytes —
+    including 0xFF (e8m0 NaN). `sf_pad_byte` sets the pad fill so tests can
+    poison it deliberately.
     Returns (a_ref_list, b_ref_list, qa, qb, a_sc_contig, b_sc_contig, cu_seqlens_k):
       a_ref_list: list of per-expert (m, k_i) fp32 dequantized A.
       b_ref_list: list of per-expert (n, k_i) fp32 dequantized B.
       qa:  (m, total_k) K-major fp8 (stride (total_k, 1)).
       qb:  (n, total_k) K-major fp8 (stride (total_k, 1)).
+      a_sc_contig: (1, rm, total_padded_rk, 32, 4, 4) K-padded SFA (tile-aligned per batch).
+      b_sc_contig: (1, rn, total_padded_rk, 32, 4, 4) K-padded SFB (tile-aligned per batch).
       cu_seqlens_k: (num_experts+1,) int32.
     """
     if not (
         f"seqlens_k length {len(seqlens_k)} != num_experts {num_experts}"
     )
     for i, k_i in enumerate(seqlens_k):
+        assert k_i > 0, f"seqlens_k[{i}]={k_i} must be positive"
     total_k = int(sum(seqlens_k))
     std = randn_std if randn_std is not None else (max(seqlens_k)) ** -0.5
+    from .quantize import to_mx_compiled
+    def quantize(mn, k_i):
+        # The quantizer reshapes K into sf_vec_size chunks, so zero-pad k_i up to a
+        # multiple of it; zeros never raise a chunk amax, so the real elements
+        # quantize identically. Values are sliced back to k_i; scales keep the
+        # ceil(k_i / sf_vec_size) blocks (the last one covers a partial chunk).
+        k_q = (k_i + sf_vec_size - 1) // sf_vec_size * sf_vec_size
+        hp = torch.zeros(mn, k_q, dtype=torch.bfloat16, device="cuda")
+        hp[:, :k_i] = torch.randn(mn, k_i, dtype=torch.bfloat16, device="cuda") * std
+        q, sc = to_mx_compiled(hp, sf_vec_size)
+        q = q[:, :k_i]
+        ref = q.float() * sc.float().repeat_interleave(sf_vec_size, dim=-1)[:, :k_i]
+        return q, sc, ref
     a_q_list, a_sc_list, a_ref_list = [], [], []
     b_q_list, b_sc_list, b_ref_list = [], [], []
     for k_i in seqlens_k:
+        a_q, a_sc, a_ref = quantize(m, k_i)
         a_q_list.append(a_q)
         a_sc_list.append(a_sc)
+        a_ref_list.append(a_ref)
+        b_q, b_sc, b_ref = quantize(n, k_i)
         b_q_list.append(b_q)
         b_sc_list.append(b_sc)
+        b_ref_list.append(b_ref)
     # Pack operand data along K: (m, total_k), (n, total_k). varlen_k's
     # ragged TMA descriptors are built for MN-major operands (stride 1 on
     total_padded_rk = (total_k + tile - 1) // tile + (num_experts - 1)
     total_padded_k = total_padded_rk * tile
     total_padded_sf_k = total_padded_k // sf_vec_size
+    sa_2d_padded = torch.full(
+        (m, total_padded_sf_k), sf_pad_byte, dtype=torch.uint8, device="cuda"
+    ).view(a_sc_list[0].dtype)
+    sb_2d_padded = torch.full(
+        (n, total_padded_sf_k), sf_pad_byte, dtype=torch.uint8, device="cuda"
+    ).view(b_sc_list[0].dtype)
     k_offset = 0
     for i, k_i in enumerate(seqlens_k):
+        sf_k_i = (k_i + sf_vec_size - 1) // sf_vec_size
         k_offset_padded = (k_offset // tile + i) * tile
         sf_k_offset_padded = k_offset_padded // sf_vec_size
         sa_2d_padded[:, sf_k_offset_padded : sf_k_offset_padded + sf_k_i] = a_sc_list[i]
     )
     stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+    from ..gemm_tvm_ffi_utils import make_fake_varlen_args
     varlen_args_fake = make_fake_varlen_args(varlen_m, varlen_k, False, None) or VarlenArguments()
     # Fake operand tensors with sym_ints (varlen-aware shapes).
     if varlen_m:
         total_m_sym = cute.sym_int()
+        n_sym, l_sym = cute.sym_int(), cute.sym_int()
+        # Sub-byte (fp4) operands need the contiguous K extent statically divisible
+        # by the packing factor; harmless for 8-bit dtypes.
+        k_sym = cute.sym_int(divisibility=div_for_dtype(ab_dtype) if ab_dtype.width < 8 else 1)
+        # Detect B's leading (stride-1) dim so n-major B is accepted for varlen_m
+        # (mxfp8 only; fp4 is always K-major). A must be K-major for varlen_m —
+        # the public API enforces this (see quack/gemm.py) for all dtypes.
         fake_mA = fake_tensor(
             ab_dtype,
             (total_m_sym, k_sym),
         varlen_args,
         stream,
     ):
+        gemm(a, b, d, None, compile_epi_args, scheduler_args, varlen_args, stream, sfa, sfb)
     compiled = cute.compile(
         runner,

build/torch-cuda/quack/broadcast_utils.py CHANGED Viewed

@@ -5,18 +5,17 @@ import cutlass
 import cutlass.cute as cute
 from cutlass import Float32, const_expr
-from .layout_utils import make_acc_tensor_mn_view
 @cute.jit
 def vec_op(tCrC: cute.Tensor, tCrVec: cute.Tensor, op: Callable, is_colvec: bool) -> None:
     if const_expr(tCrC.element_type != Float32):  # Convert to f32
-        tCrC_f32 = cute.make_rmem_tensor(tCrC.shape, Float32)
-        tCrC_f32.store(tCrC.load().to(Float32))
     else:
         tCrC_f32 = tCrC
     # this happens to work for frgA layout too, not just acc layout
-    tCrC_f32_mn = make_acc_tensor_mn_view(tCrC_f32)
     if const_expr(is_colvec):
         assert cute.size(tCrC_f32_mn, mode=[0]) == cute.size(tCrVec)
         for r in cutlass.range(cute.size(tCrC_f32_mn, mode=[0]), unroll_full=True):

 import cutlass.cute as cute
 from cutlass import Float32, const_expr
+from . import layout_utils
 @cute.jit
 def vec_op(tCrC: cute.Tensor, tCrVec: cute.Tensor, op: Callable, is_colvec: bool) -> None:
     if const_expr(tCrC.element_type != Float32):  # Convert to f32
+        tCrC_f32 = tCrC.to(Float32)
     else:
         tCrC_f32 = tCrC
     # this happens to work for frgA layout too, not just acc layout
+    tCrC_f32_mn = layout_utils.reshape_acc_to_mn(tCrC_f32)
     if const_expr(is_colvec):
         assert cute.size(tCrC_f32_mn, mode=[0]) == cute.size(tCrVec)
         for r in cutlass.range(cute.size(tCrC_f32_mn, mode=[0]), unroll_full=True):

build/torch-cuda/quack/cache/__init__.py ADDED Viewed

	@@ -0,0 +1,74 @@

+# Copyright (c) 2025-2026, Tri Dao.
+"""Persistent kernel-cache utilities for QuACK.
+Public API
+----------
+Persistent ``.o`` cache:
+* :func:`jit_cache` — decorator that wraps a kernel-compile function with
+  in-memory + persistent ``.o`` caching (see :mod:`quack.cache.jit`).
+* :data:`CACHE_ENABLED`, :data:`CACHE_DIR`, :data:`EXTRA_SOURCE_DIRS` —
+  static-config flags.
+* :class:`FileLock`, :func:`get_cache_path`, :class:`CacheInfo` —
+  supporting types.
+Async compilation (see :mod:`quack.cache.async_compile`):
+* :class:`CompilePending` — raised by ``jit_cache`` on a cold miss while a
+  compile pool is active; the caller defers and retries once the ``.o``
+  lands.
+* :func:`pool_scope` — activate a compile pool for a scoped block (used by
+  the autotuner's bench loop).
+CRITICAL ORDERING: the static-config flags below MUST be defined before the
+``from quack.cache.jit import ...`` block. ``quack/cache/jit.py`` does
+``import quack.cache as _state`` at its module top; Python returns the
+partially-initialized package object, and lookups inside ``jit_cache``'s
+wrapper rely on these names already existing at that checkpoint. Reordering
+the imports here, even via an auto-formatter, will break the first kernel
+compile with ``AttributeError``. The defensive unit tests in
+``tests/test_cache.py`` exercise this path end-to-end.
+"""
+from __future__ import annotations
+import os
+from pathlib import Path
+from typing import List, Optional
+CACHE_ENABLED: bool = os.getenv("QUACK_CACHE_ENABLED", "1") == "1"
+CACHE_DIR: Optional[str] = os.getenv("QUACK_CACHE_DIR", None)
+#: Downstream projects can append directories here to include their sources
+#: in the cache fingerprint. Must be set before the first jit_cache call.
+EXTRA_SOURCE_DIRS: List[Path] = []
+# ---------------------------------------------------------------------------
+# Public API surface. Imported AFTER the flags are defined.
+# ---------------------------------------------------------------------------
+from .jit import (  # noqa: E402
+    EXPORT_FUNC_NAME,
+    LOCK_TIMEOUT,
+    CacheInfo,
+    FileLock,
+    get_cache_path,
+    jit_cache,
+)
+from .async_compile import (  # noqa: E402
+    CompilePending,
+    pool_scope,
+)
+__all__ = [
+    # Persistent .o cache.
+    "jit_cache",
+    "CacheInfo",
+    "EXPORT_FUNC_NAME",
+    "LOCK_TIMEOUT",
+    "FileLock",
+    "get_cache_path",
+    # Async compilation.
+    "CompilePending",
+    "pool_scope",
+]

build/torch-cuda/quack/cache/_pool_preload.py ADDED Viewed

	@@ -0,0 +1,50 @@

+# Copyright (c) 2026, Tri Dao.
+"""Forkserver preload module for the async compile pool.
+Imported once inside the multiprocessing *forkserver* process (see
+``multiprocessing.set_forkserver_preload``). Every pool worker is then
+``fork()``-ed from that warm process and inherits the imported interpreter
+state via copy-on-write: worker startup drops from ~13 s (torch 4 s +
+cutlass/cute/tvm_ffi 9 s per spawn) to ~0.1 s per fork.
+This is the same architecture as PyTorch Inductor's compile-worker
+``SubprocPool``: one sidecar pays the import, workers fork from it.
+Fork-safety: nothing here may initialize CUDA (a forked child of a
+CUDA-initialized process is undefined behavior). Importing torch and
+cutlass does not create a CUDA context; workers additionally run with
+``CUDA_VISIBLE_DEVICES=""`` + ``QUACK_ARCH``/``CUTE_DSL_ARCH`` overrides so
+the compile path never touches the driver (the same mechanism the CPU-only
+compile workflow uses).
+"""
+import os
+import subprocess
+# Pin the target arch BEFORE importing quack: import-time code paths (e.g.
+# rmsnorm_config._detect_arch_major) consult QUACK_ARCH via
+# get_device_capacity and would otherwise initialize CUDA — which both makes
+# the forkserver's context leak into children and trips torch's forked-child
+# guard. nvidia-smi queries the capability without creating a CUDA context.
+if "QUACK_ARCH" not in os.environ:
+    try:
+        out = subprocess.run(
+            ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        cap = out.stdout.strip().splitlines()[0].strip()  # e.g. "9.0"
+        major, minor = cap.split(".")
+        os.environ["QUACK_ARCH"] = f"{major}{minor}"
+        os.environ.setdefault(
+            "CUTE_DSL_ARCH", f"sm_{major}{minor}a" if int(major) >= 9 else f"sm_{major}{minor}"
+        )
+    except Exception:
+        pass  # CPU-only box: rely on user-provided env, as before
+# Belt and suspenders: even if some import still tries to touch CUDA, make
+# it see no devices rather than creating a context in the forkserver.
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+from .. import cache  # noqa: F401, E402  (pulls torch, cutlass.cute, tvm_ffi)

build/torch-cuda/quack/cache/async_compile.py ADDED Viewed

	@@ -0,0 +1,413 @@

+# Copyright (c) 2026, Tri Dao.
+"""Async kernel compilation: defer-and-retry via a pool of CPU subprocesses.
+When a pool is active, ``jit_cache`` handles a ``.o``-cache miss by
+submitting the pickled ``(module, qualname, args, kwargs)`` of the
+``_compile_*`` function to the pool and raising :class:`CompilePending`
+instead of compiling in-process. The caller defers the work item, runs
+something else, and retries once the ``.o`` lands (a ~1 ms load). Two
+callers implement this loop:
+* the pytest plugin's ``--async-compile`` defer loop (tests are the work
+  items; see :mod:`quack.testing.pytest_plugin`);
+* the autotuner's bench loop under :func:`pool_scope` (candidate configs
+  are the work items; see :class:`quack.autotuner.Autotuner`).
+Design notes:
+* **The ``.o`` file is the only rendezvous** between workers and consumers —
+  compiled kernels aren't picklable, so the persistent cache doubles as the
+  IPC channel, and the per-key ``flock`` in ``jit_cache`` doubles as
+  cross-process dedupe (multiple pools / xdist workers coexist safely;
+  :func:`_flock_held_exclusively` lets a consumer defer on a key some other
+  process is already compiling).
+* **Workers never launch kernels by construction**: they call the
+  tensor-free ``_compile_*`` functions directly, GPU-blind (arch pinned via
+  ``QUACK_ARCH``/``CUTE_DSL_ARCH``, ``CUDA_VISIBLE_DEVICES=""``).
+* **Worker startup is an Inductor-style sidecar**: a ``forkserver`` preloads
+  torch/cutlass once (:mod:`quack.cache._pool_preload`, ~13 s) and workers
+  fork from it copy-on-write (~0.1 s each). :func:`_neutral_main` keeps
+  multiprocessing child prep from re-executing the user's script.
+* **Failure semantics**: a failed pool compile is never trusted — the
+  consumer falls through to an in-process compile so the real exception
+  surfaces with a local traceback.
+Env knobs: ``QUACK_ASYNC_COMPILE_START=spawn`` (disable the fork sidecar),
+``QUACK_COMPILE_WORKERS`` (shared-executor size, default 8).
+"""
+from __future__ import annotations
+import base64
+import contextlib
+import fcntl
+import importlib
+import os
+import pickle
+from concurrent.futures import Future, ProcessPoolExecutor
+from multiprocessing import get_context
+from typing import Optional
+def _flock_held_exclusively(lock_path: str) -> bool:
+    """True if some process currently holds the flock exclusively.
+    Used to detect "another process is compiling this key right now" so the
+    consumer defers instead of submitting a duplicate compile to its own
+    pool (a duplicate would occupy a pool slot blocked on the same flock).
+    """
+    try:
+        fd = os.open(lock_path, os.O_RDONLY | os.O_CREAT)
+    except OSError:
+        return False
+    try:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_SH | fcntl.LOCK_NB)
+            fcntl.flock(fd, fcntl.LOCK_UN)
+            return False
+        except OSError:
+            return True
+    finally:
+        os.close(fd)
+class CompilePending(BaseException):
+    """A jit_cache miss was submitted to the async compile pool.
+    The caller (test) cannot proceed until the ``.o`` exists; the test runner
+    should defer the test and retry it later. Carries the cache ``sha`` so
+    the runner can poll for completion without re-running the test.
+    Derives from :class:`BaseException` (like ``KeyboardInterrupt``) so that
+    test-body ``except Exception`` / ``pytest.raises(Exception)`` blocks
+    cannot swallow it and turn a not-yet-run test into a false pass. Only
+    the plugin's phase hooks are supposed to catch it.
+    """
+    def __init__(self, sha: str, qualname: str):
+        super().__init__(f"kernel compile pending in pool: {qualname} [{sha[:12]}]")
+        self.sha = sha
+        self.qualname = qualname
+def _detect_arch_env() -> tuple[Optional[str], Optional[str]]:
+    """Return (QUACK_ARCH, CUTE_DSL_ARCH) for GPU-blind pool workers.
+    An explicit ``QUACK_ARCH`` env override wins — CI cross-compiles for a
+    different arch than the runner's GPU (e.g. ``QUACK_ARCH=120`` on an
+    H100), and workers must compile for the *target* arch, not the physical
+    one. Otherwise detect from the parent's GPU. Either way the workers
+    themselves never touch the CUDA driver (no context per worker,
+    fork-safe).
+    """
+    quack_arch = os.environ.get("QUACK_ARCH")
+    if quack_arch is not None:
+        cute_arch = os.environ.get("CUTE_DSL_ARCH")
+        if cute_arch is None:
+            from ..cute_dsl_utils import _parse_arch_str
+            major, minor = _parse_arch_str(quack_arch)
+            cc = f"{major}{minor}"
+            cute_arch = f"sm_{cc}a" if major >= 9 else f"sm_{cc}"
+        return quack_arch, cute_arch
+    try:
+        import torch
+        if torch.cuda.is_available():
+            major, minor = torch.cuda.get_device_capability()
+            cc = f"{major}{minor}"
+            return cc, f"sm_{cc}a" if major >= 9 else f"sm_{cc}"
+    except Exception:
+        pass
+    return None, os.environ.get("CUTE_DSL_ARCH")
+def _pool_initializer(quack_arch: Optional[str], cute_dsl_arch: Optional[str]):
+    # GPU-blind compilation: hide devices and pin the target arch via the
+    # same overrides the CPU-only compile workflow uses. Forked workers must
+    # never initialize CUDA (fork-safety), and spawned workers save the
+    # ~1-2 s + ~300 MB of a per-worker CUDA context.
+    if quack_arch is not None:
+        os.environ["QUACK_ARCH"] = quack_arch
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    if cute_dsl_arch is not None:
+        os.environ["CUTE_DSL_ARCH"] = cute_dsl_arch
+    # Pay the heavy torch/cutlass import at worker start (no-op under
+    # forkserver: the preload already imported it before the fork).
+    from .. import cache  # noqa: F401
+def _pool_worker(mod_name: str, qualname: str, key_b64: str, o_path: str) -> Optional[str]:
+    """Compile one key. Returns None on success, error string on failure."""
+    try:
+        obj = importlib.import_module(mod_name)
+        for part in qualname.split("."):
+            obj = getattr(obj, part)
+        args, kwargs = pickle.loads(base64.b64decode(key_b64))
+        obj(*args, **kwargs)  # jit_cache wrapper: compiles + exports .o
+        if not os.path.exists(o_path):
+            return "compile succeeded but .o was not exported"
+        return None
+    except Exception as e:
+        return f"{type(e).__name__}: {e}"
+def _make_executor(jobs: int) -> ProcessPoolExecutor:
+    """Build a compile-worker executor (Inductor-style forkserver sidecar).
+    Forkserver + preload: one sidecar process pays the ~13 s torch/cutlass
+    import once, workers fork from it in ~0.1 s each (copy-on-write). The
+    forkserver singleton is shared per-process, so multiple executors (the
+    test pool, the autotuner's) fork from the same warm sidecar. Opt out
+    with QUACK_ASYNC_COMPILE_START=spawn.
+    """
+    start_method = os.environ.get("QUACK_ASYNC_COMPILE_START", "forkserver")
+    ctx = get_context(start_method)
+    if start_method == "forkserver":
+        # vendored under sonic_moe.quack: preload the sibling module by package
+        # (upstream hardcodes the top-level "quack.cache._pool_preload").
+        ctx.set_forkserver_preload([__package__ + "._pool_preload"])
+    return ProcessPoolExecutor(
+        max_workers=jobs,
+        mp_context=ctx,
+        initializer=_pool_initializer,
+        initargs=_detect_arch_env(),
+    )
+_shared_executor: Optional[ProcessPoolExecutor] = None
+def get_shared_executor() -> ProcessPoolExecutor:
+    """Executor for ad-hoc compile tasks (e.g. autotuner precompile sweeps).
+    Reuses the active :class:`CompilePool`'s executor when one exists (the
+    pytest ``--async-compile`` session pool); otherwise lazily creates a
+    process-wide executor sized by ``QUACK_COMPILE_WORKERS`` (default 8).
+    Deliberately ignores :class:`suppress_pool` — suppression turns off the
+    *defer-on-miss* behavior of jit_cache, not access to compile workers.
+    """
+    global _shared_executor
+    if _active_pool is not None:
+        return _active_pool._executor
+    if _shared_executor is None:
+        import atexit
+        _shared_executor = _make_executor(int(os.environ.get("QUACK_COMPILE_WORKERS", "8")))
+        # Explicit teardown: without this, the executor is GC'd during
+        # interpreter shutdown after its weakref machinery is already gone,
+        # printing a spurious "Exception ignored in weakref_cb".
+        atexit.register(_shared_executor.shutdown, wait=False, cancel_futures=True)
+    return _shared_executor
+@contextlib.contextmanager
+def _neutral_main():
+    """Stop multiprocessing child prep from re-executing the user's script.
+    ``Process.start()`` captures preparation data from ``sys.modules['__main__']``:
+    for a path-based script the *child* re-runs the whole file via
+    ``runpy.run_path`` (so pickles referencing ``__main__`` resolve). Our
+    tasks never reference ``__main__`` — they resolve everything by module
+    name — and a user script that, say, builds CUDA tensors at import time
+    would kill every worker at spawn with "Cannot re-initialize CUDA in
+    forked subprocess". Executor workers are spawned synchronously inside
+    ``executor.submit`` (``_adjust_process_count``), so masking ``__main__``
+    with an empty stub for the duration of the submit is sufficient and
+    scoped. Single-threaded callers only (pytest defer loop, autotune bench
+    loop).
+    """
+    import sys
+    import types
+    real_main = sys.modules.get("__main__")
+    sys.modules["__main__"] = types.ModuleType("__main__")  # no __file__/__spec__
+    try:
+        yield
+    finally:
+        if real_main is not None:
+            sys.modules["__main__"] = real_main
+class CompilePool:
+    """Process pool + in-flight bookkeeping, keyed by jit_cache sha.
+    Owns its executor by default; pass ``executor=`` to share one (e.g.
+    :func:`pool_scope` wraps the session-long shared executor so scoped
+    pools don't respawn workers per autotune sweep). A shared executor is
+    not shut down by :meth:`shutdown` — only this pool's futures are
+    cancelled.
+    """
+    def __init__(self, jobs: Optional[int] = None, executor: Optional[ProcessPoolExecutor] = None):
+        self._own_executor = executor is None
+        self._executor = executor if executor is not None else _make_executor(jobs)
+        self._futures: dict[str, Future] = {}
+        # Keys being compiled by *another process* (e.g. a different xdist
+        # worker's pool), detected via the per-key flock. We defer on them
+        # without spending one of our own pool slots on a duplicate compile.
+        # sha -> (o_path, lock_path)
+        self._external: dict[str, tuple[str, str]] = {}
+        self.n_submitted = 0
+    def mark_external(self, sha: str, o_path: str, lock_path: str) -> None:
+        """Record that some other process is compiling ``sha`` right now."""
+        if sha not in self._futures:
+            self._external[sha] = (str(o_path), str(lock_path))
+    def prewarm(self) -> None:
+        """Start the sidecar + first worker now, off the critical path.
+        Same idea as Inductor's ``warm_pool()``: the forkserver's ~13 s
+        torch/cutlass preload import starts at the first ``Process`` spawn,
+        which is lazy (first submit). Submitting a no-op at session setup
+        overlaps that import with pytest collection and the leading warm
+        tests instead of the first cold compile.
+        """
+        with _neutral_main():
+            self._executor.submit(os.getpid)
+    def submit_raw(self, sha: str, mod: str, qualname: str, key_b64: str, o_path: str) -> None:
+        if sha in self._futures:
+            return
+        with _neutral_main():
+            self._futures[sha] = self._executor.submit(_pool_worker, mod, qualname, key_b64, o_path)
+        self.n_submitted += 1
+    def submit(self, sha: str, fn, args: tuple, kwargs: dict, o_path) -> bool:
+        """Submit a live jit_cache miss. Returns False if the key can't be
+        shipped to a subprocess (unpicklable args, ``<locals>`` qualname,
+        fn defined in ``__main__``) — the caller should compile in-process
+        instead."""
+        if sha in self._futures:
+            return True
+        if "<locals>" in fn.__qualname__ or fn.__module__ == "__main__":
+            # Not resolvable by module+qualname in a worker; compile in-process.
+            return False
+        try:
+            key_b64 = base64.b64encode(pickle.dumps((args, kwargs))).decode("ascii")
+        except Exception:
+            return False
+        self.submit_raw(sha, fn.__module__, fn.__qualname__, key_b64, str(o_path))
+        return True
+    def poll(self, sha: str) -> tuple[str, Optional[str]]:
+        """Return (state, error): state in {"new", "pending", "done", "failed"}."""
+        fut = self._futures.get(sha)
+        if fut is None:
+            ext = self._external.get(sha)
+            if ext is not None:
+                o_path, lock_path = ext
+                if os.path.exists(o_path):
+                    del self._external[sha]
+                    return "done", None
+                if _flock_held_exclusively(lock_path):
+                    return "pending", None
+                # External compiler released the lock without producing a .o
+                # (crashed / failed): forget it so the next attempt submits
+                # to our own pool.
+                del self._external[sha]
+            return "new", None
+        if not fut.done():
+            return "pending", None
+        try:
+            err = fut.result()
+        except Exception as e:  # BrokenProcessPool etc.
+            err = f"pool worker died: {type(e).__name__}: {e}"
+        return ("done", None) if err is None else ("failed", err)
+    def stats(self) -> dict:
+        done = sum(1 for f in self._futures.values() if f.done())
+        errors = []
+        for sha, f in self._futures.items():
+            if not f.done() or f.cancelled():
+                continue
+            exc = f.exception()
+            err = f"{type(exc).__name__}: {exc}" if exc is not None else f.result()
+            if err:
+                errors.append((sha, err))
+        return {
+            "submitted": self.n_submitted,
+            "done": done,
+            "failed": len(errors),
+            "errors": errors,
+        }
+    def shutdown(self) -> None:
+        if self._own_executor:
+            self._executor.shutdown(wait=False, cancel_futures=True)
+        else:
+            for fut in self._futures.values():
+                fut.cancel()
+# --- module-level active pool -----------------------------------------------
+_active_pool: Optional[CompilePool] = None
+_suppress_depth = 0
+class suppress_pool:
+    """Context manager: make :func:`get_active_pool` return None inside.
+    Used by the test runner for a deferred test's final attempt: compile
+    in-process (blocking) so a key that never completes in the pool still
+    produces a real result or a real traceback instead of deferring forever.
+    """
+    def __enter__(self):
+        global _suppress_depth
+        _suppress_depth += 1
+        return self
+    def __exit__(self, *exc):
+        global _suppress_depth
+        _suppress_depth -= 1
+def activate(jobs: int) -> CompilePool:
+    """Activate the session-wide pool (idempotent). Used by the pytest plugin;
+    scoped callers should prefer :func:`pool_scope`."""
+    global _active_pool
+    if _active_pool is None:
+        _active_pool = CompilePool(jobs)
+    return _active_pool
+def deactivate() -> None:
+    global _active_pool
+    if _active_pool is not None:
+        _active_pool.shutdown()
+        _active_pool = None
+def get_active_pool() -> Optional[CompilePool]:
+    return None if _suppress_depth > 0 else _active_pool
+@contextlib.contextmanager
+def pool_scope():
+    """Activate a compile pool for the duration of the block.
+    Reuses the globally active pool when one exists (e.g. the pytest
+    ``--async-compile`` session pool); otherwise activates a temporary pool
+    backed by the shared executor and deactivates it on exit — so
+    ``CompilePending`` can only escape into code inside the block, never
+    into unrelated user code paths.
+    This is how the autotuner overlaps candidate-config compilation with
+    benchmarking: the bench loop runs inside ``pool_scope()``, catches
+    ``CompilePending`` per config, and retries a config once its ``.o``
+    lands (see ``Autotuner.__call__``).
+    """
+    global _active_pool
+    if _active_pool is not None:
+        yield _active_pool
+        return
+    pool = CompilePool(executor=get_shared_executor())
+    _active_pool = pool
+    try:
+        yield pool
+    finally:
+        _active_pool = None
+        pool.shutdown()

build/torch-cuda/quack/{cache_utils.py → cache/jit.py} RENAMED Viewed

@@ -1,15 +1,21 @@
 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
-"""Persistent .o cache for CuTe DSL compiled kernels.
-Compiled kernels are exported as object files (.o) via export_to_c.
-On subsequent runs the .o is loaded via tvm_ffi (~1ms) instead of
-re-generating IR + re-JIT'ing (~100ms per kernel).
-Controls:
-  QUACK_CACHE_ENABLED=0       — disable persistent .o cache (default: enabled)
-  QUACK_CACHE_DIR=path        — override default cache directory
 """
 import fcntl
 import functools
 import hashlib
@@ -26,26 +32,20 @@ import cutlass
 import cutlass.cute as cute
 import tvm_ffi
-CACHE_ENABLED: bool = os.getenv("QUACK_CACHE_ENABLED", "1") == "1"
-CACHE_DIR: str | None = os.getenv("QUACK_CACHE_DIR", None)
-COMPILE_ONLY: bool = False
-# Downstream projects can append directories here to include their sources
-# in the cache fingerprint. Must be set before the first jit_cache call.
-EXTRA_SOURCE_DIRS: list[Path] = []
 EXPORT_FUNC_NAME = "func"
 LOCK_TIMEOUT = 60
 CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
-def _noop_kernel(*args, **kwargs):
-    pass
 def get_cache_path() -> Path:
-    if CACHE_DIR is not None:
-        cache_dir = Path(CACHE_DIR)
     else:
         cache_dir = Path(tempfile.gettempdir()) / getuser() / "quack_cache"
     cache_dir.mkdir(parents=True, exist_ok=True)
@@ -70,8 +70,13 @@ def _compute_source_fingerprint() -> str:
     h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode())
     h.update(f"cutlass={cutlass.__version__}".encode())
     h.update(f"tvm_ffi={tvm_ffi.__version__}".encode())
-    _hash_source_dir(h, Path(__file__).resolve().parent)
-    for extra_dir in EXTRA_SOURCE_DIRS:
         _hash_source_dir(h, Path(extra_dir).resolve())
     return h.hexdigest()
@@ -126,6 +131,25 @@ def jit_cache(fn):
     The decorated function should return a compiled kernel (i.e. call cute.compile).
     The disk cache key is (fn.__qualname__, *args, **sorted_kwargs).
     """
     cache = {}
     hits = 0
@@ -136,49 +160,162 @@ def jit_cache(fn):
         nonlocal hits, misses
         cache_key = args + tuple(sorted(kwargs.items())) if kwargs else args
-        # 1. In-memory hit
         if cache_key in cache:
             hits += 1
-            return _noop_kernel if COMPILE_ONLY else cache[cache_key]
-        # 2. Disk hit
-        disk_key = (fn.__qualname__,) + cache_key
-        if CACHE_ENABLED:
-            sha = _key_to_hash(disk_key)
-            cache_path = get_cache_path() / _compute_source_fingerprint()
-            cache_path.mkdir(parents=True, exist_ok=True)
-            o_path = cache_path / f"{sha}.o"
-            lock_path = cache_path / f"{sha}.lock"
             try:
                 with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT):
                     if o_path.exists():
-                        m = cute.runtime.load_module(str(o_path), enable_tvm_ffi=True)
-                        loaded = m[EXPORT_FUNC_NAME]
                         cache[cache_key] = loaded
                         hits += 1
-                        return _noop_kernel if COMPILE_ONLY else loaded
-            except RuntimeError:
-                pass
-        # 3. Compile
-        misses += 1
-        compiled_fn = fn(*args, **kwargs)
-        # 4. Store
-        cache[cache_key] = compiled_fn
-        if CACHE_ENABLED:
-            try:
-                with FileLock(lock_path, exclusive=True, timeout=LOCK_TIMEOUT):
-                    if not o_path.exists():
-                        o_path.parent.mkdir(parents=True, exist_ok=True)
-                        compiled_fn.export_to_c(
-                            object_file_path=str(o_path),
-                            function_name=EXPORT_FUNC_NAME,
-                        )
-            except Exception as e:
-                print(f"quack cache: export failed for key {sha}: {e}")
-        return _noop_kernel if COMPILE_ONLY else compiled_fn
     def cache_clear():
         nonlocal hits, misses

 # Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+"""Persistent ``.o`` cache for CuTe DSL compiled kernels.
+Compiled kernels are exported as object files (``.o``) via ``export_to_c``. On
+subsequent runs the ``.o`` is loaded via tvm_ffi (~1 ms) instead of
+re-generating IR + re-JIT'ing (~500 ms per kernel).
+Runtime config (``CACHE_ENABLED``, ``CACHE_DIR``, ``EXTRA_SOURCE_DIRS``)
+lives in :mod:`quack.cache` (the package init).
+When an async compile pool is active (see :mod:`quack.cache.async_compile`),
+a cold miss is shipped to a CPU worker and :class:`CompilePending` is raised
+instead of compiling in-process; the caller (pytest defer loop, autotune
+bench loop) retries once the ``.o`` lands.
 """
+from __future__ import annotations
 import fcntl
 import functools
 import hashlib
 import cutlass.cute as cute
 import tvm_ffi
+# `quack.cache` (the package itself) holds the mutable runtime flags as a
+# single source of truth; reads happen via attribute access on `_state` so we
+# always see the live value, not a snapshot taken at module import.
+from .. import cache as _state  # noqa: E402  (intentional partial-import; see __init__.py)
 EXPORT_FUNC_NAME = "func"
 LOCK_TIMEOUT = 60
 CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
 def get_cache_path() -> Path:
+    if _state.CACHE_DIR is not None:
+        cache_dir = Path(_state.CACHE_DIR)
     else:
         cache_dir = Path(tempfile.gettempdir()) / getuser() / "quack_cache"
     cache_dir.mkdir(parents=True, exist_ok=True)
     h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode())
     h.update(f"cutlass={cutlass.__version__}".encode())
     h.update(f"tvm_ffi={tvm_ffi.__version__}".encode())
+    # Hash the entire `quack` package, not just `quack/cache/`. Resolving via
+    # the top-level package import keeps the fingerprint stable regardless of
+    # where inside the package this file lives.
+    import importlib as _importlib; _quack = _importlib.import_module(__package__.rsplit(".", 1)[0])
+    _hash_source_dir(h, Path(_quack.__file__).resolve().parent)
+    for extra_dir in _state.EXTRA_SOURCE_DIRS:
         _hash_source_dir(h, Path(extra_dir).resolve())
     return h.hexdigest()
     The decorated function should return a compiled kernel (i.e. call cute.compile).
     The disk cache key is (fn.__qualname__, *args, **sorted_kwargs).
+    Concurrency model
+    -----------------
+    The disk side uses a per-key ``{sha}.lock`` file (advisory ``flock``):
+    * **Fast path (warm cache).** If the ``.o`` file already exists, we take a
+      shared lock just long enough to ``load_module`` it. Many readers can
+      proceed concurrently.
+    * **Slow path (cold cache).** The actual ``fn(*args, **kwargs)`` compile
+      runs *under* the exclusive lock. This serializes redundant compilations
+      of the same key across xdist workers / processes: if N processes race
+      on a cold key, only one calls ``cute.compile``; the rest wait for the
+      lock, see the ``.o`` appear, and load it. (Previously the compile ran
+      *between* the shared-lock check and the exclusive-lock export, so all
+      N processes wasted CPU compiling the same key in parallel — wall time
+      was unchanged but compile-CPU pressure scaled with concurrency, which
+      starved other compiles when many keys were cold at once.)
+    The lock is per-key, so distinct keys never contend with each other.
     """
     cache = {}
     hits = 0
         nonlocal hits, misses
         cache_key = args + tuple(sorted(kwargs.items())) if kwargs else args
+        # Snapshot once per call so a concurrent flip of ``_state.CACHE_ENABLED``
+        # mid-call can't desync the disk-path branch.
+        enabled = _state.CACHE_ENABLED
+        # 1. In-memory hit. Same process already compiled or loaded this key.
         if cache_key in cache:
             hits += 1
+            return cache[cache_key]
+        # 2. Cache disabled: pure in-process compile, no disk side effects.
+        if not enabled:
+            misses += 1
+            compiled_fn = fn(*args, **kwargs)
+            cache[cache_key] = compiled_fn
+            return compiled_fn
+        sha = _key_to_hash((fn.__qualname__,) + cache_key)
+        cache_path = get_cache_path() / _compute_source_fingerprint()
+        cache_path.mkdir(parents=True, exist_ok=True)
+        o_path = cache_path / f"{sha}.o"
+        lock_path = cache_path / f"{sha}.lock"
+        def _load_cached() -> object:
+            """Load the .o into a callable; caller guarantees existence."""
+            m = cute.runtime.load_module(str(o_path), enable_tvm_ffi=True)
+            return m[EXPORT_FUNC_NAME]
+        def _quarantine_corrupt(exc: Exception) -> None:
+            """A cached .o that fails to load (truncated write from a killed
+            worker, missing __tvm_ffi_func, ...) is a cache miss, not an error:
+            delete it so this and future processes recompile instead of failing
+            forever (the CI cache persists across runs)."""
+            print(
+                f"quack cache: corrupt cached object for key {sha} "
+                f"({type(exc).__name__}: {exc}); deleting and recompiling"
+            )
+            try:
+                o_path.unlink()
+            except OSError:
+                pass
+        # 3. Fast path: optimistic existence check, then shared-lock load.
+        #    The unlocked ``.exists()`` is a no-cost short-circuit for warm
+        #    caches; the shared lock guards against reading a partial file
+        #    while a concurrent writer holds the exclusive lock.
+        if o_path.exists():
             try:
                 with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT):
                     if o_path.exists():
+                        try:
+                            loaded = _load_cached()
+                        except Exception as e:
+                            # Corrupt entry: recover under the exclusive lock in
+                            # the slow path (shared lock can't safely delete).
+                            _quarantine_corrupt(e)
+                        else:
+                            cache[cache_key] = loaded
+                            hits += 1
+                            return loaded
+            except RuntimeError:
+                pass  # lock timeout; fall through to slow path
+        # 3b. Async-compile pool: on a cold miss with a pool
+        #     active, ship the key to a CPU subprocess and raise
+        #     CompilePending instead of compiling in-process. The test runner
+        #     defers the test and retries once the worker has exported the
+        #     .o. Pool failures fall through to the in-process compile below
+        #     so the real exception surfaces with a local traceback.
+        from . import async_compile as _async
+        pool = _async.get_active_pool()
+        if pool is not None:
+            state, err = pool.poll(sha)
+            if state == "new":
+                # If another process (e.g. a different xdist worker's pool)
+                # holds the exclusive per-key flock, it is compiling this key
+                # right now: defer on it instead of submitting a duplicate.
+                if _async._flock_held_exclusively(str(lock_path)):
+                    pool.mark_external(sha, str(o_path), str(lock_path))
+                    raise _async.CompilePending(sha, fn.__qualname__)
+                if pool.submit(sha, fn, args, kwargs, o_path):
+                    raise _async.CompilePending(sha, fn.__qualname__)
+                # unpicklable key / <locals> qualname: compile in-process
+            elif state == "pending":
+                raise _async.CompilePending(sha, fn.__qualname__)
+            elif state == "done":
+                try:
+                    with FileLock(lock_path, exclusive=False, timeout=LOCK_TIMEOUT):
+                        if o_path.exists():
+                            try:
+                                loaded = _load_cached()
+                            except Exception as e:
+                                _quarantine_corrupt(e)
+                            else:
+                                cache[cache_key] = loaded
+                                hits += 1
+                                return loaded
+                except RuntimeError:
+                    pass  # lock timeout; fall through to slow path
+            else:  # "failed"
+                print(
+                    f"quack cache: async compile failed for {fn.__qualname__} "
+                    f"[{sha[:12]}]: {err}; recompiling in-process for a real traceback"
+                )
+        # 4. Slow path: take EXCLUSIVE lock and compile under it. The recheck
+        #    inside the lock catches the race where another process compiled
+        #    while we were waiting; in that case we just load and return
+        #    without duplicating the compile.
+        try:
+            with FileLock(lock_path, exclusive=True, timeout=LOCK_TIMEOUT):
+                if o_path.exists():
+                    try:
+                        loaded = _load_cached()
+                    except Exception as e:
+                        _quarantine_corrupt(e)  # holds the exclusive lock: safe
+                    else:
                         cache[cache_key] = loaded
                         hits += 1
+                        return loaded
+                misses += 1
+                compiled_fn = fn(*args, **kwargs)
+                # Export to a private temp file, then atomically rename into
+                # place: a process killed mid-export (xdist worker OOM-kill,
+                # timeout) must never leave a truncated .o at the final path —
+                # the advisory flock dies with the process, and a persistent
+                # cache (CI keeps one in $HOME) would then fail every future
+                # run on this key with "Symbols not found: __tvm_ffi_func".
+                tmp_path = o_path.with_suffix(f".o.tmp.{os.getpid()}")
+                try:
+                    compiled_fn.export_to_c(
+                        object_file_path=str(tmp_path),
+                        function_name=EXPORT_FUNC_NAME,
+                    )
+                    os.replace(tmp_path, o_path)
+                except Exception as e:
+                    print(f"quack cache: export failed for key {sha}: {e}")
+                    try:
+                        tmp_path.unlink()
+                    except OSError:
+                        pass
+                cache[cache_key] = compiled_fn
+                return compiled_fn
+        except RuntimeError as e:
+            # Lock acquisition timed out (heavy contention or stuck holder).
+            # Fall back to in-process compile, no disk write. Better to do
+            # the work twice than to fail the test.
+            print(
+                f"quack cache: lock timeout for key {sha}: {e}; "
+                f"falling back to in-process compile without disk cache"
+            )
+            misses += 1
+            compiled_fn = fn(*args, **kwargs)
+            cache[cache_key] = compiled_fn
+            return compiled_fn
     def cache_clear():
         nonlocal hits, misses

build/torch-cuda/quack/compile_utils.py CHANGED Viewed

@@ -6,10 +6,20 @@ import cutlass.cute as cute
 def make_fake_tensor(dtype, shape, divisibility=1, leading_dim=-1) -> Optional[cute.Tensor]:
-    if leading_dim < 0:
-        leading_dim = len(shape) + leading_dim
     if dtype is None:
         return None
     stride = tuple(
         cute.sym_int64(divisibility=divisibility) if i != leading_dim else 1
         for i in range(len(shape))
@@ -17,3 +27,8 @@ def make_fake_tensor(dtype, shape, divisibility=1, leading_dim=-1) -> Optional[c
     return cute.runtime.make_fake_tensor(
         dtype, shape, stride=stride, assumed_align=divisibility * dtype.width // 8
     )

 def make_fake_tensor(dtype, shape, divisibility=1, leading_dim=-1) -> Optional[cute.Tensor]:
+    """Build a fake CuTe tensor with dynamic (sym) strides for tensor-free compilation.
+    ``leading_dim`` selects the dim whose stride is statically 1 (matching
+    ``from_dlpack(...).mark_layout_dynamic(leading_dim=...)``). Pass
+    ``leading_dim=None`` for a fully-dynamic layout with no static stride-1 dim
+    (matching ``mark_layout_dynamic()`` on a tensor without a contiguous dim).
+    ``divisibility`` is in elements; ``assumed_align`` (bytes) is
+    ``divisibility * dtype.width // 8``.
+    """
     if dtype is None:
         return None
+    if leading_dim is not None and leading_dim < 0:
+        leading_dim = len(shape) + leading_dim
     stride = tuple(
         cute.sym_int64(divisibility=divisibility) if i != leading_dim else 1
         for i in range(len(shape))
     return cute.runtime.make_fake_tensor(
         dtype, shape, stride=stride, assumed_align=divisibility * dtype.width // 8
     )
+def make_fake_stream():
+    """Fake CUDA stream for tensor-free compilation (real stream comes from the TVM FFI env)."""
+    return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)

build/torch-cuda/quack/complex.py ADDED Viewed

	@@ -0,0 +1,292 @@

+"""Complex64 element type for CuTe-DSL kernels.
+Single-precision complex (re + imj) carried as f64-packed bits (re in the low
+32 bits, im in the high 32 bits). f64 is on `cute.MemRefType`'s element-type
+allowlist; the natural `complex<f32>` MLIR type is not. Arithmetic methods
+unpack each f64 into two Float32 lanes, compute, and repack -- the bitcasts
+are folded out by ptxas.
+Inherits from `Float32` (with `width=64, mlir_type=T.f64` overrides) so that
+Python's subclass-precedence rule routes `Float32 OP Complex64` to our
+reflected `__r*__` operators before Numeric's promotion logic sees the
+operands. Without this, Float32-on-the-LEFT would silently promote through
+Float64 conversion and corrupt the packed bits.
+Boundary convention (tvm-ffi): the compiled kernel's ABI sees f64 storage.
+At the call site, pass `torch.complex64` tensors as `t.view(torch.float64)`
+-- use `complex_storage(t)` for the conversion.
+See `AI/complex64_design_notes.md` for the why and what's been validated.
+"""
+from __future__ import annotations
+import ctypes
+import numpy as np
+import torch
+import cutlass.cute as cute
+from cutlass import Float32, Numeric
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import arith
+from cutlass._mlir.extras import types as T
+from cutlass._mlir_helpers.arith import bitcast as _bitcast
+from cutlass.base_dsl.typing import FloatMeta
+class Complex64(Float32, metaclass=FloatMeta, width=64, mlir_type=T.f64):
+    """Complex64 carried as f64-packed bits (re in low 32, im in high 32).
+    `tensor.element_type is Complex64` inside the kernel; indexing returns
+    Complex64 instances; `+`, `-`, `*`, `__neg__`, and `conj()` work natively.
+    """
+    def __init__(self, x, im=None, *, loc=None, ip=None):
+        # Two-arg lane form: Complex64(re, im) packs (re, im) into f64 bits.
+        # Coerce both through Float32 so int / float / Float32 / ir.Value(f32)
+        # all work as inputs.
+        if im is not None:
+            # Static fast path: both args are Python int/float -- no MLIR
+            # context needed (lets host-side code call Complex64(2.5, -1.5)).
+            if isinstance(x, (int, float)) and isinstance(im, (int, float)):
+                Complex64.__init__(self, complex(x, im), loc=loc, ip=ip)
+                return
+            re_ssa = Float32(x).ir_value()
+            im_ssa = Float32(im).ir_value()
+            Numeric.__init__(self, Complex64._pack_ssa(re_ssa, im_ssa))
+            return
+        # Same-type copy MUST be checked first. `_cvt_to_dest` (cute/tensor.py)
+        # calls `data.to(element_type)` on every tensor write, which becomes
+        # `Complex64(complex_instance)`; falling through to the generic-Numeric
+        # branch below would re-pack as (real_view_of_packed_bits, 0) and
+        # silently corrupt the data.
+        if isinstance(x, Complex64):
+            Numeric.__init__(self, x.value)
+            return
+        if isinstance(x, complex):
+            f64_val = _pack_python_complex(x)
+            Numeric.__init__(self, f64_val)
+            return
+        if isinstance(x, ir.Value):
+            if x.type == T.f64():
+                # Already in our storage form (loaded from a Complex64 tensor,
+                # or output of _pack_ssa).
+                Numeric.__init__(self, x)
+                return
+            if x.type == T.f32():
+                packed = Complex64._pack_ssa(x, arith.constant(T.f32(), 0.0))
+                Numeric.__init__(self, packed)
+                return
+            raise TypeError(f"Complex64: ir.Value of unsupported type {x.type}")
+        if isinstance(x, Numeric):
+            # Float32, Int32, Float64, etc. -> coerce real lane through Float32,
+            # imag lane = 0. Float32(Float32) is a no-op, so this also handles
+            # the Float32 case cleanly.
+            re_ssa = Float32(x).ir_value()
+            packed = Complex64._pack_ssa(re_ssa, arith.constant(T.f32(), 0.0))
+            Numeric.__init__(self, packed)
+            return
+        if isinstance(x, (int, float)):
+            Complex64.__init__(self, complex(x, 0.0), loc=loc, ip=ip)
+            return
+        raise TypeError(f"Complex64: unsupported source type {type(x)}")
+    # ---- packing / unpacking primitives --------------------------------
+    @staticmethod
+    def _pack_ssa(re_f32, im_f32):
+        """Pack two f32 SSA lanes into one f64 SSA value (re lo, im hi)."""
+        re_i32 = _bitcast(re_f32, T.i32())
+        im_i32 = _bitcast(im_f32, T.i32())
+        re_i64 = arith.extui(T.i64(), re_i32)
+        im_i64 = arith.extui(T.i64(), im_i32)
+        hi = arith.shli(im_i64, arith.constant(T.i64(), 32))
+        return _bitcast(arith.ori(re_i64, hi), T.f64())
+    def _unpack(self):
+        """Split self -> (re_f32, im_f32) as Float32 SSA values."""
+        i64_ssa = _bitcast(self.ir_value(), T.i64())
+        lo32 = arith.trunci(T.i32(), i64_ssa)
+        hi32 = arith.trunci(T.i32(), arith.shrui(i64_ssa, arith.constant(T.i64(), 32)))
+        return Float32(_bitcast(lo32, T.f32())), Float32(_bitcast(hi32, T.f32()))
+    @staticmethod
+    def from_re_im(re: Float32, im: Float32) -> "Complex64":
+        """Build a Complex64 from two Float32 SSA lanes.
+        Equivalent to `Complex64(re, im)`; kept as an explicit name for the
+        hot-path call sites that want to skip the Float32 coercion in __init__.
+        """
+        return Complex64(Complex64._pack_ssa(re.ir_value(), im.ir_value()))
+    # Internal alias used by arithmetic methods.
+    _from_re_im = from_re_im
+    # ---- accessors ------------------------------------------------------
+    def real(self) -> Float32:
+        re, _ = self._unpack()
+        return re
+    def imag(self) -> Float32:
+        _, im = self._unpack()
+        return im
+    def conj(self) -> "Complex64":
+        re, im = self._unpack()
+        return Complex64._from_re_im(re, -im)
+    # ---- arithmetic -----------------------------------------------------
+    def __add__(self, other, *, loc=None, ip=None):
+        a_re, a_im = self._unpack()
+        b_re, b_im = _other_lanes(other)
+        return Complex64._from_re_im(a_re + b_re, a_im + b_im)
+    def __radd__(self, other, *, loc=None, ip=None):
+        return self.__add__(other, loc=loc, ip=ip)
+    def __sub__(self, other, *, loc=None, ip=None):
+        a_re, a_im = self._unpack()
+        b_re, b_im = _other_lanes(other)
+        return Complex64._from_re_im(a_re - b_re, a_im - b_im)
+    def __rsub__(self, other, *, loc=None, ip=None):
+        a_re, a_im = self._unpack()
+        b_re, b_im = _other_lanes(other)
+        return Complex64._from_re_im(b_re - a_re, b_im - a_im)
+    def __mul__(self, other, *, loc=None, ip=None):
+        a_re, a_im = self._unpack()
+        if isinstance(other, Complex64):
+            b_re, b_im = other._unpack()
+            return Complex64._from_re_im(a_re * b_re - a_im * b_im, a_re * b_im + a_im * b_re)
+        # Real scalar: (re, im) * s = (re*s, im*s)
+        s = Float32(other)
+        return Complex64._from_re_im(a_re * s, a_im * s)
+    def __rmul__(self, other, *, loc=None, ip=None):
+        return self.__mul__(other, loc=loc, ip=ip)
+    def __neg__(self, *, loc=None, ip=None):
+        re, im = self._unpack()
+        return Complex64._from_re_im(-re, -im)
+    # ---- runtime arg passing -------------------------------------------
+    def __c_pointers__(self):
+        # Scalar Complex64 args travel as 8 bytes (the packed-as-f64 value).
+        if not isinstance(self.value, float):
+            raise ValueError(
+                "Complex64 with a dynamic SSA value cannot be passed as a "
+                "kernel argument; only static values are supported"
+            )
+        return [ctypes.cast(ctypes.pointer(ctypes.c_double(self.value)), ctypes.c_void_p)]
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+def _pack_python_complex(c: complex) -> float:
+    """Compute the f64 representation of a complex bit-packed (re, im)."""
+    re_b = int(np.float32(c.real).view(np.uint32))
+    im_b = int(np.float32(c.imag).view(np.uint32))
+    return float(np.uint64((im_b << 32) | re_b).view(np.float64))
+def _other_lanes(other):
+    """Unpack the RHS of a binary op into (re_f32, im_f32) Float32 lanes."""
+    if isinstance(other, Complex64):
+        return other._unpack()
+    return Float32(other), Float32(0.0)
+def _retag_as_complex64(t):
+    """Restore `t.element_type is Complex64` after a code path that derived it
+    from MLIR (where complex64 collapses to Float64 / Int64 because
+    `Numeric.from_mlir_type` is a many-to-one lookup)."""
+    t._dtype = Complex64
+    return t
+# ---------------------------------------------------------------------------
+# Public helpers
+# ---------------------------------------------------------------------------
+def allocate_smem_complex(
+    allocator,
+    layout_or_shape,
+    byte_alignment: int = 16,
+    swizzle=None,
+):
+    """Allocate a `Complex64` smem tensor.
+    Wraps `cutlass.utils.SmemAllocator.allocate_tensor(Complex64, ...)` and
+    re-tags the result so `tensor.element_type is Complex64`. Without the
+    re-tag, the JIT-side tensor's element_type is `Float64` (derived from the
+    f64 memref) and writes go through `Complex64.to(Float64)` and corrupt the
+    packed bits.
+    """
+    t = allocator.allocate_tensor(
+        Complex64, layout_or_shape, byte_alignment=byte_alignment, swizzle=swizzle
+    )
+    return _retag_as_complex64(t)
+def recast_to_complex64(src: cute.Tensor) -> cute.Tensor:
+    """Recast any tensor (e.g. Float32, Int64) to a `Complex64` tensor.
+    Wraps `cute.recast_tensor(src, Complex64)` and re-tags the result. Same
+    dtype-loss bug as `allocate_smem_complex` -- the underlying recast goes
+    through `make_tensor`, which derives element_type from the MLIR memref
+    (here f64) and gets back Float64.
+    """
+    return _retag_as_complex64(cute.recast_tensor(src, Complex64))
+def complex_storage(t: torch.Tensor) -> torch.Tensor:
+    """View a `torch.complex64` tensor as `torch.float64` with the same memory.
+    Compiled kernels declared with `Complex64` element type have an f64 ABI;
+    use this at the boundary to satisfy tvm-ffi's dtype check without copying.
+    """
+    if t.dtype == torch.float64:
+        return t
+    if t.dtype != torch.complex64:
+        raise TypeError(
+            f"complex_storage expects torch.complex64 (or torch.float64 for "
+            f"already-converted storage), got {t.dtype}"
+        )
+    return t.view(torch.float64)
+# ---------------------------------------------------------------------------
+# tvm-ffi registration
+# ---------------------------------------------------------------------------
+def _register_with_tvm_ffi() -> None:
+    """Teach tvm-ffi that Complex64 has an f64 ABI.
+    Both `NumericToTVMFFIDtype` (the type->dtype-string lookup) and
+    `AcceptableNumericTypesForScalar` (the allowlist for scalar kernel args)
+    are plain Python collections, so we extend them at import time.
+    """
+    from cutlass.cute import _tvm_ffi_args_spec_converter as _cv
+    _cv.NumericToTVMFFIDtype.setdefault(Complex64, "float64")
+    if Complex64 not in _cv.AcceptableNumericTypesForScalar:
+        _cv.AcceptableNumericTypesForScalar.append(Complex64)
+_register_with_tvm_ffi()

build/torch-cuda/quack/copy_utils.py CHANGED Viewed

@@ -1,17 +1,19 @@
-# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
-from typing import Optional, Type, Tuple, Callable, Sequence
 from functools import partial
 import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, Int16, Boolean, const_expr
-from cutlass.cute.nvgpu import cpasync, warp, warpgroup
 from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup  # noqa
 from cutlass.cutlass_dsl import dsl_user_op
 import cutlass.pipeline
-from cutlass._mlir.dialects import llvm
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
@@ -20,6 +22,97 @@ from .utils import make_vector
 Sm100MmaPeerBitMask = 0xFEFFFFFF
 @dsl_user_op
@@ -36,9 +129,7 @@ def cvt_copy(
 ) -> None:
     assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem
     if const_expr(src.element_type != dst.element_type):
-        src_cvt = cute.make_rmem_tensor_like(src, dst.element_type)
-        src_cvt.store(src.load().to(dst.element_type))
-        src = src_cvt
     if const_expr(retile):
         src = tiled_copy.retile(src)
     cute.copy(tiled_copy, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
@@ -75,6 +166,13 @@ def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
     return dst
 @dsl_user_op
 def load_s2r_retile(
     tiled_copy: cute.TiledCopy,
@@ -95,11 +193,22 @@ def load_s2r_retile(
 @dsl_user_op
 def load_t2r(
-    thr_copy: cute.ThrCopy, shape: cute.Shape, src: cute.Tensor, *, loc=None, ip=None
 ) -> cute.Tensor:
-    cDst = cute.make_identity_tensor(shape)
-    dst = cute.make_rmem_tensor(thr_copy.partition_D(cDst).shape, src.element_type, loc=loc, ip=ip)
-    cute.copy(thr_copy, src, dst, loc=loc, ip=ip)
     return dst
@@ -321,18 +430,14 @@ def as_position_independent_swizzle_tensor(tensor: cute.Tensor) -> cute.Tensor:
     return cute.make_tensor(cute.recast_ptr(tensor.iterator, dtype=tensor.element_type), new_layout)
-def partition_D_position_independent(
-    thr_copy: cute.core.ThrCopy, tensor: cute.Tensor
-) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_D(tensor).iterator),
         thr_copy.partition_D(as_position_independent_swizzle_tensor(tensor)).layout,
     )
-def partition_S_position_independent(
-    thr_copy: cute.core.ThrCopy, tensor: cute.Tensor
-) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_S(tensor).iterator),
         thr_copy.partition_S(as_position_independent_swizzle_tensor(tensor)).layout,
@@ -373,12 +478,12 @@ def sm90_get_smem_load_op(
 def get_smem_store_atom(
-    arch: cutlass.Constexpr[int],
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
-    if const_expr(arch < 90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
@@ -397,12 +502,12 @@ def get_smem_store_atom(
 def get_smem_load_atom(
-    arch: cutlass.Constexpr[int],
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
-    if const_expr(arch < 90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
@@ -421,26 +526,35 @@ def get_smem_load_atom(
 def get_smem_store_C(
-    tiled_mma: cute.TiledMma,
     sC: cute.Tensor,
     tidx: Int32,
-    arch: int,
     transpose: bool = False,
     position_independent=False,
     major_mode_size: Optional[int] = None,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
-    copy_atom = get_smem_store_atom(arch, dtype, transpose, major_mode_size=major_mode_size)
-    tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tRS_sC = thr_copy.partition_D(sC)
     else:
         tRS_sC = partition_D_position_independent(thr_copy, sC)
-    def copy_fn(src: cute.Tensor, dst_idx: Optional[Int32] = None, **new_kwargs):
-        dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[None, None, None, dst_idx]
         cvt_copy(tiled_copy, src, dst_tensor, retile=True, **new_kwargs)
     return copy_fn, thr_copy, tRS_sC
@@ -449,19 +563,18 @@ def get_smem_load_C(
     tiled_mma: cute.TiledMma,
     sC: cute.Tensor,
     tidx: Int32,
-    arch: int,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
-    copy_atom = get_smem_load_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tSR_sC = thr_copy.partition_S(sC)
     else:
         tSR_sC = partition_S_position_independent(thr_copy, sC)
-    copy_atom_RS = get_smem_store_atom(arch, dtype, transpose)
     thr_copy_RS = cute.make_tiled_copy_C(copy_atom_RS, tiled_mma).get_slice(tidx)
     tRS_shape = thr_copy_RS.partition_S(cute.make_identity_tensor(sC.shape[:2])).shape
@@ -475,10 +588,18 @@ def get_smem_load_C(
 def epilog_smem_copy_atom(
     tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False
 ) -> cute.TiledCopy:
-    copy_atom_C = cute.make_copy_atom(
-        warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2),
-        cutlass.Float16,  # this is just to get the right source layout
-    )
     tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
     return tiled_copy_C_atom
@@ -488,13 +609,12 @@ def get_smem_store_epi(
     epi_tile: cute.Shape,
     sC: Optional[cute.Tensor],
     tidx: Int32,
-    arch: int,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor, cute.Tensor]:
     dtype = sC.element_type if const_expr(sC is not None) else cutlass.Float16
     tiled_copy_C_atom = epilog_smem_copy_atom(tiled_mma, epi_tile)
-    copy_atom = get_smem_store_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_S(copy_atom, tiled_copy_C_atom)
     thr_copy = tiled_copy.get_slice(tidx)
     tRS_sC = None
@@ -515,11 +635,11 @@ def get_smem_store_epi(
 def get_smem_store_A(
-    tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, arch: int, position_independent=False
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
-    transpose = tiled_mma.op.a_major_mode == warpgroup.OperandMajorMode.MN
-    copy_atom = get_smem_store_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
@@ -537,13 +657,12 @@ def get_smem_load_A(
     tiled_mma: cute.TiledMma,
     sA: cute.Tensor,
     tidx: Int32,
-    arch: int,
     with_dst_tensor: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
-    transpose = tiled_mma.op.a_major_mode == warpgroup.OperandMajorMode.MN
-    copy_atom = get_smem_load_atom(arch, dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
@@ -563,6 +682,79 @@ def get_smem_load_A(
     return copy_fn if not with_dst_tensor else copy_fn_w_dst_tensor, thr_copy, tSR_sA
 @dsl_user_op
 def cpasync_reduce_bulk_add_f32(
     smem_ptr: cute.Pointer,
@@ -572,18 +764,14 @@ def cpasync_reduce_bulk_add_f32(
     loc=None,
     ip=None,
 ):
-    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value()
-    # cache_hint = cutlass.Int64(0x14F0000000000000)  # EVICT_LAST
-    llvm.inline_asm(
-        None,
-        [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()],
-        "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;",
-        "l,r,r",
-        # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()],
-        # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;",
-        # "l,r,r,l",
-        has_side_effects=True,
-        is_align_stack=False,
     )
@@ -674,33 +862,34 @@ def tma_gather4_load(
     """
     if len(row_indices) != 4:
         raise ValueError(f"gather4 requires exactly 4 row indices, got {len(row_indices)}")
-    col_val = Int32(col_idx).ir_value()
-    row_vals = [Int32(row_idx).ir_value() for row_idx in row_indices]
     # Convert pointers to integer addresses
-    desc_addr = tma_desc_ptr.toint(loc=loc, ip=ip).ir_value()
-    dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip).ir_value()
     mbar_addr = mbarrier_ptr.toint(loc=loc, ip=ip)
     if num_cta > 1:
         # Executed by both CTAs. Set peer bit to 0 so that the
         # transaction bytes will update CTA0's barrier.
         mbar_addr = mbar_addr & Sm100MmaPeerBitMask
-    mbar_addr = mbar_addr.ir_value()
     # Handle multicast_mask - may already be ir.Value or Python int
     multicast_mask_val = None
     if multicast_mask is not None:
-        multicast_mask_val = Int16(multicast_mask).ir_value()
     assert multicast_mask_val is None, "multicast is not supported yet"
     # Emit inline PTX for TMA gather4
     # PTX: cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes
     #      [dstMem], [tensorMap, {col, row0, row1, row2, row3}], [smem_bar];
     ptx = (
-        f"cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::{num_cta} "
-        "[$0], [$1, {$2, $3, $4, $5, $6}], [$7];"
     )
-    llvm.inline_asm(
-        None,
-        [
             dst_addr,
             desc_addr,
             col_val,
@@ -710,10 +899,6 @@ def tma_gather4_load(
             row_vals[3],
             mbar_addr,
         ],
-        ptx,
-        "r,l,r,r,r,r,r,r",  # constraints: register, long, 6x register
-        has_side_effects=True,
-        is_align_stack=False,
         loc=loc,
         ip=ip,
     )
@@ -723,34 +908,137 @@ def cpasync_bulk_get_copy_fn(
     src_tensor: cute.Tensor,
     dst_tensor: cute.Tensor,
     single_stage: bool = False,
     **kwargs,
 ) -> Callable:
     group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0))
     group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0))
     # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
     src = cute.group_modes(src_tensor, 0, group_rank_src)
     dst = cute.group_modes(dst_tensor, 0, group_rank_dst)
     def copy_bulk(src_idx, dst_idx, tma_bar_ptr: cute.Pointer, **new_kwargs):
         atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), src.element_type)
-        with cute.arch.elect_one():
-            cute.copy(
-                atom,
-                src[None, src_idx],
-                dst[None, dst_idx],
-                mbar_ptr=tma_bar_ptr,
-                **new_kwargs,
-                **kwargs,
-            )
     def copy_bulk_single_stage(tma_bar_ptr: cute.Pointer, **new_kwargs):
         atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), src.element_type)
-        with cute.arch.elect_one():
-            cute.copy(atom, src, dst, mbar_ptr=tma_bar_ptr, **new_kwargs, **kwargs)
     return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage
 @dsl_user_op
 def tma_get_copy_fn(
     atom: cute.CopyAtom,
@@ -800,6 +1088,236 @@ def tma_get_copy_fn(
     return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
 def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
     def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
         copy(
@@ -812,6 +1330,18 @@ def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsyn
     return copy_fn
 @cute.jit
 def gather_m_get_copy_fn(
     thr_copy_A: cute.ThrCopy,
@@ -855,7 +1385,7 @@ def gather_m_get_copy_fn(
     mA_k = cute.logical_divide(mA, (None, tile_K))
-    def copy_fn(src_idx, dst_idx, pred: bool = False):
         tApA_k = None
         if const_expr(pred):
             tApA_k = cute.make_rmem_tensor(cols_per_thread, Boolean)
@@ -964,9 +1494,7 @@ def gather_k_get_copy_fn(
         for k in cutlass.range(cols_per_thread):
             col_idx = tAcA[0, 0, k][1]
             k_idx[k] = sAIdx_cur[col_idx]
-        cute.arch.sync_warp()
-        with cute.arch.elect_one():
-            a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
         return k_idx, tApA_k
     def copy_fn(
@@ -1078,9 +1606,7 @@ def gather_k_get_tma_copy_fn(
     ) -> cute.Tensor:
         a_prefetch_pipeline.consumer_wait(a_prefetch_consumer_state)
         tSR_rAIdx = load_s2r(tSR_sAIdx[None, None, dst_idx])
-        cute.arch.sync_warp()
-        with cute.arch.elect_one():
-            a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
         return tSR_rAIdx
     def copy_fn(src_idx, dst_idx, tSR_rAIdx, tma_bar_ptr: cute.Pointer):

+# Copyright (c) 2025-2026, QuACK team.
+from typing import Any, Optional, Type, Tuple, Callable, Sequence
 from functools import partial
 import cutlass
 import cutlass.cute as cute
+import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass import Int32, Int16, Boolean, const_expr
+from cutlass.base_dsl.arch import Arch
+from cutlass.cute.nvgpu import cpasync, tcgen05, warp
 from cutlass.cute.nvgpu.tcgen05.mma import CtaGroup  # noqa
 from cutlass.cutlass_dsl import dsl_user_op
+from cutlass.utils import LayoutEnum, block_copy
 import cutlass.pipeline
 from cutlass._mlir import ir
 from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
 Sm100MmaPeerBitMask = 0xFEFFFFFF
+_TCGEN05_TMEM_OPS = (
+    tcgen05.Ld16x128bOp,
+    tcgen05.Ld16x256bOp,
+    tcgen05.Ld16x32bx2Op,
+    tcgen05.Ld16x64bOp,
+    tcgen05.Ld32x32bOp,
+    tcgen05.LdRed16x32bx2Op,
+    tcgen05.LdRed32x32bOp,
+    tcgen05.St16x128bOp,
+    tcgen05.St16x256bOp,
+    tcgen05.St16x32bx2Op,
+    tcgen05.St16x64bOp,
+    tcgen05.St32x32bOp,
+)
+_TCGEN05_TMEM_STORE_OPS = (
+    tcgen05.St16x128bOp,
+    tcgen05.St16x256bOp,
+    tcgen05.St16x32bx2Op,
+    tcgen05.St16x64bOp,
+    tcgen05.St32x32bOp,
+)
+def tmem_store_atom_from_load_atom(
+    copy_atom_t2r: Any,
+    src_dtype: Type[cutlass.Numeric],
+    dst_dtype: Type[cutlass.Numeric],
+) -> cute.CopyAtom:
+    """Return the matching tcgen05 R2T store atom for a selected T2R load atom.
+    `src_dtype` is the register fragment dtype loaded by T2R; `dst_dtype` is
+    the TMEM element dtype to store. Ratio 1 uses CUTLASS's operation-family
+    mapping directly. Ratio 2 is intentionally narrow: we allow the current
+    Ld32x32b path by halving repeat, and the widest 16dp path by halving the
+    vector width. Narrower 16dp cross-family mappings are not mirrored by
+    CUTLASS's same-family helper, so they assert until validated.
+    C++ CuTe's operation-family mapping is `cute::TMEM::tmem_load_to_store`:
+    https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/copy_traits_sm100.hpp#L3274
+    """
+    load_op = copy_atom_t2r.op if const_expr(hasattr(copy_atom_t2r, "op")) else copy_atom_t2r
+    if const_expr(hasattr(load_op, "op")):
+        load_op = load_op.op
+    assert src_dtype.width >= dst_dtype.width, "TMEM R2T helper only supports narrowing stores"
+    assert src_dtype.width % dst_dtype.width == 0, "TMEM source/destination widths must divide"
+    ratio = src_dtype.width // dst_dtype.width
+    assert ratio in (1, 2), "TMEM R2T helper only supports src/dst width ratio 1 or 2"
+    repeat = load_op.repeat
+    unpack = tcgen05.Unpack.NONE
+    if const_expr(getattr(load_op, "pack", None) == tcgen05.Pack.PACK_16b_IN_32b):
+        unpack = tcgen05.Unpack.UNPACK_32b_IN_16b
+    if const_expr(isinstance(load_op, tcgen05.Ld16x64bOp)):
+        assert ratio == 1, "No validated ratio-2 store mapping for Ld16x64bOp"
+        store_op = tcgen05.St16x64bOp(repeat, unpack)
+    elif const_expr(isinstance(load_op, tcgen05.Ld16x128bOp)):
+        assert ratio == 1, "No validated ratio-2 store mapping for Ld16x128bOp"
+        store_op = tcgen05.St16x128bOp(repeat, unpack)
+    elif const_expr(isinstance(load_op, tcgen05.Ld16x256bOp)):
+        store_op = (
+            tcgen05.St16x256bOp(repeat, unpack)
+            if const_expr(ratio == 1)
+            else tcgen05.St16x128bOp(repeat, unpack)
+        )
+    elif const_expr(isinstance(load_op, tcgen05.Ld16x32bx2Op)):
+        assert ratio == 1, "No validated ratio-2 store mapping for Ld16x32bx2Op"
+        store_op = tcgen05.St16x32bx2Op(repeat, unpack)
+    elif const_expr(isinstance(load_op, tcgen05.Ld32x32bOp)):
+        if const_expr(ratio == 2):
+            assert repeat.value % 2 == 0, "Ld32x32b ratio-2 store needs even repeat"
+            repeat = tcgen05.Repetition(repeat.value // 2)
+        store_op = tcgen05.St32x32bOp(repeat, unpack)
+    else:
+        raise TypeError(f"Unsupported TMEM load op for store conversion: {type(load_op)}")
+    return cute.make_copy_atom(store_op, dst_dtype)
+def _tmem_copy_reg_tv_layout(tiled_copy: cute.TiledCopy):
+    """Return the register-side TV layout for a tcgen05 tmem copy."""
+    op = tiled_copy.op
+    if const_expr(hasattr(op, "op")):
+        op = op.op
+    if const_expr(isinstance(op, _TCGEN05_TMEM_OPS)):
+        # TMEM stores read from registers; all other TMEM copy ops here write
+        # registers, including LdRed* reductions that upstream is_tmem_load
+        # intentionally does not classify as plain loads.
+        return (
+            tiled_copy.layout_src_tv_tiled
+            if const_expr(isinstance(op, _TCGEN05_TMEM_STORE_OPS))
+            else tiled_copy.layout_dst_tv_tiled
+        )
+    raise TypeError(f"Cannot infer tmem copy direction from tiled_copy.op={op}")
 @dsl_user_op
 ) -> None:
     assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem
     if const_expr(src.element_type != dst.element_type):
+        src = src.to(dst.element_type, loc=loc, ip=ip)
     if const_expr(retile):
         src = tiled_copy.retile(src)
     cute.copy(tiled_copy, src, dst, pred=pred, loc=loc, ip=ip, **kwargs)
     return dst
+@dsl_user_op
+def contiguous(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor:
+    dst = cute.make_rmem_tensor(src.shape, src.element_type, loc=loc, ip=ip)
+    cute.autovec_copy(src, dst, loc=loc, ip=ip)
+    return dst
 @dsl_user_op
 def load_s2r_retile(
     tiled_copy: cute.TiledCopy,
 @dsl_user_op
 def load_t2r(
+    tiled_copy: cute.TiledCopy,
+    src: cute.Tensor,
+    *,
+    fence: bool = False,
+    loc=None,
+    ip=None,
 ) -> cute.Tensor:
+    """Load one tmem tile partition into rmem, deriving the rmem shape from `src`.
+    `src` should already be indexed to the tile being copied, with any
+    stage/subtile modes removed.
+    """
+    dst = tmem_reg_frag(tiled_copy, src, loc=loc, ip=ip)
+    cute.copy(tiled_copy, src, dst, loc=loc, ip=ip)
+    if const_expr(fence):
+        cute.arch.fence_view_async_tmem_load()
     return dst
     return cute.make_tensor(cute.recast_ptr(tensor.iterator, dtype=tensor.element_type), new_layout)
+def partition_D_position_independent(thr_copy: cute.ThrCopy, tensor: cute.Tensor) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_D(tensor).iterator),
         thr_copy.partition_D(as_position_independent_swizzle_tensor(tensor)).layout,
     )
+def partition_S_position_independent(thr_copy: cute.ThrCopy, tensor: cute.Tensor) -> cute.Tensor:
     return cute.make_tensor(
         swizzle_ptr(thr_copy.partition_S(tensor).iterator),
         thr_copy.partition_S(as_position_independent_swizzle_tensor(tensor)).layout,
 def get_smem_store_atom(
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
+    arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum()
+    if const_expr(arch < Arch.sm_90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
 def get_smem_load_atom(
     element_type: Type[cute.Numeric],
     transpose: bool = False,
     major_mode_size: Optional[int] = None,
 ) -> cute.CopyAtom:
+    arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum()
+    if const_expr(arch < Arch.sm_90 or element_type.width != 16):
         return cute.make_copy_atom(
             cute.nvgpu.CopyUniversalOp(),
             element_type,
 def get_smem_store_C(
+    tiled_mma: cute.TiledMma | cute.TiledCopy,
     sC: cute.Tensor,
     tidx: Int32,
     transpose: bool = False,
     position_independent=False,
     major_mode_size: Optional[int] = None,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
+    if const_expr(isinstance(tiled_mma, cute.TiledCopy)):
+        tiled_copy_t2r = tiled_mma
+        layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR
+        copy_atom = sm100_utils.get_smem_store_op(
+            layout, dtype, tiled_copy_t2r.value_type, tiled_copy_t2r
+        )
+        tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r)
+    else:
+        copy_atom = get_smem_store_atom(dtype, transpose, major_mode_size=major_mode_size)
+        tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tRS_sC = thr_copy.partition_D(sC)
     else:
         tRS_sC = partition_D_position_independent(thr_copy, sC)
+    def copy_fn(src: cute.Tensor, dst_idx: Optional[Int32] = None, fence=False, **new_kwargs):
+        dst_tensor = tRS_sC if const_expr(dst_idx is None) else tRS_sC[..., dst_idx]
         cvt_copy(tiled_copy, src, dst_tensor, retile=True, **new_kwargs)
+        if const_expr(fence):
+            cute.arch.fence_view_async_shared()
     return copy_fn, thr_copy, tRS_sC
     tiled_mma: cute.TiledMma,
     sC: cute.Tensor,
     tidx: Int32,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sC.element_type
+    copy_atom = get_smem_load_atom(dtype, transpose)
     tiled_copy = cute.make_tiled_copy_C(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
         tSR_sC = thr_copy.partition_S(sC)
     else:
         tSR_sC = partition_S_position_independent(thr_copy, sC)
+    copy_atom_RS = get_smem_store_atom(dtype, transpose)
     thr_copy_RS = cute.make_tiled_copy_C(copy_atom_RS, tiled_mma).get_slice(tidx)
     tRS_shape = thr_copy_RS.partition_S(cute.make_identity_tensor(sC.shape[:2])).shape
 def epilog_smem_copy_atom(
     tiled_mma: cute.TiledMma, epi_tile: cute.Shape, transpose: bool = False
 ) -> cute.TiledCopy:
+    arch = cutlass.base_dsl.BaseDSL._get_dsl().get_arch_enum()
+    if const_expr(arch < Arch.sm_90):
+        copy_atom_C = cute.make_copy_atom(
+            cute.nvgpu.CopyUniversalOp(),
+            cutlass.Float16,  # this is just to get the right source layout
+            num_bits_per_copy=(2 if not transpose else 1) * cutlass.Float16.width,
+        )
+    else:
+        copy_atom_C = cute.make_copy_atom(
+            warp.StMatrix8x8x16bOp(transpose, num_matrices=4 if epi_tile[1] % 16 == 0 else 2),
+            cutlass.Float16,  # this is just to get the right source layout
+        )
     tiled_copy_C_atom = cute.make_tiled_copy_C_atom(copy_atom_C, tiled_mma)
     return tiled_copy_C_atom
     epi_tile: cute.Shape,
     sC: Optional[cute.Tensor],
     tidx: Int32,
     transpose: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor, cute.Tensor]:
     dtype = sC.element_type if const_expr(sC is not None) else cutlass.Float16
+    copy_atom = get_smem_store_atom(dtype, transpose)
     tiled_copy_C_atom = epilog_smem_copy_atom(tiled_mma, epi_tile)
     tiled_copy = cute.make_tiled_copy_S(copy_atom, tiled_copy_C_atom)
     thr_copy = tiled_copy.get_slice(tidx)
     tRS_sC = None
 def get_smem_store_A(
+    tiled_mma: cute.TiledMma, sA: cute.Tensor, tidx: Int32, position_independent=False
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
+    transpose = tiled_mma.op.a_major_mode == cute.nvgpu.OperandMajorMode.MN
+    copy_atom = get_smem_store_atom(dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
     tiled_mma: cute.TiledMma,
     sA: cute.Tensor,
     tidx: Int32,
     with_dst_tensor: bool = False,
     position_independent=False,
 ) -> Tuple[Callable, cute.TiledCopy, cute.Tensor]:
     dtype = sA.element_type
+    transpose = tiled_mma.op.a_major_mode == cute.nvgpu.OperandMajorMode.MN
+    copy_atom = get_smem_load_atom(dtype, transpose)
     tiled_copy = cute.make_tiled_copy_A(copy_atom, tiled_mma)
     thr_copy = tiled_copy.get_slice(tidx)
     if const_expr(not position_independent):
     return copy_fn if not with_dst_tensor else copy_fn_w_dst_tensor, thr_copy, tSR_sA
+def _cpasync_reduction_kind_name(reduction_kind: Any) -> str:
+    name = (
+        reduction_kind.lower() if isinstance(reduction_kind, str) else reduction_kind.name.lower()
+    )
+    assert name in {"add", "min", "max", "inc", "dec", "and", "or", "xor"}, (
+        f"Unsupported cp.reduce.async.bulk reduction kind: {reduction_kind}"
+    )
+    return name
+def _cpasync_bulk_reduce_suffix(
+    reduction_kind: Any,
+    dtype: Type[cutlass.Numeric],
+) -> str:
+    op = _cpasync_reduction_kind_name(reduction_kind)
+    if dtype is cutlass.Float16:
+        assert op in {"add", "min", "max"}, f"{op} is not supported for f16 bulk reduce"
+        return f"{op}.noftz.f16" if op == "add" else f"{op}.f16"
+    if dtype is cutlass.BFloat16:
+        assert op in {"add", "min", "max"}, f"{op} is not supported for bf16 bulk reduce"
+        return f"{op}.noftz.bf16" if op == "add" else f"{op}.bf16"
+    if dtype is cutlass.Float32:
+        assert op == "add", f"{op} is not supported for f32 bulk reduce"
+        return "add.f32"
+    if dtype is cutlass.Float64:
+        assert op == "add", f"{op} is not supported for f64 bulk reduce"
+        return "add.f64"
+    signed = getattr(dtype, "signed", None)
+    width = getattr(dtype, "width", None)
+    if signed is not None:
+        assert width in (32, 64), f"Unsupported integer bulk-reduce width: {width}"
+        if op in {"and", "or", "xor"}:
+            return f"{op}.b{width}"
+        if op in {"min", "max", "add"}:
+            return f"{op}.{'s' if signed else 'u'}{width}"
+        assert op in {"inc", "dec"} and dtype is cutlass.Uint32, (
+            f"{op} bulk reduce is only supported for u32"
+        )
+        return f"{op}.u32"
+    raise TypeError(f"Unsupported cp.reduce.async.bulk dtype: {dtype}")
+@dsl_user_op
+def cpasync_bulk_s2g(
+    smem_ptr: cute.Pointer,
+    gmem_ptr: cute.Pointer,
+    store_bytes: int | Int32,
+    *,
+    reduction_kind: Optional[Any] = None,
+    dtype: Optional[Type[cutlass.Numeric]] = None,
+    loc=None,
+    ip=None,
+):
+    smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip)
+    if reduction_kind is None:
+        ptx = "cp.async.bulk.global.shared::cta.bulk_group [{$r0}], [{$r1}], {$r2};"
+    else:
+        assert dtype is not None, "dtype is required for cp.reduce.async.bulk"
+        ptx = (
+            "cp.reduce.async.bulk.global.shared::cta.bulk_group."
+            f"{_cpasync_bulk_reduce_suffix(reduction_kind, dtype)} "
+            "[{$r0}], [{$r1}], {$r2};"
+        )
+    cute.arch.inline_ptx(
+        ptx,
+        read_only_args=[gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes)],
+        loc=loc,
+        ip=ip,
+    )
 @dsl_user_op
 def cpasync_reduce_bulk_add_f32(
     smem_ptr: cute.Pointer,
     loc=None,
     ip=None,
 ):
+    cpasync_bulk_s2g(
+        smem_ptr,
+        gmem_ptr,
+        store_bytes,
+        reduction_kind=cpasync.ReductionOp.ADD,
+        dtype=cutlass.Float32,
+        loc=loc,
+        ip=ip,
     )
     """
     if len(row_indices) != 4:
         raise ValueError(f"gather4 requires exactly 4 row indices, got {len(row_indices)}")
+    col_val = Int32(col_idx)
+    row_vals = [Int32(row_idx) for row_idx in row_indices]
     # Convert pointers to integer addresses
+    desc_addr = tma_desc_ptr.toint(loc=loc, ip=ip)
+    dst_addr = dst_smem_ptr.toint(loc=loc, ip=ip)
     mbar_addr = mbarrier_ptr.toint(loc=loc, ip=ip)
     if num_cta > 1:
         # Executed by both CTAs. Set peer bit to 0 so that the
         # transaction bytes will update CTA0's barrier.
         mbar_addr = mbar_addr & Sm100MmaPeerBitMask
+    mbar_addr = Int32(mbar_addr)
     # Handle multicast_mask - may already be ir.Value or Python int
     multicast_mask_val = None
     if multicast_mask is not None:
+        multicast_mask_val = Int16(multicast_mask)
     assert multicast_mask_val is None, "multicast is not supported yet"
     # Emit inline PTX for TMA gather4
     # PTX: cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes
     #      [dstMem], [tensorMap, {col, row0, row1, row2, row3}], [smem_bar];
     ptx = (
+        "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes."
+        f"cta_group::{num_cta} "
+        "[{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];"
     )
+    cute.arch.inline_ptx(
+        ptx,
+        read_only_args=[
             dst_addr,
             desc_addr,
             col_val,
             row_vals[3],
             mbar_addr,
         ],
         loc=loc,
         ip=ip,
     )
     src_tensor: cute.Tensor,
     dst_tensor: cute.Tensor,
     single_stage: bool = False,
+    reduction_kind: Optional[cute.nvgpu.cpasync.ReductionKind] = None,
     **kwargs,
 ) -> Callable:
+    src_is_smem = const_expr(
+        isinstance(src_tensor.iterator, cute.Pointer)
+        and src_tensor.memspace == cute.AddressSpace.smem
+    )
+    dst_is_smem = const_expr(
+        isinstance(dst_tensor.iterator, cute.Pointer)
+        and dst_tensor.memspace == cute.AddressSpace.smem
+    )
+    if const_expr(reduction_kind is not None):
+        assert src_is_smem and not dst_is_smem, "cp.reduce.async.bulk only supports SMEM -> GMEM"
     group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0))
     group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0))
     # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK)
     src = cute.group_modes(src_tensor, 0, group_rank_src)
     dst = cute.group_modes(dst_tensor, 0, group_rank_dst)
+    if const_expr(src_is_smem and not dst_is_smem):
+        def copy_bulk_s2g(src_idx, dst_idx, **new_kwargs):
+            store_bytes = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8)
+            with cute.arch.elect_one():
+                cpasync_bulk_s2g(
+                    src[None, src_idx].iterator,
+                    dst[None, dst_idx].iterator,
+                    store_bytes,
+                    reduction_kind=reduction_kind,
+                    dtype=src.element_type,
+                    **new_kwargs,
+                    **kwargs,
+                )
+        def copy_bulk_s2g_single_stage(**new_kwargs):
+            store_bytes = const_expr(cute.size(src.shape) * src.element_type.width // 8)
+            with cute.arch.elect_one():
+                cpasync_bulk_s2g(
+                    src.iterator,
+                    dst.iterator,
+                    store_bytes,
+                    reduction_kind=reduction_kind,
+                    dtype=src.element_type,
+                    **new_kwargs,
+                    **kwargs,
+                )
+        return copy_bulk_s2g if const_expr(not single_stage) else copy_bulk_s2g_single_stage
     def copy_bulk(src_idx, dst_idx, tma_bar_ptr: cute.Pointer, **new_kwargs):
+        assert dst_is_smem and not src_is_smem, "cp.async.bulk G2S expects GMEM -> SMEM"
         atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), src.element_type)
+        cute.copy(
+            atom,
+            src[None, src_idx],
+            dst[None, dst_idx],
+            mbar_ptr=tma_bar_ptr,
+            **new_kwargs,
+            **kwargs,
+        )
     def copy_bulk_single_stage(tma_bar_ptr: cute.Pointer, **new_kwargs):
+        assert dst_is_smem and not src_is_smem, "cp.async.bulk G2S expects GMEM -> SMEM"
         atom = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), src.element_type)
+        cute.copy(atom, src, dst, mbar_ptr=tma_bar_ptr, **new_kwargs, **kwargs)
     return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage
+def cpasync_bulk_get_store_or_add_fn(
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    store_first_contribution: bool,
+    single_stage: bool = False,
+    **kwargs,
+) -> Callable:
+    assert not single_stage, "store-or-add helper only supports staged SMEM -> GMEM tensors"
+    src_is_smem = const_expr(
+        isinstance(src_tensor.iterator, cute.Pointer)
+        and src_tensor.memspace == cute.AddressSpace.smem
+    )
+    dst_is_smem = const_expr(
+        isinstance(dst_tensor.iterator, cute.Pointer)
+        and dst_tensor.memspace == cute.AddressSpace.smem
+    )
+    assert src_is_smem and not dst_is_smem, "store-or-add helper only supports SMEM -> GMEM"
+    group_rank_src = const_expr(cute.rank(src_tensor))
+    group_rank_dst = const_expr(cute.rank(dst_tensor))
+    src = cute.group_modes(src_tensor, 0, group_rank_src - 1)
+    dst = cute.group_modes(dst_tensor, 0, group_rank_dst - 1)
+    @cute.jit
+    def copy_bulk_s2g_store_or_add(src_idx, dst_idx, idx, **new_kwargs):
+        store_bytes = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8)
+        src_ptr = src[None, src_idx].iterator
+        dst_ptr = dst[None, dst_idx].iterator
+        with cute.arch.elect_one():
+            if const_expr(store_first_contribution):
+                if idx == 0:
+                    cpasync_bulk_s2g(
+                        src_ptr,
+                        dst_ptr,
+                        store_bytes,
+                        reduction_kind=None,
+                        **new_kwargs,
+                        **kwargs,
+                    )
+                else:
+                    cpasync_bulk_s2g(
+                        src_ptr,
+                        dst_ptr,
+                        store_bytes,
+                        reduction_kind=cpasync.ReductionOp.ADD,
+                        dtype=src.element_type,
+                        **new_kwargs,
+                        **kwargs,
+                    )
+            else:
+                cpasync_bulk_s2g(
+                    src_ptr,
+                    dst_ptr,
+                    store_bytes,
+                    reduction_kind=cpasync.ReductionOp.ADD,
+                    dtype=src.element_type,
+                    **new_kwargs,
+                    **kwargs,
+                )
+    return copy_bulk_s2g_store_or_add
 @dsl_user_op
 def tma_get_copy_fn(
     atom: cute.CopyAtom,
     return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g
+@dsl_user_op
+def tma_get_block_copy_fn(
+    atom: cute.CopyAtom,
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    tma_multicast: Optional[dict] = None,
+    single_stage: bool = False,
+    *,
+    loc=None,
+    ip=None,
+    **kwargs,
+) -> Callable:
+    src_is_smem = const_expr(
+        isinstance(src_tensor.iterator, cute.Pointer)
+        and src_tensor.memspace == cute.AddressSpace.smem
+    )
+    if const_expr(tma_multicast is not None and "use_2cta_mma_inst" not in tma_multicast):
+        op = atom.op if const_expr(hasattr(atom, "op")) else atom
+        tma_multicast = {
+            **tma_multicast,
+            "use_2cta_mma_inst": getattr(op, "cta_group", None) == tcgen05.CtaGroup.TWO,
+        }
+    smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor)
+    group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0))
+    group_rank_gmem = const_expr(cute.rank(gmem_tensor) - (1 if not single_stage else 0))
+    s = cute.group_modes(smem_tensor, 0, group_rank_smem)
+    g = cute.group_modes(gmem_tensor, 0, group_rank_gmem)
+    src, dst = (s, g) if src_is_smem else (g, s)
+    @dsl_user_op
+    def copy_tma(src_idx, dst_idx, *, loc=None, ip=None, **new_kwargs):
+        src_cur = src[None, src_idx]
+        dst_cur = dst[None, dst_idx]
+        if const_expr(tma_multicast is None):
+            block_copy(atom, src_cur, dst_cur, **new_kwargs, **kwargs, loc=loc, ip=ip)
+        else:
+            block_copy(
+                atom,
+                src_cur,
+                dst_cur,
+                tma_multicast=tma_multicast,
+                **new_kwargs,
+                **kwargs,
+                loc=loc,
+                ip=ip,
+            )
+    @dsl_user_op
+    def copy_tma_single_stage(*, loc=None, ip=None, **new_kwargs):
+        if const_expr(tma_multicast is None):
+            block_copy(atom, src, dst, **new_kwargs, **kwargs, loc=loc, ip=ip)
+        else:
+            block_copy(
+                atom,
+                src,
+                dst,
+                tma_multicast=tma_multicast,
+                **new_kwargs,
+                **kwargs,
+                loc=loc,
+                ip=ip,
+            )
+    return copy_tma if const_expr(not single_stage) else copy_tma_single_stage
+def s2t_get_copy_fn(
+    src_tensor: cute.Tensor,
+    dst_tensor: cute.Tensor,
+    cta_group: tcgen05.CtaGroup,
+) -> Callable:
+    """
+    Make tiledCopy for smem to tmem load, then return a copy function over stages.
+    :param src_tensor: The source tensor in smem
+    :param dst_tensor: The destination tensor in tmem
+    """
+    assert src_tensor.element_type == dst_tensor.element_type
+    # (MMA, MMA_MN, MMA_K, STAGE)
+    src_compact = cute.filter_zeros(src_tensor)
+    # (MMA, MMA_MN, MMA_K)
+    dst_compact = cute.filter_zeros(dst_tensor)
+    # Make S2T CopyAtom and tiledCopy.
+    copy_atom = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(cta_group), dst_tensor.element_type)
+    tiled_copy = tcgen05.make_s2t_copy(copy_atom, dst_compact)
+    thr_copy = tiled_copy.get_slice(0)
+    # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
+    src_partition = tcgen05.get_s2t_smem_desc_tensor(tiled_copy, thr_copy.partition_S(src_compact))
+    # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+    dst_partition = thr_copy.partition_D(dst_compact)
+    @dsl_user_op
+    def copy_s2t(stage_idx, *, loc=None, ip=None, **new_kwargs):
+        # Stage slice of partitioned source tensor: ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
+        stage_coord = (None, None, None, None, stage_idx)
+        cute.copy(
+            tiled_copy, src_partition[stage_coord], dst_partition, loc=loc, ip=ip, **new_kwargs
+        )
+    return copy_s2t
+# tcgen05 TMEM <-> RMEM helpers (t2r loads / r2t stores).
+#
+# The register-side fragment of a tmem copy is derivable by layout algebra,
+# with no reference to the original (pre-partition) tile tensor:
+# - the per-thread register VALUE shape is mode 1 of the tiled copy's
+#   register-side TV layout (`layout_dst_tv_tiled` for loads,
+#   `layout_src_tv_tiled` for stores). The tmem-side partition can't supply
+#   it: tmem partitioning is warp-collective, so its value mode counts tmem
+#   cells across the whole warp, not per-thread register elements.
+# - the tile-iteration modes are shared between partition_S and partition_D
+#   (same tiler over the same tile extent), so they can be read off whichever
+#   side was already partitioned.
+# This kills the make-a-fake/identity-tensor-and-partition_D dance previously
+# needed at every t2r site.
+def tmem_reg_frag(
+    tiled_copy: cute.TiledCopy,
+    partitioned: cute.Tensor,
+    num_extra_modes: int = 0,
+    dtype: Optional[Type[cutlass.Numeric]] = None,
+    *,
+    loc=None,
+    ip=None,
+) -> cute.Tensor:
+    """Allocate the per-thread register fragment for ONE tile of a tcgen05
+    tmem copy, given any partitioned view of it (tmem or otherwise).
+    `partitioned` is (V, iter..., extra...) as produced by partition_S/_D;
+    the trailing `num_extra_modes` modes (stage, epi-subtile, ...) are
+    excluded from the fragment and indexed at copy time instead. `dtype`
+    defaults to `partitioned.element_type`.
+    The register side is inferred from the tcgen05 load/store op: destination
+    for t2r loads, source for r2t stores."""
+    tv = _tmem_copy_reg_tv_layout(tiled_copy)
+    val_shape = tv.shape[1]
+    rank = cute.rank(partitioned.shape)
+    iters = tuple(partitioned.shape[i] for i in range(1, rank - num_extra_modes))
+    frag_dtype = partitioned.element_type if const_expr(dtype is None) else dtype
+    return cute.make_rmem_tensor((val_shape, *iters), frag_dtype, loc=loc, ip=ip)
+def coord_frag(tiled_copy: cute.TiledCopy, tidx: Int32, shape) -> cute.Tensor:
+    """Per-thread (row, col) coordinates aligned with a tiled copy's register
+    fragments (`tmem_reg_frag` / `load_t2r`): the register-side partition of
+    an identity tensor over `shape`. Deliberately partition_D — partition_S of
+    a TMEM tiled copy keeps whole warp-addressed atom tiles instead of
+    distributing elements over lanes."""
+    return tiled_copy.get_slice(tidx).partition_D(cute.make_identity_tensor(shape))
+def r2s_partition_from_t2r(
+    tiled_copy_t2r: cute.TiledCopy,
+    s: cute.Tensor,
+    tidx: Int32,
+    transpose: bool = False,
+    position_independent=False,
+) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
+    """SMEM-store (r2s) side chained off a tmem-load tiled copy: the r2s copy
+    inherits the t2r copy's per-thread value ownership via make_tiled_copy_D,
+    so the loaded fragment can be stored (post-conversion) without a shuffle.
+    By default the store atom is selected like SM100 GEMM epilogues:
+    `get_smem_store_op(layout, dst_dtype, tiled_copy_t2r.value_type, tiled_copy_t2r)`,
+    so the stmatrix shape follows the tmem-load atom. `transpose=True` maps to
+    COL_MAJOR, otherwise ROW_MAJOR.
+    `s` is the staged SMEM tile; its trailing stage mode is excluded from the
+    register fragment. `position_independent=True` partitions through a
+    position-independent swizzle view, matching `get_smem_store_C`.
+    Returns `(tiled_copy, tRS_r, tRS_s)`; store via
+    `cute.copy(tiled_copy, tRS_r, tRS_s[..., idx])`."""
+    dtype = s.element_type
+    layout = LayoutEnum.COL_MAJOR if const_expr(transpose) else LayoutEnum.ROW_MAJOR
+    copy_atom = sm100_utils.get_smem_store_op(
+        layout, dtype, tiled_copy_t2r.value_type, tiled_copy_t2r
+    )
+    tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r)
+    thr_copy = tiled_copy.get_slice(tidx)
+    if const_expr(not position_independent):
+        tRS_s = thr_copy.partition_D(s)
+    else:
+        tRS_s = partition_D_position_independent(thr_copy, s)
+    rank = cute.rank(tRS_s.shape)
+    frag_shape = tuple(tRS_s.shape[i] for i in range(rank - 1))
+    tRS_r = cute.make_rmem_tensor(frag_shape, dtype)
+    return tiled_copy, tRS_r, tRS_s
+def s2r_partition_from_t2r(
+    tiled_copy_t2r: cute.TiledCopy,
+    s: cute.Tensor,
+    tidx: Int32,
+    r_layout: cute.Layout,
+    copy_atom: Optional[cute.CopyAtom] = None,
+    transpose: bool = False,
+    position_independent=False,
+) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]:
+    """SMEM-load (s2r) counterpart of `r2s_partition_from_t2r` (ldmatrix vs
+    stmatrix), for reading an epilogue input that was TMA-staged into an
+    epi-tile SMEM buffer (e.g. C in gemm, z in ssd) into registers
+    element-aligned with the t2r fragments. The register fragment is
+    allocated with `r_layout` (pass the r2s fragment's layout) so its linear
+    element order matches the t2r/r2s fragments; `tSR_r` is its retiled view
+    for the collective copy. Per-warp SMEM footprints of this load and the
+    chained r2s store coincide, so reusing one buffer for input then output
+    is warp-local (no inter-warp hazard).
+    `position_independent=True` partitions through a position-independent
+    swizzle view, matching `get_smem_load_C`.
+    Returns `(tiled_copy, tRS_r, tSR_r, tSR_s)`; load via
+    `cute.copy(tiled_copy, tSR_s[..., idx], tSR_r)` then read `tRS_r`."""
+    dtype = s.element_type
+    if const_expr(copy_atom is None):
+        copy_atom = cute.make_copy_atom(
+            warp.LdMatrix8x8x16bOp(transpose=transpose, num_matrices=4), dtype
+        )
+    tiled_copy = cute.make_tiled_copy_D(copy_atom, tiled_copy_t2r)
+    thr_copy = tiled_copy.get_slice(tidx)
+    if const_expr(not position_independent):
+        tSR_s = thr_copy.partition_S(s)
+    else:
+        tSR_s = partition_S_position_independent(thr_copy, s)
+    tRS_r = cute.make_rmem_tensor(r_layout, dtype)
+    tSR_r = tiled_copy.retile(tRS_r)
+    return tiled_copy, tRS_r, tSR_r, tSR_s
 def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync):
     def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
         copy(
     return copy_fn
+def chain_tma_producer_copy_fns(copy_fns: Sequence[Optional[Callable]]):
+    if not any(fn is not None for fn in copy_fns):
+        return None
+    def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs):
+        for fn in copy_fns:
+            if const_expr(fn is not None):
+                fn(src_idx=src_idx, producer_state=producer_state, **new_kwargs)
+    return copy_fn
 @cute.jit
 def gather_m_get_copy_fn(
     thr_copy_A: cute.ThrCopy,
     mA_k = cute.logical_divide(mA, (None, tile_K))
+    def copy_fn(src_idx, dst_idx, pred: cutlass.Constexpr[bool] = False):
         tApA_k = None
         if const_expr(pred):
             tApA_k = cute.make_rmem_tensor(cols_per_thread, Boolean)
         for k in cutlass.range(cols_per_thread):
             col_idx = tAcA[0, 0, k][1]
             k_idx[k] = sAIdx_cur[col_idx]
+        a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
         return k_idx, tApA_k
     def copy_fn(
     ) -> cute.Tensor:
         a_prefetch_pipeline.consumer_wait(a_prefetch_consumer_state)
         tSR_rAIdx = load_s2r(tSR_sAIdx[None, None, dst_idx])
+        a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
         return tSR_rAIdx
     def copy_fn(src_idx, dst_idx, tSR_rAIdx, tma_bar_ptr: cute.Pointer):

build/torch-cuda/quack/cross_entropy.py CHANGED Viewed

@@ -5,7 +5,7 @@ from functools import partial
 from typing import Optional, Type, Literal
 import torch
-from ._ops_compat import add_quack_op_namespace_prefix
 from torch import Tensor
 import cuda.bindings.driver as cuda
@@ -18,11 +18,12 @@ from . import utils as utils
 from . import copy_utils as copy_utils
 from . import layout_utils as layout_utils
 from .compile_utils import make_fake_tensor as fake_tensor
 from .reduce import row_reduce, online_softmax_reduce
 from .reduction_base import ReductionBase
-from .cache_utils import jit_cache
 from .cute_dsl_utils import torch2cute_dtype_map
-from cutlass.base_dsl import Arch
 class CrossEntropy(ReductionBase):
@@ -75,6 +76,7 @@ class CrossEntropy(ReductionBase):
         mLoss: cute.Tensor,  # (M,)
         mLSE: Optional[cute.Tensor],  # (M,)
         mdX: Optional[cute.Tensor],  # (M, N) - if provided, compute gradient
         ignore_index: Int32,  # Index to ignore in loss computation
         stream: cuda.CUstream,
     ):
@@ -97,6 +99,7 @@ class CrossEntropy(ReductionBase):
             mLoss,
             mLSE,
             mdX,
             ignore_index,
             tiler_mn,
             tiled_copy,
@@ -117,6 +120,7 @@ class CrossEntropy(ReductionBase):
         mLoss: cute.Tensor,  # (M,)
         mLSE: Optional[cute.Tensor],  # (M,)
         mdX: Optional[cute.Tensor],  # (M, N) - if provided, compute gradient
         ignore_index: Int32,  # Index to ignore in loss computation
         tiler_mn: cute.Shape,
         tiled_copy: cute.TiledCopy,
@@ -156,8 +160,16 @@ class CrossEntropy(ReductionBase):
         row = tXcX[0][0]
         target = Int32.zero
         if row < shape[0]:
             target = Int32(mTarget[row])
         if row < shape[0]:
             copy(tXgX, tXsX, is_async=True)
@@ -221,7 +233,7 @@ class CrossEntropy(ReductionBase):
         ):
             lse = max_x + cute.math.log(denom, fastmath=True)
             # Set loss to 0 if this index should be ignored, otherwise compute normally
-            loss_val = (lse - target_logit) if not should_ignore else Float32.zero
             mLoss[row] = mLoss.element_type(loss_val)
             if const_expr(mLSE is not None):
                 mLSE[row] = lse
@@ -239,7 +251,6 @@ class CrossEntropy(ReductionBase):
             probs = exp_x * denom_inv
             gdX = cute.local_tile(mdX, tiler_mn, (bidx, cluster_y))
             tXgdX = thr_copy.partition_D(gdX)
-            tXrdX = cute.make_rmem_tensor_like(tXgdX)
             tXcFull = thr_copy.partition_S(cX)
             # Compute gradient: probs for all classes, (probs - 1) for target class
             # If ignored, gradient is already zero
@@ -248,46 +259,58 @@ class CrossEntropy(ReductionBase):
             if not should_ignore:
                 for i in cutlass.range(cute.size(tXrX), unroll_full=True):
                     tXrdX_f32[i] = tXrdX_f32[i] if tXcFull[i][1] != target else tXrdX_f32[i] - 1.0
-            tXrdX.store(tXrdX_f32.load().to(tXrdX.element_type))
             if row < shape[0]:
                 copy(tXrdX, tXgdX)
-@jit_cache
-def _compile_cross_entropy_fwd(
-    dtype, target_dtype, target_logit_dtype, N, has_lse, has_dx, target_logit_ndim
-):
-    batch_sym = cute.sym_int()
-    div = math.gcd(128 // dtype.width, N)
-    x_cute = fake_tensor(dtype, (batch_sym, N), div)
-    dx_cute = fake_tensor(dtype, (batch_sym, N), div) if has_dx else None
-    target_cute = fake_tensor(target_dtype, (batch_sym,))
-    if target_logit_dtype is not None:
-        if target_logit_ndim == 2:
-            target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym, cute.sym_int()), div)
         else:
-            target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,))
-    else:
-        target_logit_cute = None
-    loss_cute = fake_tensor(Float32, (batch_sym,))
-    lse_cute = fake_tensor(Float32, (batch_sym,)) if has_lse else None
-    # If there's dx, it's faster to not use online softmax since we want the exp(x - max)
-    cross_entropy_op = CrossEntropy(dtype, N, online_softmax=not has_dx)
-    return cute.compile(
-        cross_entropy_op,
-        x_cute,
-        target_cute,
-        target_logit_cute,
-        loss_cute,
-        lse_cute,
-        dx_cute,
-        Int32(0),  # ignore_index, just for compilation
-        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
-        options="--enable-tvm-ffi",
-    )
-@torch.library.custom_op(add_quack_op_namespace_prefix("cross_entropy_fwd_out"), mutates_args={"loss", "lse", "dx"})
 def cross_entropy_fwd_out(
     x: Tensor,
     target: Tensor,
@@ -295,6 +318,7 @@ def cross_entropy_fwd_out(
     loss: Tensor,
     lse: Optional[Tensor],
     dx: Optional[Tensor],
     ignore_index: int = -100,
 ) -> None:
     """Cross entropy forward pass.
@@ -307,6 +331,7 @@ def cross_entropy_fwd_out(
         loss: Output loss tensor of shape (M,)
         lse: Optional output log-sum-exp tensor of shape (M,)
         dx: Optional output gradient tensor of shape (M, N)
         ignore_index: Index to ignore in loss computation
     Returns:
@@ -314,14 +339,12 @@ def cross_entropy_fwd_out(
     """
     assert x.dim() == 2, "Input must be 2D"
     assert target.dim() == 1, "Target must be 1D"
-    assert x.is_cuda and target.is_cuda, "Tensors must be on CUDA device"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype"
     assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64"
     if target_logit is not None:
-        assert target_logit.is_cuda, "Target logits must be on CUDA device"
         assert target_logit.dtype in [torch.float16, torch.bfloat16, torch.float32]
-    if dx is not None:
-        assert dx.is_cuda, "dx must be on CUDA device"
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
     target_dtype = torch2cute_dtype_map[target.dtype]
@@ -329,54 +352,24 @@ def cross_entropy_fwd_out(
         torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None
     )
     target_logit_ndim = target_logit.ndim if target_logit is not None else None
-    _compile_cross_entropy_fwd(
         dtype,
         target_dtype,
         target_logit_dtype,
         N,
         lse is not None,
         dx is not None,
         target_logit_ndim,
-    )(x, target, target_logit, loss, lse, dx, Int32(ignore_index))
-@cross_entropy_fwd_out.register_fake
-def _cross_entropy_fwd_out_fake(
-    x: Tensor,
-    target: Tensor,
-    target_logit: Optional[Tensor],
-    loss: Tensor,
-    lse: Optional[Tensor],
-    dx: Optional[Tensor],
-    ignore_index: int = -100,
-) -> None:
-    # See softmax.py _softmax_fwd_fake for why register_fake is needed.
-    from .cache_utils import COMPILE_ONLY
-    if COMPILE_ONLY and not isinstance(x.size(1), torch.SymInt):
-        N = x.size(1)
-        dtype = torch2cute_dtype_map[x.dtype]
-        target_dtype = torch2cute_dtype_map[target.dtype]
-        target_logit_dtype = (
-            torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None
-        )
-        target_logit_ndim = target_logit.ndim if target_logit is not None else None
-        _compile_cross_entropy_fwd(
-            dtype,
-            target_dtype,
-            target_logit_dtype,
-            N,
-            lse is not None,
-            dx is not None,
-            target_logit_ndim,
-        )
-        _compile_cross_entropy_backward(dtype, target_dtype, N)
 def cross_entropy_fwd(
     x: torch.Tensor,
     target: torch.Tensor,
     target_logit: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
     return_lse: bool = False,
     return_dx: bool = False,
@@ -387,7 +380,7 @@ def cross_entropy_fwd(
     loss = torch.empty(M, device=device, dtype=torch.float32)
     lse = torch.empty(M, device=device, dtype=torch.float32) if return_lse else None
     dx = (torch.empty_like(x) if not inplace_backward else x) if return_dx else None
-    cross_entropy_fwd_out(x, target, target_logit, loss, lse, dx, ignore_index)
     if return_lse and return_dx:
         return loss, lse, dx
     elif return_lse:
@@ -432,6 +425,7 @@ class CrossEntropyBackward:
         mDLoss: cute.Tensor,
         mdX: cute.Tensor,
         mLSE: cute.Tensor,
         ignore_index: Int32,  # Index to ignore in gradient computation
         stream: cuda.CUstream,
     ):
@@ -451,6 +445,7 @@ class CrossEntropyBackward:
             mDLoss,
             mdX,
             mLSE,
             ignore_index,
             mX.shape,
             tiler_mn,
@@ -474,6 +469,7 @@ class CrossEntropyBackward:
         mDLoss: cute.Tensor,  # (M,)
         mdX: cute.Tensor,  # (M, N)
         mLSE: cute.Tensor,  # (M,)
         ignore_index: Int32,  # Index to ignore in gradient computation
         shape: cute.Shape,
         tiler_mn: cute.Shape,
@@ -507,6 +503,18 @@ class CrossEntropyBackward:
         copy = partial(copy_utils.copy, pred=tXpX)
         row = tXcX[0][0]
         if row < shape[0]:
             copy(tXgX, tXsX, is_async=True)
         cute.arch.cp_async_commit_group()
@@ -516,13 +524,11 @@ class CrossEntropyBackward:
         cute.autovec_copy(tXsX, tXrX)
         x = tXrX.load().to(Float32)
-        target = Int32.zero
         dloss = Float32.zero
         lse = Float32.zero
         if row < shape[0]:
-            target = Int32(mTarget[row])
             should_ignore = Boolean(target == ignore_index)
-            # Set dloss to 0 if this index should be ignored
             if not should_ignore:
                 dloss = Float32(mDLoss[row])
             lse = Float32(mLSE[row])
@@ -534,32 +540,36 @@ class CrossEntropyBackward:
         for i in cutlass.range(cute.size(tXcFull), unroll_full=True):
             mask[i] = tXcFull[i][1] == target
         grad = cute.where(mask.load(), prob_shifted, probs)
-        grad = grad * dloss
         tXrdX.store(grad.to(tXrdX.element_type))
         if row < shape[0]:
             copy(tXrdX, tXgdX)
-@jit_cache
-def _compile_cross_entropy_backward(dtype, target_dtype, N):
-    batch_sym = cute.sym_int()
-    div = math.gcd(128 // dtype.width, N)
-    x_cute, dx_cute = [fake_tensor(dtype, (batch_sym, N), div)] * 2
-    target_cute = fake_tensor(target_dtype, (batch_sym,))
-    dloss_cute, lse_cute = [fake_tensor(Float32, (batch_sym,))] * 2
-    cross_entropy_backward_op = CrossEntropyBackward(dtype, N)
-    return cute.compile(
-        cross_entropy_backward_op,
-        x_cute,
-        target_cute,
-        dloss_cute,
-        dx_cute,
-        lse_cute,
-        Int32(0),  # ignore_index, just for compilation
-        cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
-        options="--enable-tvm-ffi",
-    )
 def _cross_entropy_backward(
@@ -568,6 +578,7 @@ def _cross_entropy_backward(
     dloss: torch.Tensor,
     lse: torch.Tensor,
     dx: torch.Tensor,
     ignore_index=-100,
 ) -> None:
     """Cross entropy backward pass.
@@ -576,8 +587,11 @@ def _cross_entropy_backward(
         target: Target class indices tensor of shape (M,)
         dloss: Upstream gradients tensor of shape (M,)
         lse: Log-sum-exp values tensor of shape (M,)
     Returns:
-        Input gradients tensor of shape (M, N)
     """
     assert x.dim() == 2, "Input must be 2D"
     assert target.dim() == 1, "Target must be 1D"
@@ -586,48 +600,33 @@ def _cross_entropy_backward(
     assert x.shape[0] == target.shape[0], "Batch dimensions must match"
     assert x.shape[0] == dloss.shape[0], "Batch dimensions must match"
     assert x.shape[0] == lse.shape[0], "Batch dimensions must match"
-    assert x.is_cuda and target.is_cuda and dloss.is_cuda and lse.is_cuda, (
-        "Tensors must be on CUDA device"
-    )
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype"
     assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64"
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
     target_dtype = torch2cute_dtype_map[target.dtype]
-    _compile_cross_entropy_backward(dtype, target_dtype, N)(
-        x, target, dloss, dx, lse, Int32(ignore_index)
     )
-@torch.library.custom_op(add_quack_op_namespace_prefix("cross_entropy_bwd_out"), mutates_args={"dx"})
 def cross_entropy_bwd_out(
     x: torch.Tensor,
     target: torch.Tensor,
     dloss: torch.Tensor,
     lse: torch.Tensor,
     dx: torch.Tensor,
     ignore_index: int = -100,
 ) -> None:
-    _cross_entropy_backward(x, target, dloss, lse, dx, ignore_index)
-@cross_entropy_bwd_out.register_fake
-def _cross_entropy_bwd_out_fake(
-    x: torch.Tensor,
-    target: torch.Tensor,
-    dloss: torch.Tensor,
-    lse: torch.Tensor,
-    dx: torch.Tensor,
-    ignore_index: int = -100,
-) -> None:
-    # See softmax.py _softmax_fwd_fake for why register_fake is needed.
-    from .cache_utils import COMPILE_ONLY
-    if COMPILE_ONLY and not isinstance(x.size(1), torch.SymInt):
-        N = x.size(1)
-        dtype = torch2cute_dtype_map[x.dtype]
-        target_dtype = torch2cute_dtype_map[target.dtype]
-        _compile_cross_entropy_backward(dtype, target_dtype, N)
 def cross_entropy_bwd(
@@ -635,51 +634,90 @@ def cross_entropy_bwd(
     target: torch.Tensor,
     dloss: torch.Tensor,
     lse: torch.Tensor,
     ignore_index: int = -100,
     inplace_backward: bool = False,
 ) -> None:
     if inplace_backward and not torch.compiler.is_compiling():
         dx = x
         _cross_entropy_backward(
-            x=x, target=target, dloss=dloss, lse=lse, dx=x, ignore_index=ignore_index
         )
     else:
         dx = torch.empty_like(x)
         cross_entropy_bwd_out(
-            x=x, target=target, dloss=dloss, lse=lse, dx=dx, ignore_index=ignore_index
         )
     return dx
 class CrossEntropyFunction(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, x, target, lse_partial=None, ignore_index=-100, inplace_backward=False):
         if lse_partial is None:
-            loss, lse = cross_entropy_fwd(x, target, ignore_index=ignore_index, return_lse=True)
         else:
             # if we already compute partial lse, then to compute the final lse we treat
             # @lse_partial as @x and @x as @target_logit
             loss, lse = cross_entropy_fwd(
-                lse_partial, target, target_logit=x, ignore_index=ignore_index, return_lse=True
             )
-        ctx.save_for_backward(x, target, lse)
         ctx.ignore_index = ignore_index
         ctx.inplace_backward = inplace_backward
         return loss
     @staticmethod
     def backward(ctx, dloss):
-        x, target, lse = ctx.saved_tensors
         dx = cross_entropy_bwd(
-            x, target, dloss, lse, ctx.ignore_index, inplace_backward=ctx.inplace_backward
         )
-        return dx, None, None, None, None
 def cross_entropy(
     x: torch.Tensor,
     target: torch.Tensor,
     lse_partial: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
     reduction: Literal["none", "mean", "sum"] = "mean",
     inplace_backward: bool = False,
@@ -690,12 +728,13 @@ def cross_entropy(
         x: Input logits tensor of shape (M, N)
         target: Target class indices tensor of shape (M,)
         lse_partial: Optional precomputed log-sum-exp partial results
         reduction: Specifies the reduction to apply to the output:
             'none': no reduction will be applied (default)
             'mean': the sum of the output will be divided by the number of elements
             'sum': the output will be summed
         inplace_backward: Whether to perform backward pass in-place
-        ignore_index: Index to ignore in loss computation (loss will be 0 for these indices)
     Returns:
         Cross entropy loss tensor:
@@ -703,8 +742,19 @@ def cross_entropy(
             - If reduction='mean': scalar tensor with mean loss
             - If reduction='sum': scalar tensor with sum of losses
     """
-    loss = CrossEntropyFunction.apply(x, target, lse_partial, ignore_index, inplace_backward)
     if reduction == "mean":
         return loss.sum() / (target != ignore_index).sum().float()
     elif reduction == "sum":
         return loss.sum()

 from typing import Optional, Type, Literal
 import torch
+from ._ops_compat import add_op_namespace_prefix
 from torch import Tensor
 import cuda.bindings.driver as cuda
 from . import copy_utils as copy_utils
 from . import layout_utils as layout_utils
 from .compile_utils import make_fake_tensor as fake_tensor
+from .dsl import cute_op
 from .reduce import row_reduce, online_softmax_reduce
 from .reduction_base import ReductionBase
+from .cache import jit_cache
 from .cute_dsl_utils import torch2cute_dtype_map
+from cutlass.base_dsl.arch import Arch
 class CrossEntropy(ReductionBase):
         mLoss: cute.Tensor,  # (M,)
         mLSE: Optional[cute.Tensor],  # (M,)
         mdX: Optional[cute.Tensor],  # (M, N) - if provided, compute gradient
+        mWeight: Optional[cute.Tensor],
         ignore_index: Int32,  # Index to ignore in loss computation
         stream: cuda.CUstream,
     ):
             mLoss,
             mLSE,
             mdX,
+            mWeight,
             ignore_index,
             tiler_mn,
             tiled_copy,
         mLoss: cute.Tensor,  # (M,)
         mLSE: Optional[cute.Tensor],  # (M,)
         mdX: Optional[cute.Tensor],  # (M, N) - if provided, compute gradient
+        mWeight: Optional[cute.Tensor],
         ignore_index: Int32,  # Index to ignore in loss computation
         tiler_mn: cute.Shape,
         tiled_copy: cute.TiledCopy,
         row = tXcX[0][0]
         target = Int32.zero
+        target_weight = Float32.zero
         if row < shape[0]:
             target = Int32(mTarget[row])
+            if const_expr(mWeight is not None):
+                # Gate on target != ignore_index: ignore_index may be negative
+                # (PyTorch default -100), and indexing mWeight at that offset is OOB.
+                if target != ignore_index:
+                    target_weight = Float32(mWeight[target])
+            else:
+                target_weight = 1.0
         if row < shape[0]:
             copy(tXgX, tXsX, is_async=True)
         ):
             lse = max_x + cute.math.log(denom, fastmath=True)
             # Set loss to 0 if this index should be ignored, otherwise compute normally
+            loss_val = target_weight * (lse - target_logit) if not should_ignore else Float32.zero
             mLoss[row] = mLoss.element_type(loss_val)
             if const_expr(mLSE is not None):
                 mLSE[row] = lse
             probs = exp_x * denom_inv
             gdX = cute.local_tile(mdX, tiler_mn, (bidx, cluster_y))
             tXgdX = thr_copy.partition_D(gdX)
             tXcFull = thr_copy.partition_S(cX)
             # Compute gradient: probs for all classes, (probs - 1) for target class
             # If ignored, gradient is already zero
             if not should_ignore:
                 for i in cutlass.range(cute.size(tXrX), unroll_full=True):
                     tXrdX_f32[i] = tXrdX_f32[i] if tXcFull[i][1] != target else tXrdX_f32[i] - 1.0
+            if const_expr(mWeight is not None):
+                tXrdX_f32.store(tXrdX_f32.load() * target_weight)
+            tXrdX = tXrdX_f32.to(tXgdX.element_type)
             if row < shape[0]:
                 copy(tXrdX, tXgdX)
+    @staticmethod
+    @jit_cache
+    def compile(
+        dtype,
+        target_dtype,
+        target_logit_dtype,
+        N,
+        has_lse,
+        has_dx,
+        weight_dtype,
+        target_logit_ndim,
+    ):
+        batch_sym = cute.sym_int()
+        div = math.gcd(128 // dtype.width, N)
+        x_cute = fake_tensor(dtype, (batch_sym, N), div)
+        dx_cute = fake_tensor(dtype, (batch_sym, N), div) if has_dx else None
+        target_cute = fake_tensor(target_dtype, (batch_sym,))
+        if target_logit_dtype is not None:
+            if target_logit_ndim == 2:
+                target_logit_cute = fake_tensor(
+                    target_logit_dtype, (batch_sym, cute.sym_int()), div
+                )
+            else:
+                target_logit_cute = fake_tensor(target_logit_dtype, (batch_sym,))
         else:
+            target_logit_cute = None
+        loss_cute = fake_tensor(Float32, (batch_sym,))
+        lse_cute = fake_tensor(Float32, (batch_sym,)) if has_lse else None
+        weight_cute = fake_tensor(weight_dtype, (N,)) if weight_dtype is not None else None
+        # If there's dx, it's faster to not use online softmax since we want the exp(x - max)
+        return cute.compile(
+            CrossEntropy(dtype, N, online_softmax=not has_dx),
+            x_cute,
+            target_cute,
+            target_logit_cute,
+            loss_cute,
+            lse_cute,
+            dx_cute,
+            weight_cute,
+            Int32(0),  # ignore_index, just for compilation
+            cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
+@cute_op(add_op_namespace_prefix("cross_entropy_fwd_out"), mutates_args={"loss", "lse", "dx"})
 def cross_entropy_fwd_out(
     x: Tensor,
     target: Tensor,
     loss: Tensor,
     lse: Optional[Tensor],
     dx: Optional[Tensor],
+    weight: Optional[Tensor],
     ignore_index: int = -100,
 ) -> None:
     """Cross entropy forward pass.
         loss: Output loss tensor of shape (M,)
         lse: Optional output log-sum-exp tensor of shape (M,)
         dx: Optional output gradient tensor of shape (M, N)
+        weight: Optional weight vector of shape (N,)
         ignore_index: Index to ignore in loss computation
     Returns:
     """
     assert x.dim() == 2, "Input must be 2D"
     assert target.dim() == 1, "Target must be 1D"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype"
     assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64"
     if target_logit is not None:
         assert target_logit.dtype in [torch.float16, torch.bfloat16, torch.float32]
+    if x.size(0) == 0:
+        return
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
     target_dtype = torch2cute_dtype_map[target.dtype]
         torch2cute_dtype_map[target_logit.dtype] if target_logit is not None else None
     )
     target_logit_ndim = target_logit.ndim if target_logit is not None else None
+    weight_dtype = torch2cute_dtype_map[weight.dtype] if weight is not None else None
+    CrossEntropy.compile(
         dtype,
         target_dtype,
         target_logit_dtype,
         N,
         lse is not None,
         dx is not None,
+        weight_dtype,
         target_logit_ndim,
+    )(x, target, target_logit, loss, lse, dx, weight, Int32(ignore_index))
 def cross_entropy_fwd(
     x: torch.Tensor,
     target: torch.Tensor,
     target_logit: Optional[torch.Tensor] = None,
+    weight: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
     return_lse: bool = False,
     return_dx: bool = False,
     loss = torch.empty(M, device=device, dtype=torch.float32)
     lse = torch.empty(M, device=device, dtype=torch.float32) if return_lse else None
     dx = (torch.empty_like(x) if not inplace_backward else x) if return_dx else None
+    cross_entropy_fwd_out(x, target, target_logit, loss, lse, dx, weight, ignore_index)
     if return_lse and return_dx:
         return loss, lse, dx
     elif return_lse:
         mDLoss: cute.Tensor,
         mdX: cute.Tensor,
         mLSE: cute.Tensor,
+        mWeight: Optional[cute.Tensor],
         ignore_index: Int32,  # Index to ignore in gradient computation
         stream: cuda.CUstream,
     ):
             mDLoss,
             mdX,
             mLSE,
+            mWeight,
             ignore_index,
             mX.shape,
             tiler_mn,
         mDLoss: cute.Tensor,  # (M,)
         mdX: cute.Tensor,  # (M, N)
         mLSE: cute.Tensor,  # (M,)
+        mWeight: Optional[cute.Tensor],
         ignore_index: Int32,  # Index to ignore in gradient computation
         shape: cute.Shape,
         tiler_mn: cute.Shape,
         copy = partial(copy_utils.copy, pred=tXpX)
         row = tXcX[0][0]
+        target = Int32.zero
+        target_weight = Float32.zero
+        if row < shape[0]:
+            target = Int32(mTarget[row])
+            if const_expr(mWeight is not None):
+                # Gate on target != ignore_index: ignore_index may be negative
+                # (PyTorch default -100), and indexing mWeight at that offset is OOB.
+                if target != ignore_index:
+                    target_weight = Float32(mWeight[target])
+            else:
+                target_weight = 1.0
         if row < shape[0]:
             copy(tXgX, tXsX, is_async=True)
         cute.arch.cp_async_commit_group()
         cute.autovec_copy(tXsX, tXrX)
         x = tXrX.load().to(Float32)
         dloss = Float32.zero
         lse = Float32.zero
         if row < shape[0]:
             should_ignore = Boolean(target == ignore_index)
+            # dloss is set to 0 if this index should be ignored
             if not should_ignore:
                 dloss = Float32(mDLoss[row])
             lse = Float32(mLSE[row])
         for i in cutlass.range(cute.size(tXcFull), unroll_full=True):
             mask[i] = tXcFull[i][1] == target
         grad = cute.where(mask.load(), prob_shifted, probs)
+        grad = grad * dloss * target_weight
         tXrdX.store(grad.to(tXrdX.element_type))
         if row < shape[0]:
             copy(tXrdX, tXgdX)
+    @staticmethod
+    @jit_cache
+    def compile(dtype, target_dtype, N, weight_dtype):
+        batch_sym = cute.sym_int()
+        div = math.gcd(128 // dtype.width, N)
+        x_cute, dx_cute = [fake_tensor(dtype, (batch_sym, N), div)] * 2
+        target_cute = fake_tensor(target_dtype, (batch_sym,))
+        dloss_cute = cute.runtime.make_fake_tensor(
+            Float32, (batch_sym,), stride=(cute.sym_int64(),)
+        )
+        lse_cute = fake_tensor(Float32, (batch_sym,))
+        weight_cute = fake_tensor(weight_dtype, (N,)) if weight_dtype is not None else None
+        return cute.compile(
+            CrossEntropyBackward(dtype, N),
+            x_cute,
+            target_cute,
+            dloss_cute,
+            dx_cute,
+            lse_cute,
+            weight_cute,
+            Int32(0),  # ignore_index, just for compilation
+            cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True),
+            options="--enable-tvm-ffi",
+        )
 def _cross_entropy_backward(
     dloss: torch.Tensor,
     lse: torch.Tensor,
     dx: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
     ignore_index=-100,
 ) -> None:
     """Cross entropy backward pass.
         target: Target class indices tensor of shape (M,)
         dloss: Upstream gradients tensor of shape (M,)
         lse: Log-sum-exp values tensor of shape (M,)
+        dx: Output gradient tensor of shape (M, N)
+        weight: Optional per-class weight tensor of shape (N,)
+        ignore_index: Index to ignore in gradient computation
     Returns:
+        None (mutates dx in-place)
     """
     assert x.dim() == 2, "Input must be 2D"
     assert target.dim() == 1, "Target must be 1D"
     assert x.shape[0] == target.shape[0], "Batch dimensions must match"
     assert x.shape[0] == dloss.shape[0], "Batch dimensions must match"
     assert x.shape[0] == lse.shape[0], "Batch dimensions must match"
     assert x.dtype in [torch.float16, torch.bfloat16, torch.float32], "Unsupported input dtype"
     assert target.dtype in [torch.int32, torch.int64], "Target must be int32 or int64"
+    if weight is not None:
+        assert weight.is_cuda, "weight must be on CUDA device"
+        assert weight.is_floating_point(), "weight must be a floating-point tensor"
+    if x.size(0) == 0:
+        return
     N = x.size(1)
     dtype = torch2cute_dtype_map[x.dtype]
     target_dtype = torch2cute_dtype_map[target.dtype]
+    weight_dtype = torch2cute_dtype_map[weight.dtype] if weight is not None else None
+    CrossEntropyBackward.compile(dtype, target_dtype, N, weight_dtype)(
+        x, target, dloss, dx, lse, weight, Int32(ignore_index)
     )
+@cute_op(add_op_namespace_prefix("cross_entropy_bwd_out"), mutates_args={"dx"})
 def cross_entropy_bwd_out(
     x: torch.Tensor,
     target: torch.Tensor,
     dloss: torch.Tensor,
     lse: torch.Tensor,
     dx: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
 ) -> None:
+    _cross_entropy_backward(x, target, dloss, lse, dx, weight, ignore_index)
 def cross_entropy_bwd(
     target: torch.Tensor,
     dloss: torch.Tensor,
     lse: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
     inplace_backward: bool = False,
 ) -> None:
     if inplace_backward and not torch.compiler.is_compiling():
         dx = x
         _cross_entropy_backward(
+            x=x,
+            target=target,
+            dloss=dloss,
+            lse=lse,
+            dx=x,
+            weight=weight,
+            ignore_index=ignore_index,
         )
     else:
         dx = torch.empty_like(x)
         cross_entropy_bwd_out(
+            x=x,
+            target=target,
+            dloss=dloss,
+            lse=lse,
+            dx=dx,
+            weight=weight,
+            ignore_index=ignore_index,
         )
     return dx
 class CrossEntropyFunction(torch.autograd.Function):
     @staticmethod
+    def forward(
+        ctx,
+        x,
+        target,
+        lse_partial=None,
+        weight=None,
+        ignore_index=-100,
+        inplace_backward=False,
+    ):
         if lse_partial is None:
+            loss, lse = cross_entropy_fwd(
+                x,
+                target,
+                weight=weight,
+                ignore_index=ignore_index,
+                return_lse=True,
+            )
         else:
             # if we already compute partial lse, then to compute the final lse we treat
             # @lse_partial as @x and @x as @target_logit
             loss, lse = cross_entropy_fwd(
+                lse_partial,
+                target,
+                target_logit=x,
+                weight=weight,
+                ignore_index=ignore_index,
+                return_lse=True,
             )
+        ctx.save_for_backward(x, target, lse, weight)
         ctx.ignore_index = ignore_index
         ctx.inplace_backward = inplace_backward
         return loss
     @staticmethod
     def backward(ctx, dloss):
+        x, target, lse, weight = ctx.saved_tensors
         dx = cross_entropy_bwd(
+            x,
+            target,
+            dloss,
+            lse,
+            weight=weight,
+            ignore_index=ctx.ignore_index,
+            inplace_backward=ctx.inplace_backward,
         )
+        return dx, None, None, None, None, None
 def cross_entropy(
     x: torch.Tensor,
     target: torch.Tensor,
     lse_partial: Optional[torch.Tensor] = None,
+    weight: Optional[torch.Tensor] = None,
     ignore_index: int = -100,
     reduction: Literal["none", "mean", "sum"] = "mean",
     inplace_backward: bool = False,
         x: Input logits tensor of shape (M, N)
         target: Target class indices tensor of shape (M,)
         lse_partial: Optional precomputed log-sum-exp partial results
+        weight: Optional per-class weight tensor of shape (N,)
+        ignore_index: Index to ignore in loss computation (loss will be 0 for these indices)
         reduction: Specifies the reduction to apply to the output:
             'none': no reduction will be applied (default)
             'mean': the sum of the output will be divided by the number of elements
             'sum': the output will be summed
         inplace_backward: Whether to perform backward pass in-place
     Returns:
         Cross entropy loss tensor:
             - If reduction='mean': scalar tensor with mean loss
             - If reduction='sum': scalar tensor with sum of losses
     """
+    loss = CrossEntropyFunction.apply(
+        x,
+        target,
+        lse_partial,
+        weight,
+        ignore_index,
+        inplace_backward,
+    )
     if reduction == "mean":
+        if weight is not None:
+            valid = target != ignore_index
+            denom = (weight[target.clamp(min=0)] * valid).sum()
+            return loss.sum() / denom
         return loss.sum() / (target != ignore_index).sum().float()
     elif reduction == "sum":
         return loss.sum()

build/torch-cuda/quack/cute_dsl_utils.py CHANGED Viewed

@@ -59,12 +59,32 @@ torch2cute_dtype_map = {
     torch.float32: Float32,
     torch.int32: Int32,
     torch.int64: Int64,
 }
 @lru_cache
-def get_max_active_clusters(cluster_size):
-    return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size)
 def _parse_arch_str(arch_str: str) -> Tuple[int, int]:
@@ -148,8 +168,46 @@ def _namedtuple_new_from_mlir_values(self, values):
     return self.__class__(*new_fields)
 def mlir_namedtuple(cls):
-    """Decorator that adds MLIR value reconstruction to a NamedTuple class.
     Usage::
@@ -158,6 +216,8 @@ def mlir_namedtuple(cls):
             tensor_arg: cute.Tensor
             const_arg: cutlass.Constexpr[int] = 0
     """
     cls.__new_from_mlir_values__ = _namedtuple_new_from_mlir_values
     return cls

     torch.float32: Float32,
     torch.int32: Int32,
     torch.int64: Int64,
+    torch.float8_e4m3fn: cutlass.Float8E4M3FN,
+    torch.float8_e5m2: cutlass.Float8E5M2,
+    torch.float8_e8m0fnu: cutlass.Float8E8M0FNU,
+    # Packed fp4: dlpack presents the logical (doubled) K extent to the DSL.
+    torch.float4_e2m1fn_x2: cutlass.Float4E2M1FN,
 }
 @lru_cache
+def get_device_multiprocessor_count(device_id: int = 0) -> int:
+    return cutlass.utils.HardwareInfo(device_id).get_device_multiprocessor_count()
+@lru_cache
+def get_max_active_clusters(
+    cluster_size: int,
+    device_capacity: Tuple[int, int] | None = None,
+    device_id: int = 0,
+) -> int:
+    if device_capacity is None:
+        device_capacity = get_device_capacity()
+    if device_capacity[0] < 9:
+        if cluster_size != 1:
+            raise ValueError("SM8x kernels do not support CTA clusters; cluster_size must be 1")
+        return get_device_multiprocessor_count(device_id)
+    return cutlass.utils.HardwareInfo(device_id).get_max_active_clusters(cluster_size=cluster_size)
 def _parse_arch_str(arch_str: str) -> Tuple[int, int]:
     return self.__class__(*new_fields)
+def _namedtuple_dynamic_fields(self):
+    """Yield fields that are represented by MLIR/runtime values.
+    Keep this in sync with ``_namedtuple_new_from_mlir_values``: fields that
+    are ``None`` or plain Python compile-time constants are preserved from the
+    compile-time template and do not consume MLIR block arguments.
+    """
+    for field_val in self:
+        if field_val is None or isinstance(field_val, StaticTypes):
+            continue
+        yield field_val
+def _namedtuple_c_pointers(self):
+    """Generic ``JitArgument.__c_pointers__`` for ``@mlir_namedtuple`` classes."""
+    from cutlass.base_dsl.typing import get_c_pointers
+    ptrs = []
+    for field_val in _namedtuple_dynamic_fields(self):
+        ptrs.extend(get_c_pointers(field_val))
+    return ptrs
+def _namedtuple_get_mlir_types(self):
+    """Generic ``JitArgument.__get_mlir_types__`` for ``@mlir_namedtuple`` classes."""
+    from cutlass.base_dsl.typing import get_mlir_types
+    types = []
+    for field_val in _namedtuple_dynamic_fields(self):
+        types.extend(get_mlir_types(field_val))
+    return types
 def mlir_namedtuple(cls):
+    """Decorator that makes a NamedTuple usable as a CuTe JIT argument.
+    Adds the full ``JitArgument`` protocol.  This matters even when a concrete
+    instance has no dynamic fields: newer CuTe DSL versions warn for non-
+    constexpr compile-only arguments that flatten to zero MLIR values unless
+    they explicitly implement the protocol.
     Usage::
             tensor_arg: cute.Tensor
             const_arg: cutlass.Constexpr[int] = 0
     """
+    cls.__c_pointers__ = _namedtuple_c_pointers
+    cls.__get_mlir_types__ = _namedtuple_get_mlir_types
     cls.__new_from_mlir_values__ = _namedtuple_new_from_mlir_values
     return cls

build/torch-cuda/quack/dsl/__init__.py ADDED Viewed

	@@ -0,0 +1,17 @@

+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+"""CuTe DSL helpers and integration hooks."""
+from . import cute_tensor_indexing  # noqa: F401
+from . import cute_tensor  # noqa: F401
+from .torch_library_op import cute_op
+__all__ = ["cute_op"]
+def __getattr__(name: str):
+    if name == "cute_op":
+        from .torch_library_op import cute_op
+        return cute_op
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

build/torch-cuda/quack/{cute_dsl_ptxas.py → dsl/cute_dsl_ptxas.py} RENAMED Viewed

@@ -3,35 +3,61 @@ System ptxas replacement for CUTLASS DSL.
 Usage::
-    CUTE_DSL_KEEP_PTX=1 CUTE_DSL_PTXAS_PATH=/usr/local/cuda/bin/ptxas pytest tests/
 Environment variables:
     CUTE_DSL_PTXAS_PATH    - Path to ptxas (e.g., /usr/local/cuda/bin/ptxas)
-    CUTE_DSL_KEEP_PTX      - Must be set to 1 before cutlass is imported
     CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output
     CUTE_DSL_DUMP_DIR      - Directory for dumped PTX files (default: cwd)
-    CUTE_DSL_KEEP_CUBIN    - Set to 1 to save compiled cubin files
 """
 import os
-import sys
 import re
-import ctypes
 import subprocess
 from pathlib import Path
-import cutlass
 CUTE_DSL_PTXAS_PATH = os.environ.get("CUTE_DSL_PTXAS_PATH", None)
 if CUTE_DSL_PTXAS_PATH:
-    os.environ["CUTE_DSL_KEEP_PTX"] = "1"
 VERBOSE = os.environ.get("CUTE_DSL_PTXAS_VERBOSE", "0") == "1"
 _original_load_cuda_library = None
 _original_create_tvm_ffi_function = None
-_user_wanted_ptx = False  # True if user originally set CUTE_DSL_KEEP_PTX=1
 def _log(msg: str):
@@ -201,18 +227,36 @@ def _patched_create_tvm_ffi_function(self):
     return _original_create_tvm_ffi_function(self)
 def patch():
-    """Install system ptxas hook. Call before importing cutlass."""
     global _original_load_cuda_library, _original_create_tvm_ffi_function, _user_wanted_ptx
     assert CUTE_DSL_PTXAS_PATH is not None
     if not os.path.isfile(CUTE_DSL_PTXAS_PATH) or not os.access(CUTE_DSL_PTXAS_PATH, os.X_OK):
         raise RuntimeError(f"ptxas not found: {CUTE_DSL_PTXAS_PATH}")
-    _user_wanted_ptx = os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1"
-    assert os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1", (
-        "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas"
-    )
     patched = False
     cuda_jit_function_cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction

 Usage::
+    CUTE_DSL_PTXAS_PATH=/usr/local/cuda/bin/ptxas pytest tests/
 Environment variables:
     CUTE_DSL_PTXAS_PATH    - Path to ptxas (e.g., /usr/local/cuda/bin/ptxas)
+    CUTE_DSL_KEEP=ptx      - Optional; keep PTX files instead of deleting them
     CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output
     CUTE_DSL_DUMP_DIR      - Directory for dumped PTX files (default: cwd)
+    CUTE_DSL_KEEP_CUBIN    - Set to 1 to save system-ptxas cubin files
 """
+import ctypes
 import os
 import re
 import subprocess
+import sys
 from pathlib import Path
 CUTE_DSL_PTXAS_PATH = os.environ.get("CUTE_DSL_PTXAS_PATH", None)
+def _keep_tokens() -> set[str]:
+    keep = os.environ.get("CUTE_DSL_KEEP", "")
+    return {token.strip().lower() for token in keep.split(",") if token.strip()}
+def _env_requests_ptx() -> bool:
+    tokens = _keep_tokens()
+    return "all" in tokens or "ptx" in tokens or os.environ.get("CUTE_DSL_KEEP_PTX") == "1"
+_USER_WANTED_PTX = _env_requests_ptx()
+def _force_keep_ptx_env() -> None:
+    tokens = _keep_tokens()
+    if "all" not in tokens and "ptx" not in tokens:
+        tokens.add("ptx")
+        os.environ["CUTE_DSL_KEEP"] = ",".join(sorted(tokens))
+    # Keep the deprecated switch too so older CUTLASS DSL builds that do not
+    # understand CUTE_DSL_KEEP still dump PTX.
+    os.environ.setdefault("CUTE_DSL_KEEP_PTX", "1")
 if CUTE_DSL_PTXAS_PATH:
+    # Must happen before CuTeDSL's EnvironmentVarManager is instantiated.
+    _force_keep_ptx_env()
+import cutlass  # noqa: E402
 VERBOSE = os.environ.get("CUTE_DSL_PTXAS_VERBOSE", "0") == "1"
 _original_load_cuda_library = None
 _original_create_tvm_ffi_function = None
+_user_wanted_ptx = False  # True if user originally requested KEEP=ptx / KEEP_PTX=1
 def _log(msg: str):
     return _original_create_tvm_ffi_function(self)
+def _force_live_keep_ptx() -> None:
+    """Update already-created CuTe DSL environment managers, if any.
+    quack imports this module before importing its kernels, but keeping this
+    fallback makes direct/late calls to ``patch()`` less fragile.
+    """
+    for cls_name in ("CuTeDSL", "CuteExperimentalDSL"):
+        dsl_cls = getattr(cutlass.cutlass_dsl, cls_name, None)
+        if dsl_cls is None:
+            continue
+        try:
+            envar = dsl_cls._get_dsl().envar
+        except Exception:
+            continue
+        envar.keep_ptx = True
+        if hasattr(envar, "keep_tokens"):
+            envar.keep_tokens = frozenset(set(envar.keep_tokens) | {"ptx"})
 def patch():
+    """Install system ptxas hook."""
     global _original_load_cuda_library, _original_create_tvm_ffi_function, _user_wanted_ptx
     assert CUTE_DSL_PTXAS_PATH is not None
     if not os.path.isfile(CUTE_DSL_PTXAS_PATH) or not os.access(CUTE_DSL_PTXAS_PATH, os.X_OK):
         raise RuntimeError(f"ptxas not found: {CUTE_DSL_PTXAS_PATH}")
+    _user_wanted_ptx = _USER_WANTED_PTX
+    _force_keep_ptx_env()
+    _force_live_keep_ptx()
     patched = False
     cuda_jit_function_cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction

build/torch-cuda/quack/dsl/cute_tensor.py ADDED Viewed

	@@ -0,0 +1,167 @@

+# Copyright (c) 2026, Tri Dao.
+"""Small CuTe tensor convenience helpers.
+Importing this module intentionally mutates CuTe's tensor class process-wide so
+fragments can be written in a more PyTorch-like style::
+    rmem_f32 = rmem_f16.to(Float32)
+    rmem_copy = rmem_view.clone()
+    rmem_contig = rmem_view.contiguous()
+``tensor.to(dtype)`` is exactly the explicit CuTe sequence::
+    dst = cute.make_rmem_tensor_like(src, dtype)
+    dst.store(src.load().to(dtype))
+``tensor.to(dtype, force_materialize=True)`` uses the same value conversion, then inserts an
+opaque no-op SSA boundary on packed f16/bf16 lanes.  This is useful when downstream codegen
+would otherwise rematerialize a vector truncation for multiple consumers.
+``tensor.clone()`` materializes into ``cute.make_rmem_tensor_like(src)``.
+``tensor.contiguous()`` mirrors ``quack.copy_utils.contiguous``.
+"""
+from __future__ import annotations
+from typing import Any
+import cutlass
+import cutlass.cute as cute
+import cutlass.cute.tensor as _cute_tensor
+from cutlass.cutlass_dsl import T, dsl_user_op
+from cutlass._mlir.dialects import llvm
+_ORIGINAL_TO_ATTR = "_quack_original_to"
+_PATCHED_TO_ATTR = "_quack_rmem_tensor_to"
+_ORIGINAL_CLONE_ATTR = "_quack_original_clone"
+_PATCHED_CLONE_ATTR = "_quack_tensor_clone"
+_ORIGINAL_CONTIGUOUS_ATTR = "_quack_original_contiguous"
+_PATCHED_CONTIGUOUS_ATTR = "_quack_tensor_contiguous"
+@dsl_user_op
+def _black_box_b32(x: cutlass.Int32, *, loc: Any = None, ip: Any = None) -> cutlass.Int32:
+    """Opaque identity for packed registers.
+    The empty asm emits no PTX instruction.  The tied input constraint (``0``) makes the
+    output use the same register class/value as the input, while still creating an SSA boundary.
+    """
+    return cutlass.Int32(
+        llvm.inline_asm(
+            T.i32(),
+            [cutlass.Int32(x).ir_value(loc=loc, ip=ip)],
+            "",
+            "=r,0",
+            has_side_effects=False,
+            is_align_stack=False,
+            asm_dialect=llvm.AsmDialect.AD_ATT,
+        )
+    )
+@cute.jit
+def _to_f16_materialized(src: Any, dtype: Any) -> Any:
+    assert src.element_type is cutlass.Float32, "src must be Float32"
+    assert dtype in (cutlass.BFloat16, cutlass.Float16), "dtype must be BFloat16 or Float16"
+    assert cute.size(src.shape) % 2 == 0, "src must have an even number of elements"
+    # Why this exists: plain Tensor.to lowers f32 -> f16/bf16 as a vector truncation.
+    # In some kernels the converted fragment has multiple consumers (e.g. STSM and RS WGMMA
+    # in SM90 FlashAttention backward).  Leaving the truncation as a high-level vector value
+    # lets later codegen rematerialize it for each consumer, producing extra packed converts
+    # and a worse WGMMA schedule.  Storing the .to result, viewing the f16/bf16 pairs as i32,
+    # then passing each packed lane through an empty tied-operand asm creates an opaque SSA
+    # boundary.  The asm emits no PTX instruction, but it forces one materialized packed value
+    # that downstream users share.
+    tmp = cute.make_rmem_tensor_like(src, dtype)
+    tmp.store(src.load().to(dtype))
+    dst = cute.make_rmem_tensor_like(src, dtype)
+    tmp_i32 = cute.recast_tensor(tmp, cutlass.Int32)
+    dst_i32 = cute.recast_tensor(dst, cutlass.Int32)
+    assert cute.size(dst_i32.shape) * 2 == cute.size(src.shape)
+    for i in cutlass.range(cute.size(dst_i32), unroll_full=True):
+        dst_i32[i] = _black_box_b32(tmp_i32[i])
+    return dst
+def _make_to() -> Any:
+    @dsl_user_op
+    def _to(
+        self: Any,
+        dtype: Any,
+        *,
+        force_materialize: bool = False,
+        loc: Any = None,
+        ip: Any = None,
+    ) -> Any:
+        if self.memspace != cute.AddressSpace.rmem:
+            raise ValueError("Tensor.to(dtype) is only supported for rmem tensors")
+        if force_materialize:
+            return _to_f16_materialized(self, dtype)
+        dst = cute.make_rmem_tensor_like(self, dtype, loc=loc, ip=ip)
+        dst.store(self.load(loc=loc, ip=ip).to(dtype, loc=loc, ip=ip), loc=loc, ip=ip)
+        return dst
+    return _to
+def _make_clone() -> Any:
+    @dsl_user_op
+    def _clone(self: Any, *, loc: Any = None, ip: Any = None) -> Any:
+        dst = cute.make_rmem_tensor_like(self, loc=loc, ip=ip)
+        cute.autovec_copy(self, dst, loc=loc, ip=ip)
+        return dst
+    return _clone
+def _make_contiguous() -> Any:
+    @dsl_user_op
+    def _contiguous(self: Any, *, loc: Any = None, ip: Any = None) -> Any:
+        dst = cute.make_rmem_tensor(self.shape, self.element_type, loc=loc, ip=ip)
+        cute.autovec_copy(self, dst, loc=loc, ip=ip)
+        return dst
+    return _contiguous
+def patch_cute_tensor() -> None:
+    """Monkey patch CuTe tensors with QuACK convenience methods.
+    The patch is idempotent. CuTe's immutable ``TensorSSA.to`` already handles
+    value conversion; this installs the analogous materializing conversion on
+    mutable register-backed ``_Tensor`` fragments, plus a ``contiguous`` method
+    equivalent to :func:`quack.copy_utils.contiguous`, and a ``clone`` method
+    that copies into a matching compact rmem tensor.
+    """
+    tensor_cls = _cute_tensor._Tensor
+    if _PATCHED_TO_ATTR not in tensor_cls.__dict__:
+        original_to = getattr(tensor_cls, "to", None)
+        if original_to is not None:
+            setattr(tensor_cls, _ORIGINAL_TO_ATTR, original_to)
+        tensor_cls.to = _make_to()  # type: ignore[method-assign]
+        setattr(tensor_cls, _PATCHED_TO_ATTR, True)
+    if _PATCHED_CLONE_ATTR not in tensor_cls.__dict__:
+        original_clone = getattr(tensor_cls, "clone", None)
+        if original_clone is not None:
+            setattr(tensor_cls, _ORIGINAL_CLONE_ATTR, original_clone)
+        tensor_cls.clone = _make_clone()  # type: ignore[method-assign]
+        setattr(tensor_cls, _PATCHED_CLONE_ATTR, True)
+    if _PATCHED_CONTIGUOUS_ATTR not in tensor_cls.__dict__:
+        original_contiguous = getattr(tensor_cls, "contiguous", None)
+        if original_contiguous is not None:
+            setattr(tensor_cls, _ORIGINAL_CONTIGUOUS_ATTR, original_contiguous)
+        tensor_cls.contiguous = _make_contiguous()  # type: ignore[method-assign]
+        setattr(tensor_cls, _PATCHED_CONTIGUOUS_ATTR, True)
+patch_cute_tensor()
+__all__ = ["patch_cute_tensor"]

build/torch-cuda/quack/dsl/cute_tensor_indexing.py ADDED Viewed

	@@ -0,0 +1,139 @@

+# Copyright (c) 2026, Tri Dao.
+"""Small compatibility layer for more Pythonic CuTe tensor indexing.
+CuTe uses ``None`` as its underscore/full-mode slice marker, e.g. ``A[i, None]``.
+This module teaches CuTe tensors the equivalent Python spelling ``:`` and expands
+``...`` to the right number of full-mode slices.
+Importing this module intentionally mutates CuTe's tensor classes process-wide.
+The original methods are retained as ``_quack_original_getitem`` and
+``_quack_original_setitem`` for debugging or manual rollback.
+"""
+from __future__ import annotations
+from typing import Any
+from cutlass.cutlass_dsl import dsl_user_op
+import cutlass.cute.tensor as _cute_tensor
+_ORIGINAL_GETITEM_ATTR = "_quack_original_getitem"
+_ORIGINAL_SETITEM_ATTR = "_quack_original_setitem"
+_PATCHED_ATTR = "_quack_extended_indexing"
+_PATCHED_SETITEM_ATTR = f"{_PATCHED_ATTR}_setitem"
+def _is_full_slice(idx: Any) -> bool:
+    return isinstance(idx, slice) and idx.start is None and idx.stop is None and idx.step is None
+def _index_uses_ellipsis(idx: Any) -> bool:
+    if idx is Ellipsis:
+        return True
+    if isinstance(idx, tuple):
+        return any(_index_uses_ellipsis(item) for item in idx)
+    return False
+def _shape_rank(shape: Any, idx: Any = None) -> int:
+    if shape is None:
+        suffix = f" in {idx!r}" if idx is not None else ""
+        raise ValueError(f"tensor shape is required to expand ellipsis{suffix}")
+    return _cute_tensor.rank(shape)
+def _shape_mode(shape: Any, mode: int) -> Any:
+    if isinstance(shape, tuple) and mode < len(shape):
+        return shape[mode]
+    return None
+def _canonicalize_cute_tensor_index(idx: Any, tensor_shape: Any = None) -> Any:
+    """Convert Python indexing sugar to CuTe's coordinate convention.
+    ``:`` becomes ``None`` (CuTe's full-mode/underscore marker) and ``...`` expands
+    within the current hierarchy level using ``tensor_shape``. Other slices like
+    ``1:4`` are intentionally rejected because CuTe tensor slicing only supports
+    keeping an entire mode or selecting a single coordinate.
+    """
+    if idx is Ellipsis:
+        return (None,) * _shape_rank(tensor_shape, idx)
+    if _is_full_slice(idx):
+        return None
+    if isinstance(idx, slice):
+        raise ValueError(f"CuTe Tensor indexing only supports full slices ':', got {idx!r}")
+    if not isinstance(idx, tuple):
+        return idx
+    ellipsis_count = sum(item is Ellipsis for item in idx)
+    if ellipsis_count > 1:
+        raise ValueError("CuTe Tensor indexing supports at most one ellipsis per tuple level")
+    explicit_modes = len(idx) - ellipsis_count
+    fill_modes = 0
+    if ellipsis_count:
+        tensor_rank = _shape_rank(tensor_shape, idx)
+        fill_modes = tensor_rank - explicit_modes
+        if fill_modes < 0:
+            raise ValueError(
+                f"ellipsis cannot expand index {idx!r} for rank-{tensor_rank} CuTe Tensor mode"
+            )
+    result: list[Any] = []
+    mode = 0
+    for item in idx:
+        if item is Ellipsis:
+            result.extend([None] * fill_modes)
+            mode += fill_modes
+        else:
+            result.append(_canonicalize_cute_tensor_index(item, _shape_mode(tensor_shape, mode)))
+            mode += 1
+    return tuple(result)
+def _make_getitem(original_getitem: Any) -> Any:
+    @dsl_user_op
+    def _getitem(self: Any, idx: Any, *, loc: Any = None, ip: Any = None) -> Any:
+        tensor_shape = self.shape if _index_uses_ellipsis(idx) else None
+        idx = _canonicalize_cute_tensor_index(idx, tensor_shape)
+        return original_getitem(self, idx, loc=loc, ip=ip)
+    return _getitem
+def _make_setitem(original_setitem: Any) -> Any:
+    @dsl_user_op
+    def _setitem(self: Any, idx: Any, data: Any, *, loc: Any = None, ip: Any = None) -> Any:
+        tensor_shape = self.shape if _index_uses_ellipsis(idx) else None
+        idx = _canonicalize_cute_tensor_index(idx, tensor_shape)
+        return original_setitem(self, idx, data, loc=loc, ip=ip)
+    return _setitem
+def patch_cute_tensor_indexing() -> None:
+    """Monkey patch CuTe Tensor indexing with ``:``, ``...`` sugar.
+    The patch is idempotent and keeps the original CuTe implementation for all
+    canonical coordinates, so existing ``A[i, j, None]`` code continues to behave
+    exactly as before. It is a process-wide mutation of CuTe's tensor classes.
+    """
+    for cls in (_cute_tensor._Tensor, _cute_tensor.TensorSSA):
+        if _PATCHED_ATTR not in cls.__dict__:
+            setattr(cls, _ORIGINAL_GETITEM_ATTR, cls.__getitem__)
+            cls.__getitem__ = _make_getitem(cls.__getitem__)  # type: ignore[method-assign]
+            setattr(cls, _PATCHED_ATTR, True)
+    # TensorSSA has no upstream __setitem__, so only _Tensor needs the store path patched.
+    tensor_cls = _cute_tensor._Tensor
+    if _PATCHED_SETITEM_ATTR not in tensor_cls.__dict__:
+        setattr(tensor_cls, _ORIGINAL_SETITEM_ATTR, tensor_cls.__setitem__)
+        tensor_cls.__setitem__ = _make_setitem(tensor_cls.__setitem__)  # type: ignore[method-assign]
+        setattr(tensor_cls, _PATCHED_SETITEM_ATTR, True)
+patch_cute_tensor_indexing()
+__all__ = ["patch_cute_tensor_indexing"]

build/torch-cuda/quack/dsl/smem_struct.py ADDED Viewed

	@@ -0,0 +1,99 @@

+# Copyright (c) 2026, Tri Dao.
+# SPDX-License-Identifier: BSD-3-Clause
+"""Per-field smem-partition annotations for SharedStorage declarations.
+A `cute.struct` lowers to ONE smem_alloca op carrying ONE `smem.partition_id`
+attribute, so a single struct cannot mix RESERVED and USER fields at the IR
+level. `Reserved[...]` + `@partitioned_struct` provide the single-declaration
+sugar instead: at trace time the decorator splits the annotated class into two
+plain cute.structs — fields wrapped in `Reserved[...]` go to a struct
+allocated with `partition=SmemPartition.RESERVED` (low addresses, packing with
+the pipeline mbarriers and the TMEM holding buf under the 1KB that is
+otherwise alignment pad ahead of the 1024-aligned USER buffers), everything
+else stays USER — and `.allocate(smem)` returns one namespace exposing all
+fields uniformly.
+    @partitioned_struct
+    class SharedStorage:
+        sdt: Reserved[spec.smem_struct(128)]   # RESERVED partition
+        sX: X.smem_struct(1024)                # USER partition
+        ...
+    storage = SharedStorage.allocate(smem)
+    storage.sdt, storage.sX                    # fields, regardless of partition
+This is trace-time-only machinery (no monkey-patching of the DSL).
+"""
+from types import SimpleNamespace
+import cutlass.cute as cute
+from cutlass.utils import SmemPartition
+class Reserved:
+    """Annotation marker: allocate this field in the RESERVED smem partition."""
+    def __init__(self, inner):
+        self.inner = inner
+    def __class_getitem__(cls, inner):
+        return cls(inner)
+class PartitionedStruct:
+    """A SharedStorage declaration split by partition. Not a cute.struct itself:
+    holds one cute.struct per partition and allocates/merges them."""
+    def __init__(self, cls):
+        annotations = dict(cls.__annotations__)
+        reserved_ann = {k: v.inner for k, v in annotations.items() if isinstance(v, Reserved)}
+        user_ann = {k: v for k, v in annotations.items() if not isinstance(v, Reserved)}
+        self._user_struct = (
+            cute.struct(type(cls.__name__, (), {"__annotations__": user_ann})) if user_ann else None
+        )
+        self._reserved_struct = (
+            cute.struct(type(cls.__name__ + "Reserved", (), {"__annotations__": reserved_ann}))
+            if reserved_ann
+            else None
+        )
+        self._user_fields = list(user_ann)
+        self._reserved_fields = list(reserved_ann)
+    def size_in_bytes(self) -> int:
+        """USER-partition footprint (what counts against smem_capacity - 1KB)."""
+        return self._user_struct.size_in_bytes() if self._user_struct is not None else 0
+    def reserved_size_in_bytes(self) -> int:
+        """RESERVED-partition footprint of the declared fields (the pipeline
+        mbarriers / TMEM holding buf allocate there separately)."""
+        return self._reserved_struct.size_in_bytes() if self._reserved_struct is not None else 0
+    def allocate(self, smem) -> SimpleNamespace:
+        """Allocate both partitions (RESERVED first, at the partition base) and
+        return a namespace exposing every declared field. A partition whose
+        struct is empty for this config (every field zero-sized) is skipped —
+        smem_alloca rejects 0-byte layouts — and its fields come back as None;
+        callers only touch such fields under the same has_* guards that made
+        them zero-sized."""
+        fields = {}
+        if self._reserved_struct is not None:
+            if self._reserved_struct.size_in_bytes() > 0:
+                inst = smem.allocate(self._reserved_struct, partition=SmemPartition.RESERVED)
+                for name in self._reserved_fields:
+                    fields[name] = getattr(inst, name)
+            else:
+                fields.update(dict.fromkeys(self._reserved_fields))
+        if self._user_struct is not None:
+            if self._user_struct.size_in_bytes() > 0:
+                inst = smem.allocate(self._user_struct)
+                for name in self._user_fields:
+                    fields[name] = getattr(inst, name)
+            else:
+                fields.update(dict.fromkeys(self._user_fields))
+        return SimpleNamespace(**fields)
+def partitioned_struct(cls) -> PartitionedStruct:
+    return PartitionedStruct(cls)

build/torch-cuda/quack/dsl/torch_library_op.py ADDED Viewed

	@@ -0,0 +1,65 @@

+# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao.
+"""``cute_op``: ``torch.library.custom_op`` for CuTe DSL kernels.
+Same trick as ``torch.library.triton_op`` (register the impl as the fake/meta
+kernel too), specialized for our setup: the fake is a pure no-op. Our ops
+only mutate their inputs, so Dynamo / AOT autograd need no shape effect from
+the fake, and kernel compilation is owned entirely by ``jit_cache`` (plus
+the async compile pool) at real execution time.
+This removes the need for hand-written ``_*_fake`` twins on each op.
+Note: we deliberately do NOT gate on ``torch.compiler.is_compiling()`` —
+that flag's underlying ``_is_compiling_flag`` is only set during
+``torch.export``, never during ``torch.compile``. Dynamo's
+``_get_fake_value_impl`` would otherwise run the body and surface
+any ``_compile_*`` ``ValueError`` as a ``TorchRuntimeError`` graph break.
+"""
+from __future__ import annotations
+from typing import Any, Callable, Iterable, Optional, Union
+import torch
+__all__ = ["cute_op"]
+def cute_op(
+    name: str,
+    *,
+    mutates_args: Union[str, Iterable[str]],
+    schema: Optional[str] = None,
+    device_types: Optional[Union[str, Iterable[str]]] = None,
+) -> Callable:
+    """Like ``torch.library.triton_op``, but for CuTe DSL kernels.
+    Args:
+        name: ``"namespace::op_name"``.
+        mutates_args: Names of mutated tensor args.
+        schema: Optional explicit schema. Required when mutating an
+            ``Optional[Tensor]`` arg (PyTorch can't infer those).
+        device_types: Optional device-type restriction.
+    """
+    def dec(fn: Callable) -> Any:
+        kwargs: dict[str, Any] = {"mutates_args": mutates_args}
+        if schema is not None:
+            kwargs["schema"] = schema
+        if device_types is not None:
+            kwargs["device_types"] = device_types
+        op = torch.library.custom_op(name, fn, **kwargs)
+        @op.register_fake
+        def _fake(*args, **kw):
+            # Pure no-op: our ops only mutate their input tensors, so under
+            # torch.compile / AOT autograd tracing there is no fake output to
+            # produce, and running the body would pay compile latency at
+            # dynamo trace time (or crash for shape/dtype combos the kernel
+            # intentionally rejects). Kernel compilation is handled by
+            # jit_cache + the async compile pool at real execution time.
+            return
+        return op
+    return dec

build/torch-cuda/quack/epi_composable.py CHANGED Viewed

@@ -1,35 +1,45 @@
 # Copyright (c) 2025, Tri Dao.
 """ComposableEpiMixin: composes EpiOps into epilogue hook methods.
-Subclasses declare _epi_ops as a tuple of EpiOp instances. The mixin auto-generates
-epi_smem_bytes_per_stage, epi_get_smem_struct, epi_get_smem_tensors, epi_begin,
-epi_begin_loop, epi_end, and EpilogueParams by querying each op.
-epi_begin and epi_begin_loop return dicts keyed by op name, so epi_visit_subtile
-can access values by name (e.g. epi_loop_tensors["alpha"]).
-EpilogueParams is auto-generated from _epi_ops (via param_fields()) plus any
-_extra_param_fields declared on the subclass. Subclasses still define
-EpilogueArguments and epi_to_underlying_arguments manually.
 """
 from dataclasses import make_dataclass, MISSING
 import cutlass.cute as cute
-from cutlass import const_expr
-from .epi_ops import EpiContext, Scalar
-def _compute_smem_map(ops):
-    """Pre-compute name → smem tensor index for each non-Scalar op."""
-    smem_map = {}
-    idx = 0
-    for op in ops:
-        if not isinstance(op, Scalar):
-            smem_map[op.name] = idx
-            idx += 1
-    return smem_map
 def _make_epi_params(epi_ops, extra_fields, bases):
@@ -52,15 +62,11 @@ class ComposableEpiMixin:
     _epi_ops = ()
     _extra_param_fields = ()  # [(name, type, default), ...] for non-op params (e.g. act_fn)
-    _epi_param_bases = ()  # Base classes for EpilogueParams (e.g. (ParamsBase,))
-    _epi_smem_map = {}
-    _epi_has_async_ops = False
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         if cls._epi_ops:
-            cls._epi_smem_map = _compute_smem_map(cls._epi_ops)
-            cls._epi_has_async_ops = any(op.needs_async_fence() for op in cls._epi_ops)
             # Auto-generate EpilogueParams if not explicitly defined on this class
             if "EpilogueParams" not in cls.__dict__:
                 cls.EpilogueParams = _make_epi_params(
@@ -69,39 +75,89 @@ class ComposableEpiMixin:
     # --- Host-side: args → params ---
     def _epi_ops_to_params_dict(self, args):
-        """Merge each op's to_params into a single dict. Subclasses call this,
-        add custom fields, then construct self.EpilogueParams(**d)."""
         d = {}
         for op in self._epi_ops:
             d.update(op.to_params(self, args))
         return d
     # --- Host-side: smem allocation (queried from ops) ---
     @classmethod
-    def epi_smem_bytes_per_stage(cls, args, cta_tile_shape_mnk, epi_tile):
-        return sum(
-            op.smem_bytes(getattr(args, op.name, None), cta_tile_shape_mnk, epi_tile)
-            for op in cls._epi_ops
-        )
     def epi_get_smem_struct(self, params):
-        fields = {}
         for op in self._epi_ops:
             result = op.smem_struct_field(self, params)
             if result is not None:
-                name, ftype = result
-                fields[name] = ftype
-        EpiSharedStorage = type("EpiSharedStorage", (), {"__annotations__": fields})
         return cute.struct(EpiSharedStorage)
     def epi_get_smem_tensors(self, params, storage):
-        return tuple(
-            op.get_smem_tensor(self, params, storage.epi)
             for op in self._epi_ops
             if not isinstance(op, Scalar)
-        )
     def epi_get_tma_atoms(self, params, *, loc=None, ip=None):
         atoms = []
@@ -123,6 +179,7 @@ class ComposableEpiMixin:
         varlen_manager,
         epilogue_barrier,
         tidx,
     ):
         ctx = EpiContext(
             self,
@@ -133,27 +190,24 @@ class ComposableEpiMixin:
             varlen_manager,
             epilogue_barrier,
             tidx,
         )
-        smem_map = self._epi_smem_map
         results = {
             op.name: op.begin(
                 self,
-                getattr(params, op.name, None),
-                epi_smem_tensors[smem_map[op.name]] if op.name in smem_map else None,
                 ctx,
             )
             for op in self._epi_ops
         }
-        if const_expr(self._epi_has_async_ops):
-            has_async_data = any(
-                getattr(params, op.name, None) is not None
-                for op in self._epi_ops
-                if op.needs_async_fence()
-            )
-            if const_expr(has_async_data):
-                cute.arch.cp_async_commit_group()
-                cute.arch.cp_async_wait_group(0)
-                epilogue_barrier.arrive_and_wait()
         return results
     def epi_begin_loop(self, params, epi_tensors, epi_coord):
@@ -161,6 +215,59 @@ class ComposableEpiMixin:
             op.name: op.begin_loop(self, epi_tensors[op.name], epi_coord) for op in self._epi_ops
         }
     @cute.jit
     def epi_end(
         self,
@@ -176,7 +283,7 @@ class ComposableEpiMixin:
         for op in self._epi_ops:
             op.end(
                 self,
-                getattr(params, op.name, None),
                 epi_tensors[op.name],
                 epi_tile,
                 tiled_copy_t2r,

 # Copyright (c) 2025, Tri Dao.
 """ComposableEpiMixin: composes EpiOps into epilogue hook methods.
+Subclasses declare _epi_ops as a class-level tuple of EpiOp instances — the
+static *schema* for the epilogue. The mixin auto-generates epi_smem_bytes,
+epi_get_smem_struct, epi_get_smem_tensors, epi_begin, epi_begin_loop,
+epi_end_loop, epi_end, and EpilogueParams by querying each op.
+Host-side, `_epi_ops_to_params_dict` (called from each subclass's
+`epi_to_underlying_arguments`) filters `_epi_ops` automatically: it shadows
+the class-level tuple with an instance-level tuple containing only the ops
+whose argument tensor is non-None. All later iteration (host- and
+device-side) walks the filtered tuple, so each EpiOp's hook methods can
+assume their `param`/`arg_tensor` is non-None.
+The two host-side hooks that run *before* `epi_to_underlying_arguments`
+(`resolve_epi_m_major` and the classmethod `epi_smem_bytes`) filter inline
+from `args`, preserving the same non-None invariant for `op.epi_m_major_score`
+and `op.smem_bytes`. They have to run first because `epi_to_underlying_arguments`
+itself depends on static attributes set up before it (e.g. `gemm.epi_tile`,
+`gemm.epi_c_stage`), and those attributes are themselves derived from
+`epi_m_major` and the epi smem budget — a chicken-and-egg ordering we resolve by
+letting these two hooks see the raw `args` and filter inline.
+epi_get_smem_tensors, epi_begin, and epi_begin_loop all return dicts keyed by
+op name, so consumers access values by name (e.g. epi_smem_tensors["mAuxOut"],
+epi_loop_tensors["alpha"]). Because inactive ops are filtered out, consumers
+must use `.get(name)` (returns None for inactive ops) rather than `[name]`.
+EpilogueParams is auto-generated from the full class-level _epi_ops (via
+param_fields()) plus any _extra_param_fields declared on the subclass.
+Subclasses still define EpilogueArguments and epi_to_underlying_arguments
+manually.
 """
 from dataclasses import make_dataclass, MISSING
 import cutlass.cute as cute
+from cutlass import Int32, const_expr
+from .cute_dsl_utils import ParamsBase
+from .epi_ops import EpiContext, EpiSmemBytes, Scalar
 def _make_epi_params(epi_ops, extra_fields, bases):
     _epi_ops = ()
     _extra_param_fields = ()  # [(name, type, default), ...] for non-op params (e.g. act_fn)
+    _epi_param_bases = (ParamsBase,)  # Base classes for the auto-generated EpilogueParams
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         if cls._epi_ops:
             # Auto-generate EpilogueParams if not explicitly defined on this class
             if "EpilogueParams" not in cls.__dict__:
                 cls.EpilogueParams = _make_epi_params(
     # --- Host-side: args → params ---
+    def _filter_epi_ops(self, args):
+        """Shadow `_epi_ops` with an instance-level tuple of only the ops whose
+        arg is non-None. Called automatically by `_epi_ops_to_params_dict`, so
+        subclass `epi_to_underlying_arguments` methods don't need to invoke it
+        directly. After this runs, op hook methods can assume their
+        `param`/`arg_tensor` is non-None."""
+        self._epi_ops = tuple(
+            op for op in type(self)._epi_ops if getattr(args, op.name, None) is not None
+        )
     def _epi_ops_to_params_dict(self, args):
+        """Filter `_epi_ops` to active ops, then merge each op's to_params into
+        a single dict. Subclasses call this from epi_to_underlying_arguments,
+        add custom fields, then construct self.EpilogueParams(**d). Filtering
+        here means every later iteration of self._epi_ops (host- and
+        device-side) walks only active ops, and each op hook can assume its
+        arg is non-None."""
+        self._filter_epi_ops(args)
         d = {}
         for op in self._epi_ops:
             d.update(op.to_params(self, args))
         return d
+    def resolve_epi_m_major(self, args):
+        # Runs inside _setup_attributes, before epi_to_underlying_arguments,
+        # because epi_m_major drives epi_tile / smem layout choices that
+        # epi_to_underlying_arguments later consumes. self._epi_ops is still
+        # the class-level schema at this point, so we filter inline from args
+        # to keep op.epi_m_major_score's non-None invariant.
+        score = 0
+        for op in type(self)._epi_ops:
+            arg = getattr(args, op.name, None)
+            if arg is not None:
+                score += op.epi_m_major_score(arg, self)
+        return score >= 0
     # --- Host-side: smem allocation (queried from ops) ---
     @classmethod
+    def epi_smem_bytes(cls, args, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None):
+        # Runs inside _compute_stages, before epi_to_underlying_arguments,
+        # because the AB/epi stage counts (and therefore epi_c_stage) depend
+        # on the epi smem budget that this returns; epi_to_underlying_arguments
+        # then consumes epi_c_stage to build TileLoad's staged smem layout.
+        # Stays a classmethod because _compute_stages is a classmethod and may
+        # be invoked without an instance, so we filter inline from args.
+        result = EpiSmemBytes()
+        for op in cls._epi_ops:
+            arg = getattr(args, op.name, None)
+            if arg is not None:
+                result += op.smem_bytes(arg, cta_tile_shape_mnk, epi_tile, warp_shape_mnk)
+        return result
     def epi_get_smem_struct(self, params):
+        fields = []
         for op in self._epi_ops:
             result = op.smem_struct_field(self, params)
             if result is not None:
+                fields.append(result)
+        # cute.struct rejects empty annotations. When every active op contributes
+        # no smem (e.g. only Scalar ops, or no active ops at all), return a
+        # zero-byte placeholder matching gemm_base's default epi struct.
+        if not fields:
+            return cute.struct.MemRange[Int32, 0]
+        # Sort smallest-to-largest so smaller fields pack ahead of larger
+        # higher-aligned fields, reducing smem wasted to alignment padding.
+        def _field_bytes(name_ftype):
+            wrapper = type("_F", (), {"__annotations__": {name_ftype[0]: name_ftype[1]}})
+            return cute.struct(wrapper).size_in_bytes()
+        fields.sort(key=_field_bytes)
+        annotations = {name: ftype for name, ftype in fields}
+        EpiSharedStorage = type("EpiSharedStorage", (), {"__annotations__": annotations})
         return cute.struct(EpiSharedStorage)
     def epi_get_smem_tensors(self, params, storage):
+        return {
+            op.name: op.get_smem_tensor(self, params, storage.epi)
             for op in self._epi_ops
             if not isinstance(op, Scalar)
+        }
     def epi_get_tma_atoms(self, params, *, loc=None, ip=None):
         atoms = []
         varlen_manager,
         epilogue_barrier,
         tidx,
+        tRS_rD_layout=None,
     ):
         ctx = EpiContext(
             self,
             varlen_manager,
             epilogue_barrier,
             tidx,
+            tRS_rD_layout,
         )
         results = {
             op.name: op.begin(
                 self,
+                getattr(params, op.name),
+                epi_smem_tensors.get(op.name),
                 ctx,
             )
             for op in self._epi_ops
         }
+        # self._epi_ops is filtered to active ops, so any op needing a fence
+        # has a non-None tensor; no inner None check required.
+        has_async_data = any(op.needs_async_fence() for op in self._epi_ops)
+        if const_expr(has_async_data):
+            cute.arch.cp_async_commit_group()
+            cute.arch.cp_async_wait_group(0)
+            epilogue_barrier.arrive_and_wait()
         return results
     def epi_begin_loop(self, params, epi_tensors, epi_coord):
             op.name: op.begin_loop(self, epi_tensors[op.name], epi_coord) for op in self._epi_ops
         }
+    def epi_tile_load_g2s_copy_fns(
+        self,
+        params,
+        epi_smem_tensors,
+        tile_coord_mnkl,
+        varlen_manager,
+        epi_pipeline,
+    ):
+        return tuple(
+            op.load_g2s_copy_fn(
+                self,
+                params,
+                epi_smem_tensors.get(op.name),
+                tile_coord_mnkl,
+                varlen_manager,
+                epi_pipeline,
+            )
+            for op in self._epi_ops
+            if op.is_tile_load()
+        )
+    @cute.jit
+    def epi_tile_load_s2r(self, params, epi_tensors, stage_idx):
+        for op in self._epi_ops:
+            op.load_s2r(self, getattr(params, op.name), epi_tensors[op.name], stage_idx)
+    @cute.jit
+    def epi_end_loop(
+        self,
+        params,
+        epi_tensors,
+        epi_coord,
+        epi_tile,
+        tiled_copy_t2r,
+        tiled_copy_r2s,
+        tile_coord_mnkl,
+        varlen_manager,
+        tidx,
+    ):
+        for op in self._epi_ops:
+            op.end_loop(
+                self,
+                getattr(params, op.name),
+                epi_tensors[op.name],
+                epi_coord,
+                epi_tile,
+                tiled_copy_t2r,
+                tiled_copy_r2s,
+                tile_coord_mnkl,
+                varlen_manager,
+                tidx,
+            )
     @cute.jit
     def epi_end(
         self,
         for op in self._epi_ops:
             op.end(
                 self,
+                getattr(params, op.name),
                 epi_tensors[op.name],
                 epi_tile,
                 tiled_copy_t2r,

build/torch-cuda/quack/epi_ops.py CHANGED Viewed

@@ -5,14 +5,18 @@ Each EpiOp encapsulates a single tensor kind's behavior across the epilogue life
 smem allocation, begin (one-time per-tile setup), begin_loop (per-subtile extraction),
 end (cleanup).
-The ops are composed via ComposableEpiMixin which iterates over a static _epi_ops tuple
-to generate epi_smem_bytes_per_stage, epi_get_smem_struct, epi_get_smem_tensors,
-epi_begin, and epi_begin_loop automatically.
 """
 import math
 import operator
 from functools import partial
 import cutlass
 import cutlass.cute as cute
@@ -26,7 +30,12 @@ from . import layout_utils as layout_utils
 class EpiContext:
-    """Shared context passed to EpiOp.begin methods. Bundles common arguments."""
     __slots__ = (
         "epi_tile",
@@ -36,6 +45,7 @@ class EpiContext:
         "varlen_manager",
         "epilogue_barrier",
         "tidx",
         "partition_for_epilogue_fn",
         "num_epi_threads",
         "batch_idx",
@@ -53,6 +63,7 @@ class EpiContext:
         varlen_manager,
         epilogue_barrier,
         tidx,
     ):
         self.epi_tile = epi_tile
         self.tiled_copy_t2r = tiled_copy_t2r
@@ -61,6 +72,7 @@ class EpiContext:
         self.varlen_manager = varlen_manager
         self.epilogue_barrier = epilogue_barrier
         self.tidx = tidx
         self.tile_M = gemm.cta_tile_shape_mnk[0]
         self.tile_N = gemm.cta_tile_shape_mnk[1]
         self.batch_idx = tile_coord_mnkl[3]
@@ -120,6 +132,31 @@ def _get_lane_warp_layouts(tiled_copy, reference_src=True):
     return lane_layout_MN, warp_layout_MN
 class EpiOp:
     """Base class for composable epilogue operations."""
@@ -137,11 +174,15 @@ class EpiOp:
         Returns dict of {param_name: value}. Like EVT's to_underlying_arguments."""
         return {}
-    # --- Host-side: smem allocation ---
-    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile):
-        """Bytes of smem needed per stage. arg_tensor is the EpilogueArguments field."""
         return 0
     def smem_struct_field(self, gemm, params):
         """Return (field_name, field_type) for @cute.struct, or None if no smem needed.
         params is the full EpilogueParams object."""
@@ -156,6 +197,22 @@ class EpiOp:
         """Return list of TMA atoms for this op."""
         return []
     # --- Device-side: kernel execution ---
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
@@ -166,6 +223,27 @@ class EpiOp:
         """Per-subtile extraction. Returns value for epi_visit_subtile."""
         return state
     def needs_async_fence(self):
         """Whether this op issues async copies that need a fence."""
         return False
@@ -201,14 +279,9 @@ class Scalar(EpiOp):
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
-        result = None
-        if const_expr(param is not None):
-            result = (
-                utils.load_scalar_or_pointer(param, dtype=self.dtype)
-                if const_expr(self.dtype is not None)
-                else utils.load_scalar_or_pointer(param)
-            )
-        return result
 class VecLoad(EpiOp):
@@ -239,23 +312,20 @@ class VecLoad(EpiOp):
     def _coord_idx(self):
         return 1 if self.dim == 1 else 0
-    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile):
-        if arg_tensor is None:
-            return 0
-        return self._tile_size(cta_tile_shape_mnk) * (arg_tensor.element_type.width // 8)
     def smem_struct_field(self, gemm, params):
-        tensor = getattr(params, self.name, None)
-        if tensor is None:
-            size, dtype = 0, Float32
-        else:
-            size = self._tile_size(gemm.cta_tile_shape_mnk)
-            dtype = tensor.element_type
-        return (f"s_{self.name}", cute.struct.Align[cute.struct.MemRange[dtype, size], 16])
     def get_smem_tensor(self, gemm, params, storage_epi):
-        if getattr(params, self.name, None) is None:
-            return None
         return getattr(storage_epi, f"s_{self.name}").get_tensor(
             cute.make_layout(self._tile_size(gemm.cta_tile_shape_mnk))
         )
@@ -263,49 +333,61 @@ class VecLoad(EpiOp):
     def needs_async_fence(self):
         return True
     def _get_gmem_vec(self, param, ctx):
         """Get the global memory vector for this tile. Override for varlen."""
         return param[ctx.batch_idx, None]
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
-        tDsV = None
-        if const_expr(param is not None):
-            dtype = param.element_type
-            num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width
-            thr_copy = copy_utils.tiled_copy_1d(
-                dtype, ctx.num_epi_threads, num_copy_elems, is_async=True
-            ).get_slice(ctx.tidx)
-            mVec = self._get_gmem_vec(param, ctx)
-            tile_dim = self._tile_dim(ctx)
-            coord_idx = ctx.tile_coord_mnkl[self._coord_idx()]
-            gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,))
-            tVgV = thr_copy.partition_S(gVec)
-            tVsV = thr_copy.partition_D(smem_tensor)
-            tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim))
-            limit = min(cute.size(mVec, mode=[0]) - coord_idx * tile_dim, tile_dim)
-            pred = cute.make_rmem_tensor((1, cute.size(tVsV.shape[1])), Boolean)
-            for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True):
-                pred[0, m] = tVcV[0, m] < limit
-            cute.copy(thr_copy, tVgV, tVsV, pred=pred)
-            tDsV = ctx.partition_for_epilogue_fn(
-                cute.make_tensor(
-                    smem_tensor.iterator,
-                    cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()),
-                )
             )
-            if const_expr(ctx.tiled_copy_t2r is not None):
-                tDsV = ctx.tiled_copy_r2s.retile(tDsV)
-        return tDsV
     @cute.jit
     def begin_loop(self, gemm, state, epi_coord):
-        tDrV_cvt = None
-        if const_expr(state is not None):
-            tDsV_cur = cute.group_modes(state, 3, cute.rank(state))[None, None, None, epi_coord]
             tDrV = cute.make_rmem_tensor(tDsV_cur.layout, tDsV_cur.element_type)
             cute.autovec_copy(cute.filter_zeros(tDsV_cur), cute.filter_zeros(tDrV))
-            tDrV_cvt = cute.make_rmem_tensor_like(tDrV, gemm.acc_dtype)
             tDrV_cvt.store(tDrV.load().to(gemm.acc_dtype))
         return tDrV_cvt
@@ -338,63 +420,47 @@ class ColVecLoad(VecLoad):
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
-        tDsV = None
-        tDrV_cvt = None
-        if const_expr(param is not None):
-            dtype = param.element_type
-            num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width
-            thr_copy = copy_utils.tiled_copy_1d(
-                dtype, ctx.num_epi_threads, num_copy_elems, is_async=True
-            ).get_slice(ctx.tidx)
-            mVec = self._get_gmem_vec(param, ctx)
-            tile_dim = self._tile_dim(ctx)
-            coord_idx = ctx.tile_coord_mnkl[self._coord_idx()]
-            gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,))
-            tVgV = thr_copy.partition_S(gVec)
-            tVsV = thr_copy.partition_D(smem_tensor)
-            tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim))
-            # ColVec uses varlen-aware limit
-            limit = min(
-                ctx.varlen_manager.len_m(ctx.batch_idx) - coord_idx * tile_dim,
-                tile_dim,
-            )
-            pred = cute.make_rmem_tensor((1, cute.size(tVsV.shape[1])), Boolean)
-            for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True):
-                pred[0, m] = tVcV[0, m] < limit
-            cute.copy(thr_copy, tVgV, tVsV, pred=pred)
-            tDsV = ctx.partition_for_epilogue_fn(
-                cute.make_tensor(
-                    smem_tensor.iterator,
-                    cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()),
-                )
             )
-            if const_expr(ctx.tiled_copy_t2r is not None):
-                tDsV = ctx.tiled_copy_r2s.retile(tDsV)
-            # Pre-allocate register tensor reused across begin_loop calls
-            tDsV_sub = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, 0]
-            tDrV_cvt = cute.make_rmem_tensor(tDsV_sub.layout, gemm.acc_dtype)
         return [tDsV, tDrV_cvt]
-    @cute.jit
-    def begin_loop(self, gemm, state, epi_coord):
-        tDsV, tDrV_cvt = state[0], state[1]
-        if const_expr(tDsV is not None):
-            # Col vector is constant across N subtiles — only copy on first N subtile.
-            # Assumes N-major epi subtile order: epi_tile_layout = ordered_layout(..., order=(1,0))
-            epi_n = epi_coord[1]
-            if epi_n == 0:
-                tDsV_cur = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, epi_coord]
-                tDrV = cute.make_rmem_tensor(tDsV_cur.layout, tDsV_cur.element_type)
-                cute.autovec_copy(cute.filter_zeros(tDsV_cur), cute.filter_zeros(tDrV))
-                tDrV_cvt.store(tDrV.load().to(gemm.acc_dtype))
-        return tDrV_cvt
 class TileStore(EpiOp):
     """Tile-sized output tensor stored via TMA (e.g. postact).
     Args:
-        name: field name in EpilogueArguments/Params (e.g. "mPostAct")
         epi_tile_fn: optional (gemm, epi_tile) -> epi_tile for half-tile (GemmGated)
     """
@@ -412,13 +478,13 @@ class TileStore(EpiOp):
         return f"epi_tile_{self.name}"
     def param_fields(self):
-        from dataclasses import MISSING
         return [
-            (self._tma_atom_key(), object, MISSING),
-            (self.name, object, MISSING),
-            (self._smem_layout_key(), object, MISSING),
-            (self._epi_tile_key(), object, MISSING),
         ]
     def to_params(self, gemm, args):
@@ -434,50 +500,207 @@ class TileStore(EpiOp):
             self._epi_tile_key(): epi_tile_out,
         }
-    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile):
-        if arg_tensor is None:
-            return 0
         if self.epi_tile_fn is not None:
             epi_tile = self.epi_tile_fn(None, epi_tile)
-        return cute.size(cute.shape(epi_tile)) * (arg_tensor.element_type.width // 8)
     def smem_struct_field(self, gemm, params):
-        smem_layout_key = self._smem_layout_key()
-        if not hasattr(params, smem_layout_key):
-            return (f"s_{self.name}", cute.struct.MemRange[Float32, 0])
         return (
             f"s_{self.name}",
             cute.struct.Align[
                 cute.struct.MemRange[
-                    gemm.postact_dtype,
-                    cute.cosize(getattr(params, smem_layout_key)),
                 ],
                 gemm.buffer_align_bytes,
             ],
         )
     def get_smem_tensor(self, gemm, params, storage_epi):
-        smem_layout_key = self._smem_layout_key()
-        if not hasattr(params, smem_layout_key):
-            return None
-        smem_layout = getattr(params, smem_layout_key)
         return getattr(storage_epi, f"s_{self.name}").get_tensor(
             smem_layout.outer,
             swizzle=smem_layout.inner,
         )
     def tma_atoms(self, gemm, params):
-        tma_key = self._tma_atom_key()
-        if hasattr(params, tma_key):
-            return [getattr(params, tma_key)]
-        return []
 @cute.jit
 def vec_multiply(gemm, tRS_rD, tDrColVec, tDrRowVec):
-    """Multiply tRS_rD by colvec and/or rowvec in-place. Uses packed f32x2 on SM100+."""
     if const_expr(tDrColVec is not None):
-        if const_expr(gemm.arch < 100):
             for i in cutlass.range(cute.size(tDrColVec), unroll_full=True):
                 tRS_rD[i] *= tDrColVec[i]
         else:
@@ -487,7 +710,7 @@ def vec_multiply(gemm, tRS_rD, tDrColVec, tDrRowVec):
                     (tDrColVec[2 * i], tDrColVec[2 * i + 1]),
                 )
     if const_expr(tDrRowVec is not None):
-        if const_expr(gemm.arch < 100):
             for i in cutlass.range(cute.size(tDrRowVec), unroll_full=True):
                 tRS_rD[i] *= tDrRowVec[i]
         else:
@@ -503,13 +726,13 @@ def colvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rSc
     """Accumulate transform_fn(input) or input * rScale into a ColVecReduce buffer.
     If transform_fn is provided, accumulates transform_fn(input[i]).
-    If rScale is provided, accumulates input[i] * rScale[i] (uses mul/fma for SM100).
     If neither, accumulates input directly (identity).
     """
     if const_expr(tDrReduce is not None):
         if const_expr(transform_fn is None):
             transform_fn = lambda x: x
-        if const_expr(gemm.arch < 100):
             for i in cutlass.range(cute.size(tDrReduce), unroll_full=True):
                 val = transform_fn(tRS_rInput[i])
                 tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val
@@ -521,6 +744,7 @@ def colvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rSc
             for m in cutlass.range(cute.size(tDrReduce_mn, mode=[0]), unroll_full=True):
                 inp = lambda n: (tRS_rInput_mn[m, 2 * n], tRS_rInput_mn[m, 2 * n + 1])
                 val0 = transform_fn(inp(0))
                 if const_expr(rScale is not None):
                     row_sum = cute.arch.mul_packed_f32x2(val0, (rScale_mn[m, 0], rScale_mn[m, 1]))
                 else:
@@ -536,14 +760,48 @@ def colvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rSc
                 tDrReduce_mn[m, 0] += row_sum[0] + row_sum[1]
-class ColVecReduce(EpiOp):
-    """Column vector reduction: accumulates across N subtiles in registers,
-    then warp-reduces and writes to gmem in epi_end.
-    No smem. The accumulation itself happens in epi_visit_subtile (user code).
-    This op handles the register allocation (begin), per-subtile slicing (begin_loop),
-    and final warp reduction + gmem write (end).
     """
     def param_fields(self):
         return [(self.name, object, None)]
@@ -551,31 +809,83 @@ class ColVecReduce(EpiOp):
     def to_params(self, gemm, args):
         return {self.name: assume_stride_divisibility(getattr(args, self.name))}
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
-        tDrReduce = None
-        if const_expr(param is not None):
-            colvec_mma_layout = cute.make_layout((ctx.tile_M, ctx.tile_N), stride=(1, 0))
-            tDrReduce_layout = ctx.partition_for_epilogue_fn(
-                cute.make_rmem_tensor(colvec_mma_layout, Float32)
-            ).layout
-            tDrReduce = cute.make_rmem_tensor(tDrReduce_layout, Float32)
-            cute.filter_zeros(tDrReduce).fill(0.0)
-        return tDrReduce
     @cute.jit
     def begin_loop(self, gemm, state, epi_coord):
-        result = None
-        if const_expr(state is not None):
-            result = cute.group_modes(state, 3, cute.rank(state))[None, None, None, epi_coord]
         return result
     @cute.jit
-    def end(
         self,
         gemm,
         param,
         state,
         epi_tile,
         tiled_copy_t2r,
         tiled_copy_r2s,
@@ -583,9 +893,13 @@ class ColVecReduce(EpiOp):
         varlen_manager,
         tidx,
     ):
-        """Intra-warp shuffle reduction across N lanes, then direct gmem write."""
-        if const_expr(param is not None):
-            tDrReduce = state
             tiled_copy = tiled_copy_t2r if tiled_copy_t2r is not None else tiled_copy_r2s
             reference_src = tiled_copy_t2r is None
@@ -593,26 +907,147 @@ class ColVecReduce(EpiOp):
             lane_layout_MN, warp_layout_MN = _get_lane_warp_layouts(tiled_copy, reference_src)
             # For ColVecReduce: reduce across N lanes (lanes_in_N threads share same M row)
             lanes_in_N = cute.size(lane_layout_MN, mode=[1])
             # Typically lanes_in_N is 4 for Sm90
             assert lanes_in_N == 1 << int(math.log2(lanes_in_N)), (
                 "lanes_in_N must be a power of 2 for butterfly reduction"
             )
-            # ── Intra-warp shuffle reduction across N lanes ──
             if const_expr(lanes_in_N > 1):
                 assert lane_layout_MN.stride[1] == 1
-                tDrReduce_flt = cute.filter_zeros(tDrReduce)
                 for i in cutlass.range(cute.size(tDrReduce_flt), unroll_full=True):
                     tDrReduce_flt[i] = cute.arch.warp_reduction(
                         tDrReduce_flt[i], operator.add, threads_in_group=lanes_in_N
                     )
             warp_N = warp_layout_MN[1]
-            assert cute.size(warp_N) == 1, (
-                "ColVecReduce assumes all reduction cols are within the same warp"
             )
-            # ── Direct gmem write (no inter-warp reduction needed: warps_in_N == 1) ──
             partition_for_epilogue_fn = partial(
                 partition_for_epilogue,
                 epi_tile=epi_tile,
@@ -621,28 +1056,46 @@ class ColVecReduce(EpiOp):
                 reference_src=tiled_copy_t2r is None,
             )
             tile_M, tile_N = gemm.cta_tile_shape_mnk[:2]
             batch_idx = tile_coord_mnkl[3]
-            limit_n = param.shape[2] if not varlen_manager.varlen_m else param.shape[1]
-            if tile_coord_mnkl[1] < limit_n:
-                if const_expr(not varlen_manager.varlen_m):
-                    mColVec = param[batch_idx, None, tile_coord_mnkl[1]]
-                else:
-                    mColVec = cute.domain_offset(
-                        (varlen_manager.params.cu_seqlens_m[batch_idx],),
-                        param[None, tile_coord_mnkl[1]],
-                    )
-                gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],))
-                limit_m = min(
-                    varlen_manager.len_m(batch_idx) - tile_coord_mnkl[0] * tile_M,
-                    tile_M,
-                )
-                tDcD = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N)))
-                tDrReduce_m = layout_utils.convert_layout_zero_stride(tDrReduce, tDrReduce.layout)[
-                    None, 0
-                ]
-                tDcD_m = layout_utils.convert_layout_zero_stride(tDcD, tDrReduce.layout)[None, 0]
-                if tDcD_m[0][1] == 0:
-                    for m in cutlass.range(cute.size(tDcD_m, mode=[0])):
-                        row_idx = tDcD_m[m][0]
-                        if row_idx < limit_m:
-                            gColVec[row_idx] = tDrReduce_m[m]

 smem allocation, begin (one-time per-tile setup), begin_loop (per-subtile extraction),
 end (cleanup).
+The ops are composed via ComposableEpiMixin. Class-level `_epi_ops` is the
+static schema; `_epi_ops_to_params_dict` (called from each subclass's
+`epi_to_underlying_arguments`) shadows it with an instance-level tuple of only
+the active ops (those whose arg tensor is non-None). All EpiOp hook methods
+below therefore assume their `param` / `arg_tensor` is non-None — the
+framework guarantees inactive ops are never iterated.
 """
 import math
 import operator
 from functools import partial
+from typing import NamedTuple
 import cutlass
 import cutlass.cute as cute
 class EpiContext:
+    """Shared context passed to EpiOp.begin methods. Bundles common arguments.
+    `tRS_rD_layout` is only populated by callers that need TileLoad — it's the
+    register layout of the matmul output tile, which TileLoad uses to shape its
+    own register tile so it lines up element-wise with tRS_rD in epi_visit_subtile.
+    """
     __slots__ = (
         "epi_tile",
         "varlen_manager",
         "epilogue_barrier",
         "tidx",
+        "tRS_rD_layout",
         "partition_for_epilogue_fn",
         "num_epi_threads",
         "batch_idx",
         varlen_manager,
         epilogue_barrier,
         tidx,
+        tRS_rD_layout=None,
     ):
         self.epi_tile = epi_tile
         self.tiled_copy_t2r = tiled_copy_t2r
         self.varlen_manager = varlen_manager
         self.epilogue_barrier = epilogue_barrier
         self.tidx = tidx
+        self.tRS_rD_layout = tRS_rD_layout
         self.tile_M = gemm.cta_tile_shape_mnk[0]
         self.tile_N = gemm.cta_tile_shape_mnk[1]
         self.batch_idx = tile_coord_mnkl[3]
     return lane_layout_MN, warp_layout_MN
+class EpiSmemBytes(NamedTuple):
+    """Shared-memory accounting for one epilogue op.
+    unstaged: allocated once per CTA tile.
+    d_stage: allocated per D/store epilogue stage.
+    c_stage: allocated per C/load epilogue stage.
+    """
+    unstaged: int = 0
+    d_stage: int = 0
+    c_stage: int = 0
+    def __add__(self, other):
+        return EpiSmemBytes(
+            self.unstaged + other.unstaged,
+            self.d_stage + other.d_stage,
+            self.c_stage + other.c_stage,
+        )
+    def __radd__(self, other):
+        if other == 0:
+            return self
+        return self.__add__(other)
 class EpiOp:
     """Base class for composable epilogue operations."""
         Returns dict of {param_name: value}. Like EVT's to_underlying_arguments."""
         return {}
+    def epi_m_major_score(self, arg_tensor, gemm):
+        """Preference for epilogue subtile order. Positive prefers M-major, negative N-major."""
         return 0
+    # --- Host-side: smem allocation ---
+    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None):
+        """Bytes of smem needed by unstaged / D-stage / C-stage storage."""
+        return EpiSmemBytes()
     def smem_struct_field(self, gemm, params):
         """Return (field_name, field_type) for @cute.struct, or None if no smem needed.
         params is the full EpilogueParams object."""
         """Return list of TMA atoms for this op."""
         return []
+    def is_tile_load(self):
+        """Whether this op is a tile-sized epilogue input loaded through the C pipeline."""
+        return False
+    def load_g2s_copy_fn(
+        self,
+        gemm,
+        params,
+        smem_tensor,
+        tile_coord_mnkl,
+        varlen_manager,
+        epi_pipeline,
+    ):
+        """Return a per-subtile gmem->smem copy function, or None."""
+        return None
     # --- Device-side: kernel execution ---
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
         """Per-subtile extraction. Returns value for epi_visit_subtile."""
         return state
+    @cute.jit
+    def load_s2r(self, gemm, param, state, stage_idx):
+        """Issue this op's tile-load smem->register copy for one epilogue stage."""
+        pass
+    def end_loop(
+        self,
+        gemm,
+        param,
+        state,
+        epi_coord,
+        epi_tile,
+        tiled_copy_t2r,
+        tiled_copy_r2s,
+        tile_coord_mnkl,
+        varlen_manager,
+        tidx,
+    ):
+        """Per-subtile cleanup after epi_visit_subtile."""
+        pass
     def needs_async_fence(self):
         """Whether this op issues async copies that need a fence."""
         return False
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
+        if const_expr(self.dtype is not None):
+            return utils.load_scalar_or_pointer(param, dtype=self.dtype)
+        return utils.load_scalar_or_pointer(param)
 class VecLoad(EpiOp):
     def _coord_idx(self):
         return 1 if self.dim == 1 else 0
+    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None):
+        return EpiSmemBytes(
+            unstaged=self._tile_size(cta_tile_shape_mnk) * (arg_tensor.element_type.width // 8)
+        )
     def smem_struct_field(self, gemm, params):
+        tensor = getattr(params, self.name)
+        size = self._tile_size(gemm.cta_tile_shape_mnk)
+        return (
+            f"s_{self.name}",
+            cute.struct.Align[cute.struct.MemRange[tensor.element_type, size], 16],
+        )
     def get_smem_tensor(self, gemm, params, storage_epi):
         return getattr(storage_epi, f"s_{self.name}").get_tensor(
             cute.make_layout(self._tile_size(gemm.cta_tile_shape_mnk))
         )
     def needs_async_fence(self):
         return True
+    def epi_m_major_score(self, arg_tensor, gemm):
+        # It costs more registers (say 4x) to keep rowvec in register vs keeping colvec in register
+        return 4 if self.dim == 1 else -1
     def _get_gmem_vec(self, param, ctx):
         """Get the global memory vector for this tile. Override for varlen."""
         return param[ctx.batch_idx, None]
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
+        dtype = param.element_type
+        num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width
+        thr_copy = copy_utils.tiled_copy_1d(
+            dtype, ctx.num_epi_threads, num_copy_elems, is_async=True
+        ).get_slice(ctx.tidx)
+        mVec = self._get_gmem_vec(param, ctx)
+        tile_dim = self._tile_dim(ctx)
+        coord_idx = ctx.tile_coord_mnkl[self._coord_idx()]
+        gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,))
+        tVgV = thr_copy.partition_S(gVec)
+        tVsV = thr_copy.partition_D(smem_tensor)
+        tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim))
+        limit = min(cute.size(mVec, mode=[0]) - coord_idx * tile_dim, tile_dim)
+        for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True):
+            if tVcV[0, m] < tile_dim:  # Guard to avoid writing beyond the smem we've allocated
+                pred = cute.make_rmem_tensor(1, Boolean)
+                pred[0] = tVcV[0, m] < limit
+                cute.copy(thr_copy, tVgV[None, m], tVsV[None, m], pred=pred)
+        tDsV = ctx.partition_for_epilogue_fn(
+            cute.make_tensor(
+                smem_tensor.iterator,
+                cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()),
             )
+        )
+        if const_expr(ctx.tiled_copy_t2r is not None):
+            tDsV = ctx.tiled_copy_r2s.retile(tDsV)
+        # Pre-allocate register tensor reused across begin_loop calls
+        tDsV_sub = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, 0]
+        tDrV_cvt = cute.make_rmem_tensor(tDsV_sub.layout, gemm.acc_dtype)
+        return [tDsV, tDrV_cvt]
     @cute.jit
     def begin_loop(self, gemm, state, epi_coord):
+        tDsV, tDrV_cvt = state[0], state[1]
+        should_load = Boolean(True)
+        if const_expr(self.dim == 1):
+            if const_expr(gemm.epi_m_major):
+                should_load = epi_coord[0] == 0
+        else:
+            if const_expr(not gemm.epi_m_major):
+                should_load = epi_coord[1] == 0
+        if should_load:
+            tDsV_cur = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, epi_coord]
             tDrV = cute.make_rmem_tensor(tDsV_cur.layout, tDsV_cur.element_type)
             cute.autovec_copy(cute.filter_zeros(tDsV_cur), cute.filter_zeros(tDrV))
             tDrV_cvt.store(tDrV.load().to(gemm.acc_dtype))
         return tDrV_cvt
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
+        dtype = param.element_type
+        num_copy_elems = const_expr(max(32, dtype.width)) // dtype.width
+        thr_copy = copy_utils.tiled_copy_1d(
+            dtype, ctx.num_epi_threads, num_copy_elems, is_async=True
+        ).get_slice(ctx.tidx)
+        mVec = self._get_gmem_vec(param, ctx)
+        tile_dim = self._tile_dim(ctx)
+        coord_idx = ctx.tile_coord_mnkl[self._coord_idx()]
+        gVec = cute.local_tile(mVec, (tile_dim,), (coord_idx,))
+        tVgV = thr_copy.partition_S(gVec)
+        tVsV = thr_copy.partition_D(smem_tensor)
+        tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim))
+        # ColVec uses varlen-aware limit
+        limit = min(
+            ctx.varlen_manager.len_m(ctx.batch_idx) - coord_idx * tile_dim,
+            tile_dim,
+        )
+        for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True):
+            if tVcV[0, m] < tile_dim:  # Guard to avoid writing beyond the smem we've allocated
+                pred = cute.make_rmem_tensor(1, Boolean)
+                pred[0] = tVcV[0, m] < limit
+                cute.copy(thr_copy, tVgV[None, m], tVsV[None, m], pred=pred)
+        tDsV = ctx.partition_for_epilogue_fn(
+            cute.make_tensor(
+                smem_tensor.iterator,
+                cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride()),
             )
+        )
+        if const_expr(ctx.tiled_copy_t2r is not None):
+            tDsV = ctx.tiled_copy_r2s.retile(tDsV)
+        # Pre-allocate register tensor reused across begin_loop calls
+        tDsV_sub = cute.group_modes(tDsV, 3, cute.rank(tDsV))[None, None, None, 0]
+        tDrV_cvt = cute.make_rmem_tensor(tDsV_sub.layout, gemm.acc_dtype)
         return [tDsV, tDrV_cvt]
 class TileStore(EpiOp):
     """Tile-sized output tensor stored via TMA (e.g. postact).
     Args:
+        name: field name in EpilogueArguments/Params (e.g. "mAuxOut")
         epi_tile_fn: optional (gemm, epi_tile) -> epi_tile for half-tile (GemmGated)
     """
         return f"epi_tile_{self.name}"
     def param_fields(self):
+        # Defaults are None so EpilogueParams can be constructed when this op is
+        # filtered out (inactive). Active calls always set all four via to_params.
         return [
+            (self._tma_atom_key(), object, None),
+            (self.name, object, None),
+            (self._smem_layout_key(), object, None),
+            (self._epi_tile_key(), object, None),
         ]
     def to_params(self, gemm, args):
             self._epi_tile_key(): epi_tile_out,
         }
+    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None):
         if self.epi_tile_fn is not None:
             epi_tile = self.epi_tile_fn(None, epi_tile)
+        # epi_tile may contain Layout entries (from SM100's compute_epilogue_tile_shape
+        # fixup path), so extract the int shape first.
+        return EpiSmemBytes(
+            d_stage=cute.size(cute.shape(epi_tile)) * (arg_tensor.element_type.width // 8)
+        )
     def smem_struct_field(self, gemm, params):
+        smem_layout = getattr(params, self._smem_layout_key())
         return (
             f"s_{self.name}",
             cute.struct.Align[
                 cute.struct.MemRange[
+                    gemm.aux_out_dtype,
+                    cute.cosize(smem_layout),
                 ],
                 gemm.buffer_align_bytes,
             ],
         )
     def get_smem_tensor(self, gemm, params, storage_epi):
+        smem_layout = getattr(params, self._smem_layout_key())
         return getattr(storage_epi, f"s_{self.name}").get_tensor(
             smem_layout.outer,
             swizzle=smem_layout.inner,
         )
     def tma_atoms(self, gemm, params):
+        return [getattr(params, self._tma_atom_key())]
+class _TileLoadState(NamedTuple):
+    """Per-tile register state produced by TileLoad.begin and consumed by load_s2r /
+    begin_loop. tRS_rTile is the register tile partitioned to match tRS_rD's layout;
+    tSR_sTile / tSR_rTile drive the per-stage smem→register copy."""
+    tiled_copy_s2r: object
+    tRS_rTile: object
+    tSR_rTile: object
+    tSR_sTile: object
+class TileLoad(EpiOp):
+    """Tile-sized auxiliary input loaded through the epilogue load pipeline.
+    TileLoad uses the same staged gmem->smem->register pipeline as GEMM's C operand,
+    but it is exposed to the epilogue as ``epi_loop_tensors[name]`` instead of as
+    ``tRS_rC``. That lets custom epilogues consume extra MxN tensors without using
+    the GEMM C argument.
+    Its shared memory is accounted as ``EpiSmemBytes.c_stage``, so it is allocated
+    per epilogue load stage. Multiple TileLoads are supported: each has its own TMA
+    descriptor and smem buffer, and the pipeline transaction count includes C plus
+    all enabled TileLoad buffers. Supported on SM90, SM100, and SM120.
+    """
+    def __init__(self, name, epi_tile_fn=None):
+        super().__init__(name)
+        self.epi_tile_fn = epi_tile_fn
+    def _tma_atom_key(self):
+        return f"tma_atom_{self.name}"
+    def _smem_layout_key(self):
+        return f"epi_{self.name}_smem_layout_staged"
+    def _epi_tile_key(self):
+        return f"epi_tile_{self.name}"
+    # The original LayoutEnum and element_type can't be recovered from the
+    # TMA-prepared tensor that ends up in params (`from_tensor` returns a typing
+    # annotation post-TMA, not a Numeric class). We stash both on the gemm at
+    # to_params time and read them back in begin(). The dtype is also exposed on
+    # the params dataclass for smem_struct_field.
+    def _layout_gemm_attr(self):
+        return f"_tile_load_layout_{self.name}"
+    def _dtype_gemm_attr(self):
+        return f"_tile_load_dtype_{self.name}"
+    def _dtype_field(self):
+        return f"{self.name}_dtype"
+    def param_fields(self):
+        # Defaults are None so EpilogueParams can be constructed when this op is
+        # filtered out (inactive). Active calls always set all five via to_params.
+        return [
+            (self._tma_atom_key(), object, None),
+            (self.name, object, None),
+            (self._smem_layout_key(), object, None),
+            (self._epi_tile_key(), object, None),
+            (self._dtype_field(), object, None),
+        ]
+    def to_params(self, gemm, args):
+        tensor = getattr(args, self.name)
+        setattr(gemm, self._layout_gemm_attr(), cutlass.utils.LayoutEnum.from_tensor(tensor))
+        setattr(gemm, self._dtype_gemm_attr(), tensor.element_type)
+        epi_tile = self.epi_tile_fn(gemm, gemm.epi_tile) if self.epi_tile_fn else None
+        tma_atom, tma_tensor, smem_layout, epi_tile_out = setup_epi_tensor(
+            gemm, tensor, epi_tile=epi_tile, op_type="load", stage=gemm.epi_c_stage
+        )
+        return {
+            self._tma_atom_key(): tma_atom,
+            self.name: tma_tensor,
+            self._smem_layout_key(): smem_layout,
+            self._epi_tile_key(): epi_tile_out,
+            self._dtype_field(): tensor.element_type,
+        }
+    def is_tile_load(self):
+        return True
+    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None):
+        if self.epi_tile_fn is not None:
+            epi_tile = self.epi_tile_fn(None, epi_tile)
+        # epi_tile may contain Layout entries from SM100's compute_epilogue_tile_shape
+        # fixup; extract the int shape first.
+        return EpiSmemBytes(
+            c_stage=cute.size(cute.shape(epi_tile)) * (arg_tensor.element_type.width // 8)
+        )
+    def smem_struct_field(self, gemm, params):
+        smem_layout = getattr(params, self._smem_layout_key())
+        dtype = getattr(params, self._dtype_field())
+        return (
+            f"s_{self.name}",
+            cute.struct.Align[
+                cute.struct.MemRange[dtype, cute.cosize(smem_layout)],
+                gemm.buffer_align_bytes,
+            ],
+        )
+    def get_smem_tensor(self, gemm, params, storage_epi):
+        smem_layout = getattr(params, self._smem_layout_key())
+        return getattr(storage_epi, f"s_{self.name}").get_tensor(
+            smem_layout.outer,
+            swizzle=smem_layout.inner,
+        )
+    def tma_atoms(self, gemm, params):
+        return [getattr(params, self._tma_atom_key())]
+    def load_g2s_copy_fn(
+        self,
+        gemm,
+        params,
+        smem_tensor,
+        tile_coord_mnkl,
+        varlen_manager,
+        epi_pipeline,
+    ):
+        tensor = getattr(params, self.name)
+        batch_idx = tile_coord_mnkl[3]
+        copy_tile_fn, _, _ = gemm.epilog_gmem_copy_and_partition(
+            getattr(params, self._tma_atom_key()),
+            varlen_manager.offset_batch_epi(tensor, batch_idx),
+            gemm.cta_tile_shape_mnk[:2],
+            getattr(params, self._epi_tile_key()),
+            smem_tensor,
+            tile_coord_mnkl,
+        )
+        return copy_utils.tma_producer_copy_fn(copy_tile_fn, epi_pipeline)
+    @cute.jit
+    def begin(self, gemm, param, smem_tensor, ctx):
+        assert gemm.arch in (90, 100, 120), "TileLoad requires the SM90/SM100/SM120 epilogue path"
+        assert ctx.tRS_rD_layout is not None
+        smem_load_ref = ctx.tiled_copy_t2r if const_expr(gemm.arch == 100) else gemm.tiled_mma
+        tiled_copy_s2r, tRS_rTile, tSR_rTile, tSR_sTile = gemm.epilog_smem_load_and_partition(
+            smem_load_ref,
+            getattr(gemm, self._layout_gemm_attr()),
+            getattr(gemm, self._dtype_gemm_attr()),
+            smem_tensor,
+            ctx.tRS_rD_layout,
+            ctx.tidx,
+        )
+        # Shape: (s2r-copy-handle, register-tile-as-rD-layout, smem→r retile target,
+        # smem→r staged source). begin_loop returns tRS_rTile; load_s2r uses the rest.
+        return _TileLoadState(tiled_copy_s2r, tRS_rTile, tSR_rTile, tSR_sTile)
+    @cute.jit
+    def load_s2r(self, gemm, param, state, stage_idx):
+        cute.copy(
+            state.tiled_copy_s2r,
+            state.tSR_sTile[None, None, None, stage_idx],
+            state.tSR_rTile,
+        )
+    @cute.jit
+    def begin_loop(self, gemm, state, epi_coord):
+        return state.tRS_rTile
 @cute.jit
 def vec_multiply(gemm, tRS_rD, tDrColVec, tDrRowVec):
+    """Multiply tRS_rD by colvec and/or rowvec in-place. Uses packed f32x2 on SM100."""
     if const_expr(tDrColVec is not None):
+        if const_expr(gemm.arch != 100):
             for i in cutlass.range(cute.size(tDrColVec), unroll_full=True):
                 tRS_rD[i] *= tDrColVec[i]
         else:
                     (tDrColVec[2 * i], tDrColVec[2 * i + 1]),
                 )
     if const_expr(tDrRowVec is not None):
+        if const_expr(gemm.arch != 100):
             for i in cutlass.range(cute.size(tDrRowVec), unroll_full=True):
                 tRS_rD[i] *= tDrRowVec[i]
         else:
     """Accumulate transform_fn(input) or input * rScale into a ColVecReduce buffer.
     If transform_fn is provided, accumulates transform_fn(input[i]).
+    If rScale is provided, accumulates input[i] * rScale[i] (uses packed mul/fma for SM100).
     If neither, accumulates input directly (identity).
     """
     if const_expr(tDrReduce is not None):
         if const_expr(transform_fn is None):
             transform_fn = lambda x: x
+        if const_expr(gemm.arch != 100):
             for i in cutlass.range(cute.size(tDrReduce), unroll_full=True):
                 val = transform_fn(tRS_rInput[i])
                 tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val
             for m in cutlass.range(cute.size(tDrReduce_mn, mode=[0]), unroll_full=True):
                 inp = lambda n: (tRS_rInput_mn[m, 2 * n], tRS_rInput_mn[m, 2 * n + 1])
                 val0 = transform_fn(inp(0))
+                assert cute.size(tDrReduce_mn, mode=[1]) % 2 == 0
                 if const_expr(rScale is not None):
                     row_sum = cute.arch.mul_packed_f32x2(val0, (rScale_mn[m, 0], rScale_mn[m, 1]))
                 else:
                 tDrReduce_mn[m, 0] += row_sum[0] + row_sum[1]
+@cute.jit
+def rowvec_reduce_accumulate(gemm, tDrReduce, tRS_rInput, transform_fn=None, rScale=None):
+    """Accumulate transform_fn(input) or input * rScale into a RowVecReduce buffer.
+    Reduces along M dimension, keeping N. The zero-stride layout on M ensures
+    elements at different M positions but same N column accumulate correctly.
     """
+    if const_expr(tDrReduce is not None):
+        if const_expr(transform_fn is None):
+            transform_fn = lambda x: x
+        if const_expr(gemm.arch != 100):
+            for i in cutlass.range(cute.size(tDrReduce), unroll_full=True):
+                val = transform_fn(tRS_rInput[i])
+                tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val
+        else:
+            # Keep CUTLASS's linear fragment indexing, but use packed f32x2 arithmetic
+            # for any transform that accepts and returns an f32x2 tuple.
+            # We have to be careful to avoid tDrReduce[2 * i] and tDrReduce[2 * i + 1] aliasing
+            # each other. For SM100, tDrReduce has layout ((32,1),1,1):((1,0),0,0) or
+            # (((2,2,4),1),2,1):(((1,0,8),0),0,0), so this works. But it's error-prone.
+            for i in cutlass.range(cute.size(tRS_rInput) // 2, unroll_full=True):
+                acc = (tDrReduce[2 * i], tDrReduce[2 * i + 1])
+                val = (tRS_rInput[2 * i], tRS_rInput[2 * i + 1])
+                val = transform_fn(val)
+                if const_expr(rScale is not None):
+                    scale = (rScale[2 * i], rScale[2 * i + 1])
+                    tDrReduce[2 * i], tDrReduce[2 * i + 1] = cute.arch.fma_packed_f32x2(
+                        val, scale, acc
+                    )
+                else:
+                    tDrReduce[2 * i], tDrReduce[2 * i + 1] = cute.arch.add_packed_f32x2(val, acc)
+            if const_expr(cute.size(tRS_rInput) % 2 != 0):
+                i = cute.size(tRS_rInput) - 1
+                val = transform_fn(tRS_rInput[i])
+                tDrReduce[i] += val * rScale[i] if const_expr(rScale is not None) else val
+class VecReduce(EpiOp):
+    """Base class for row/column vector reductions."""
+    dim = 0  # 0 for colvec output along M, 1 for rowvec output along N
+    epi_m_major_preference = 0
     def param_fields(self):
         return [(self.name, object, None)]
     def to_params(self, gemm, args):
         return {self.name: assume_stride_divisibility(getattr(args, self.name))}
+    def epi_m_major_score(self, arg_tensor, gemm):
+        return self.epi_m_major_preference
+    def _tile_size(self, cta_tile_shape_mnk):
+        return cta_tile_shape_mnk[self.dim]
+    def _broadcast_stride(self):
+        # Col: stride (1,0) broadcasts along N. Row: stride (0,1) broadcasts along M.
+        return (1, 0) if self.dim == 0 else (0, 1)
+    def _reduce_dim(self):
+        return 1 - self.dim
+    def _smem_warps(self, warp_shape_mnk):
+        warps = warp_shape_mnk[self._reduce_dim()] if warp_shape_mnk is not None else 1
+        return max(warps - 1, 0)
+    def smem_bytes(self, arg_tensor, cta_tile_shape_mnk, epi_tile, warp_shape_mnk=None):
+        smem_warps = self._smem_warps(warp_shape_mnk)
+        if smem_warps == 0:
+            return EpiSmemBytes()
+        return EpiSmemBytes(
+            unstaged=self._tile_size(cta_tile_shape_mnk) * smem_warps * (Float32.width // 8)
+        )
+    def smem_struct_field(self, gemm, params):
+        smem_warps = self._smem_warps(gemm.epi_smem_warp_shape_mnk())
+        if smem_warps == 0:
+            return None
+        size = self._tile_size(gemm.cta_tile_shape_mnk) * smem_warps
+        return (f"s_{self.name}", cute.struct.Align[cute.struct.MemRange[Float32, size], 16])
+    def get_smem_tensor(self, gemm, params, storage_epi):
+        smem_warps = self._smem_warps(gemm.epi_smem_warp_shape_mnk())
+        if smem_warps == 0:
+            return None
+        return getattr(storage_epi, f"s_{self.name}").get_tensor(
+            cute.make_layout((self._tile_size(gemm.cta_tile_shape_mnk), smem_warps))
+        )
     @cute.jit
     def begin(self, gemm, param, smem_tensor, ctx):
+        vec_mma_layout = cute.make_layout((ctx.tile_M, ctx.tile_N), stride=self._broadcast_stride())
+        tDrReduce_layout = ctx.partition_for_epilogue_fn(
+            cute.make_rmem_tensor(vec_mma_layout, Float32)
+        ).layout
+        tDrReduce = cute.make_rmem_tensor(tDrReduce_layout, Float32)
+        return (tDrReduce, smem_tensor)
     @cute.jit
     def begin_loop(self, gemm, state, epi_coord):
+        tDrReduce = state[0]
+        result = tDrReduce[None, None, None, epi_coord[0], epi_coord[1]]
+        if const_expr(epi_coord[self._reduce_dim()] == 0):
+            cute.filter_zeros(result).fill(0.0)
         return result
+class ColVecReduce(VecReduce):
+    """Column vector reduction: accumulates across N subtiles in registers,
+    then reduces across N lanes/warps and writes to gmem per completed M stripe.
+    The accumulation itself happens in epi_visit_subtile (user code).
+    This op handles the register allocation (begin), per-subtile slicing (begin_loop),
+    and reduction + gmem write (end_loop).
+    """
+    dim = 0
+    epi_m_major_preference = -1
     @cute.jit
+    def end_loop(
         self,
         gemm,
         param,
         state,
+        epi_coord,
         epi_tile,
         tiled_copy_t2r,
         tiled_copy_r2s,
         varlen_manager,
         tidx,
     ):
+        """Flush the current M stripe when the last N subtile has accumulated."""
+        epi_tile_shape = cute.zipped_divide(
+            cute.make_layout(gemm.cta_tile_shape_mnk[:2]), epi_tile
+        ).shape[1]
+        if const_expr(epi_coord[1] == epi_tile_shape[1] - 1):
+            tDrReduce, sDrReduce = state[0], state[1]
+            tDrReduce_cur = tDrReduce[None, None, None, epi_coord[0], epi_coord[1]]
             tiled_copy = tiled_copy_t2r if tiled_copy_t2r is not None else tiled_copy_r2s
             reference_src = tiled_copy_t2r is None
             lane_layout_MN, warp_layout_MN = _get_lane_warp_layouts(tiled_copy, reference_src)
             # For ColVecReduce: reduce across N lanes (lanes_in_N threads share same M row)
             lanes_in_N = cute.size(lane_layout_MN, mode=[1])
+            is_lane_n_leader = cute.arch.lane_idx() % lanes_in_N == 0
             # Typically lanes_in_N is 4 for Sm90
             assert lanes_in_N == 1 << int(math.log2(lanes_in_N)), (
                 "lanes_in_N must be a power of 2 for butterfly reduction"
             )
+            # Intra-warp shuffle reduction across N lanes
             if const_expr(lanes_in_N > 1):
+                # Assumes threads for each M row are contiguous along N, so
+                # warp_reduction over groups of lanes_in_N matches lane_layout_MN.
                 assert lane_layout_MN.stride[1] == 1
+                tDrReduce_flt = cute.filter_zeros(tDrReduce_cur)
                 for i in cutlass.range(cute.size(tDrReduce_flt), unroll_full=True):
                     tDrReduce_flt[i] = cute.arch.warp_reduction(
                         tDrReduce_flt[i], operator.add, threads_in_group=lanes_in_N
                     )
             warp_N = warp_layout_MN[1]
+            warps_in_N = const_expr(cute.size(warp_N))
+            partition_for_epilogue_fn = partial(
+                partition_for_epilogue,
+                epi_tile=epi_tile,
+                tiled_copy=tiled_copy,
+                tidx=tidx,
+                reference_src=tiled_copy_t2r is None,
+            )
+            tile_M, tile_N = gemm.cta_tile_shape_mnk[:2]
+            tDcD = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N)))
+            tDcD_cur = tDcD[None, None, None, epi_coord[0], epi_coord[1]]
+            tDrReduce_m = layout_utils.convert_layout_zero_stride(
+                tDrReduce_cur, tDrReduce_cur.layout
+            )[None, 0]
+            tDcD_m = layout_utils.convert_layout_zero_stride(tDcD_cur, tDrReduce_cur.layout)[
+                None, 0
+            ]
+            # Inter-warp reduction through smem
+            warp_idx = cute.arch.make_warp_uniform(tidx // cute.arch.WARP_SIZE)
+            warp_n_idx = warp_layout_MN.get_hier_coord(warp_idx)[1]
+            if const_expr(warps_in_N > 1):
+                if warp_n_idx > 0 and is_lane_n_leader:
+                    for m in cutlass.range(cute.size(tDcD_m, mode=[0])):
+                        row_idx = tDcD_m[m][0]
+                        sDrReduce[row_idx, warp_n_idx - 1] = tDrReduce_m[m]
+                gemm.epilogue_barrier.arrive_and_wait()
+                if warp_n_idx == 0 and is_lane_n_leader:
+                    for m in cutlass.range(cute.size(tDcD_m, mode=[0])):
+                        row_idx = tDcD_m[m][0]
+                        for warp_n in cutlass.range_constexpr(1, warps_in_N):
+                            tDrReduce_m[m] += sDrReduce[row_idx, warp_n - 1]
+            # Write to gmem
+            batch_idx = tile_coord_mnkl[3]
+            limit_m = min(varlen_manager.len_m(batch_idx) - tile_coord_mnkl[0] * tile_M, tile_M)
+            limit_n_tiles = param.shape[2] if not varlen_manager.varlen_m else param.shape[1]
+            if const_expr(not varlen_manager.varlen_m):
+                mColVec = param[batch_idx, None, tile_coord_mnkl[1]]
+            else:
+                mColVec = cute.domain_offset(
+                    (varlen_manager.params.cu_seqlens_m[batch_idx],),
+                    param[None, tile_coord_mnkl[1]],
+                )
+            gColVec = cute.local_tile(mColVec, (tile_M,), (tile_coord_mnkl[0],))
+            should_write_gmem = (
+                is_lane_n_leader
+                if const_expr(warps_in_N == 1)
+                else warp_n_idx == 0 and is_lane_n_leader
             )
+            if tile_coord_mnkl[1] < limit_n_tiles and should_write_gmem:
+                for m in cutlass.range(cute.size(tDcD_m, mode=[0])):
+                    row_idx = tDcD_m[m][0]
+                    if row_idx < limit_m:
+                        gColVec[row_idx] = tDrReduce_m[m]
+class RowVecReduce(VecReduce):
+    """Row vector reduction: accumulates across M subtiles in registers,
+    then reduces across M lanes/warps and writes to gmem per completed N stripe.
+    Output shape is (L, ceildiv(M, tile_M), N): one partial sum per CTA-M tile per
+    N column. This mirrors ColVecReduce with M/N swapped.
+    """
+    dim = 1
+    epi_m_major_preference = 4
+    @cute.jit
+    def end_loop(
+        self,
+        gemm,
+        param,
+        state,
+        epi_coord,
+        epi_tile,
+        tiled_copy_t2r,
+        tiled_copy_r2s,
+        tile_coord_mnkl,
+        varlen_manager,
+        tidx,
+    ):
+        """Flush the current N stripe when the last M subtile has accumulated."""
+        epi_tile_shape = cute.zipped_divide(
+            cute.make_layout(gemm.cta_tile_shape_mnk[:2]), epi_tile
+        ).shape[1]
+        if const_expr(epi_coord[0] == epi_tile_shape[0] - 1):
+            tDrReduce, sDrReduce = state[0], state[1]
+            tDrReduce_cur = tDrReduce[None, None, None, epi_coord[0], epi_coord[1]]
+            tiled_copy = tiled_copy_t2r if tiled_copy_t2r is not None else tiled_copy_r2s
+            reference_src = tiled_copy_t2r is None
+            # ── Derive lane layout from tiled_copy ──
+            lane_layout_MN, warp_layout_MN = _get_lane_warp_layouts(tiled_copy, reference_src)
+            # For RowVecReduce: reduce across M lanes (lanes_in_M threads share same N col)
+            lanes_in_M = cute.size(lane_layout_MN, mode=[0])
+            lanes_in_N = cute.size(lane_layout_MN, mode=[1])
+            is_lane_m_leader = cute.arch.lane_idx() < lanes_in_N
+            assert lanes_in_M == 1 << int(math.log2(lanes_in_M)), (
+                "lanes_in_M must be a power of 2 for butterfly reduction"
+            )
+            if const_expr(lanes_in_N > 1):
+                assert lane_layout_MN.stride[1] == 1, (
+                    "RowVecReduce assumes contiguous N lanes when lanes_in_N > 1"
+                )
+            # Intra-warp shuffle reduction across M lanes. M lanes may be either contiguous
+            # (SM100 N-major output) or strided by N lanes (SM100 M-major output).
+            tDrReduce_n = layout_utils.convert_layout_zero_stride(
+                tDrReduce_cur, tDrReduce_cur.layout
+            )[None, 0]
+            if const_expr(lanes_in_M > 1):
+                for n in cutlass.range(cute.size(tDrReduce_n), unroll_full=True):
+                    reduction_rows = lanes_in_M // 2
+                    while reduction_rows > 0:
+                        tDrReduce_n[n] += cute.arch.shuffle_sync_bfly(
+                            tDrReduce_n[n],
+                            offset=cute.crd2idx((reduction_rows, 0), lane_layout_MN),
+                        )
+                        reduction_rows = reduction_rows // 2
+            warp_M = warp_layout_MN[0]
+            warps_in_M = const_expr(cute.size(warp_M))
             partition_for_epilogue_fn = partial(
                 partition_for_epilogue,
                 epi_tile=epi_tile,
                 reference_src=tiled_copy_t2r is None,
             )
             tile_M, tile_N = gemm.cta_tile_shape_mnk[:2]
+            tDcD = partition_for_epilogue_fn(cute.make_identity_tensor((tile_M, tile_N)))
+            tDcD_cur = tDcD[None, None, None, epi_coord[0], epi_coord[1]]
+            tDcD_n = layout_utils.convert_layout_zero_stride(tDcD_cur, tDrReduce_cur.layout)[
+                None, 0
+            ]
+            # Inter-warp reduction through smem
+            warp_idx = cute.arch.make_warp_uniform(tidx // cute.arch.WARP_SIZE)
+            warp_m_idx = warp_layout_MN.get_hier_coord(warp_idx)[0]
+            if const_expr(warps_in_M > 1):
+                if warp_m_idx > 0 and is_lane_m_leader:
+                    for n in cutlass.range(cute.size(tDcD_n, mode=[0])):
+                        col_idx = tDcD_n[n][1]
+                        sDrReduce[col_idx, warp_m_idx - 1] = tDrReduce_n[n]
+                gemm.epilogue_barrier.arrive_and_wait()
+                if warp_m_idx == 0 and is_lane_m_leader:
+                    for n in cutlass.range(cute.size(tDcD_n, mode=[0])):
+                        col_idx = tDcD_n[n][1]
+                        for warp_m in cutlass.range_constexpr(1, warps_in_M):
+                            tDrReduce_n[n] += sDrReduce[col_idx, warp_m - 1]
+            # Write to gmem
             batch_idx = tile_coord_mnkl[3]
+            limit_m_tiles = param.shape[1] if not varlen_manager.varlen_m else param.shape[0]
+            if const_expr(not varlen_manager.varlen_m):
+                mRowVec = param[batch_idx, tile_coord_mnkl[0], None]
+            else:
+                mRowVec = param[tile_coord_mnkl[0], None]
+            gRowVec = cute.local_tile(mRowVec, (tile_N,), (tile_coord_mnkl[1],))
+            limit_n = min(
+                cute.size(mRowVec, mode=[0]) - tile_coord_mnkl[1] * tile_N,
+                tile_N,
+            )
+            should_write_gmem = (
+                is_lane_m_leader
+                if const_expr(warps_in_M == 1)
+                else warp_m_idx == 0 and is_lane_m_leader
+            )
+            if tile_coord_mnkl[0] < limit_m_tiles and should_write_gmem:
+                for n in cutlass.range(cute.size(tDcD_n, mode=[0])):
+                    col_idx = tDcD_n[n][1]
+                    if col_idx < limit_n:
+                        gRowVec[col_idx] = tDrReduce_n[n]

build/torch-cuda/quack/epi_utils.py CHANGED Viewed

@@ -32,27 +32,32 @@ def assume_broadcast_strides(*tensors):
     return [assume_stride_divisibility(t) for t in tensors]
-def setup_epi_tensor(gemm, tensor, epi_tile=None, op_type="store"):
-    """Create TMA atom + smem layout for a supplemental epilogue tensor.
     Args:
-        gemm: The GEMM object (provides arch, epi_stage, _make_tma_epi_atoms_and_tensors).
-        tensor: The global memory tensor to set up TMA for.
         epi_tile: Epilogue tile shape. Defaults to gemm.epi_tile.
         op_type: "store" or "load".
     Returns:
-        (tma_atom, tma_tensor, smem_layout_staged, epi_tile)
     """
     if epi_tile is None:
         epi_tile = gemm.epi_tile
     dtype = tensor.element_type
     layout = cutlass.utils.LayoutEnum.from_tensor(tensor)
     utils_cls = sm100_utils if gemm.arch >= 100 else sm90_utils
-    smem_layout_staged = utils_cls.make_smem_layout_epi(dtype, layout, epi_tile, gemm.epi_stage)
     tma_input = (
         copy_utils.create_ragged_tensor_for_tma(tensor, ragged_dim=0, ptr_shift=True)
-        if cute.rank(tensor) == 2
         else tensor
     )
     tma_atom, tma_tensor = gemm._make_tma_epi_atoms_and_tensors(

     return [assume_stride_divisibility(t) for t in tensors]
+def setup_epi_tensor(gemm, tensor, epi_tile=None, op_type="store", stage=None):
+    """Create copy metadata + smem layout for a supplemental epilogue tensor.
     Args:
+        gemm: The GEMM object (provides arch, epi_stage, and epilogue layout helpers).
+        tensor: The global memory tensor to set up for the epilogue.
         epi_tile: Epilogue tile shape. Defaults to gemm.epi_tile.
         op_type: "store" or "load".
     Returns:
+        (copy_atom, tensor, smem_layout_staged, epi_tile). copy_atom is None for pre-TMA archs.
     """
     if epi_tile is None:
         epi_tile = gemm.epi_tile
+    if stage is None:
+        stage = gemm.epi_stage
     dtype = tensor.element_type
     layout = cutlass.utils.LayoutEnum.from_tensor(tensor)
     utils_cls = sm100_utils if gemm.arch >= 100 else sm90_utils
+    smem_layout_staged = utils_cls.make_smem_layout_epi(dtype, layout, epi_tile, stage)
+    # Ragging-for-TMA is for varlen_m stores that need a per-batch row offset baked
+    # into the TMA descriptor. Loads don't currently support varlen_m, so skip the
+    # ragging conversion.
     tma_input = (
         copy_utils.create_ragged_tensor_for_tma(tensor, ragged_dim=0, ptr_shift=True)
+        if op_type != "load" and cute.rank(tensor) == 2
         else tensor
     )
     tma_atom, tma_tensor = gemm._make_tma_epi_atoms_and_tensors(

build/torch-cuda/quack/fast_math.py CHANGED Viewed

@@ -2,32 +2,78 @@
 import cutlass
 import cutlass.cute as cute
 from cutlass.base_dsl.typing import Integer
-from cutlass.cutlass_dsl import dsl_user_op
-class FastDivmod(cute.FastDivmodDivisor):
-    """We store the divisor along with the FastDivmodDivisor."""
-    @dsl_user_op
-    def __init__(
-        self,
-        divisor: Integer,
-        is_power_of_2: bool = None,
-        *,
-        loc=None,
-        ip=None,
-    ):
-        super().__init__(divisor, is_power_of_2=is_power_of_2, loc=loc, ip=ip)
         self.divisor = divisor
     def __extract_mlir_values__(self):
-        """Extract MLIR values for Host->Device transfer."""
-        return [self._divisor] + cutlass.extract_mlir_values(self.divisor)
     def __new_from_mlir_values__(self, values):
-        """Reconstruct FastDivmodDivisor from MLIR values."""
         new_obj = object.__new__(FastDivmod)
-        new_obj._divisor = values[0]
-        new_obj.divisor = cutlass.new_from_mlir_values(self.divisor, values[1:])
         return new_obj

 import cutlass
 import cutlass.cute as cute
+from cutlass import Int32, Uint32, Uint64
+from cutlass._mlir.dialects import llvm
 from cutlass.base_dsl.typing import Integer
+def ceil_log2(x: Integer) -> Int32:
+    """ceil(log2(x)) for 1 <= x < 2^31, as 32 - ctlz(x - 1). The llvm.intr.ctlz
+    lowers to lzcnt on the host (launch-prep path) and clz on device;
+    is_zero_poison=False makes ctlz(0) == 32, so x == 1 correctly yields 0."""
+    xm1 = (Int32(x) - 1).ir_value()
+    return Int32(32) - Int32(llvm.intr_ctlz(xm1, False))
+class FastDivmod:
+    """Magic-number unsigned divmod: q = umulhi(n, magic) >> shift, r = n - q * divisor.
+    The multiplier and shift are precomputed on the host (kernel params), so the
+    device-side divmod is 3 uniform-datapath ops (IMAD.WIDE.U32 + USHF + IMAD) plus
+    a select handling divisor == 1 (magic == 0 sentinel, like nvjet). This is the
+    lean form without the add-back correction the stock cute FastDivmodDivisor
+    emits: with shift = max(ceil_log2(d) - 1, 0) the multiplier ceil(2^(32+s)/d)
+    fits 32 bits (worst case d = 2^(c-1)+1 gives m <= 2^32 - 1) and the result is
+    exact for all dividends < 2^31, i.e. any non-negative Int32. Same contract and
+    algorithm as C++ cutlass::FastDivmod. Negative dividends are OUT of contract
+    (they reinterpret through Uint32 to values >= 2^31 and silently divide wrong);
+    use cute.FastDivmodDivisor if signed or full-u32 dividends are ever needed.
+    """
+    def __init__(self, divisor: Integer):
+        if isinstance(divisor, int):
+            assert 0 < divisor < 1 << 31
+            divisor = Int32(divisor)  # constants fold through the arithmetic below
         self.divisor = divisor
+        # Runs on the CPU at launch prep (host-side jit); the udiv is once per launch.
+        s = cutlass.max(ceil_log2(divisor) - 1, Int32(0))
+        pow_s = Uint64(Uint32(1) << Uint32(s))
+        numer = Uint64(0x100000000) * pow_s
+        magic = (numer + Uint64(Uint32(divisor)) - 1) // Uint64(Uint32(divisor))
+        self.magic = Uint32(magic & 0xFFFFFFFF)  # 0 when divisor == 1
+        self.shift = Uint32(s)
+    def __rdivmod__(self, dividend: Integer):
+        q = Uint32(cute.arch.mul_hi(Uint32(dividend), self.magic)) >> self.shift
+        # divisor == 1 sentinel: magic wrapped to 0, so q is 0; select the dividend
+        # instead. The remainder then self-corrects: r = n - n * 1 = 0.
+        q = Int32(cutlass.select_(self.magic == Uint32(0), Int32(dividend), Int32(q)))
+        r = Int32(dividend) - q * Int32(self.divisor)
+        return q, r
+    def __rfloordiv__(self, dividend: Integer) -> Int32:
+        q, _ = self.__rdivmod__(dividend)
+        return q
+    def __rmod__(self, dividend: Integer) -> Int32:
+        _, r = self.__rdivmod__(dividend)
+        return r
     def __extract_mlir_values__(self):
+        values = []
+        self._values_pos = []
+        for obj in [self.magic, self.shift, self.divisor]:
+            obj_values = cutlass.extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
     def __new_from_mlir_values__(self, values):
         new_obj = object.__new__(FastDivmod)
+        for name, n_items in zip(["magic", "shift", "divisor"], self._values_pos):
+            setattr(
+                new_obj, name, cutlass.new_from_mlir_values(getattr(self, name), values[:n_items])
+            )
+            values = values[n_items:]
+        new_obj._values_pos = self._values_pos
         return new_obj

build/torch-cuda/quack/gemm.py CHANGED Viewed

@@ -1,4 +1,4 @@
-# Copyright (c) 2025-2026, Tri Dao.
 # GEMM compilation via TVM-FFI with fake tensors and NamedTuple args.
 from typing import Optional
@@ -9,11 +9,12 @@ import cutlass.cute as cute
 from cutlass import Int32, Float32
 from cutlass.cute.runtime import make_ptr
-from .cache_utils import jit_cache
 from .compile_utils import make_fake_tensor as fake_tensor
 from .cute_dsl_utils import get_device_capacity, get_max_active_clusters, torch2cute_dtype_map
 from .gemm_default_epi import (
     GemmDefaultEpiMixin,
     GemmDefaultSm90,
     GemmDefaultSm100,
     GemmDefaultSm120,
@@ -28,7 +29,9 @@ from .gemm_tvm_ffi_utils import (
     make_fake_scheduler_args,
     make_fake_varlen_args,
     make_fake_gemm_tensors,
     compile_gemm_kernel,
 )
@@ -62,9 +65,12 @@ def _compile_gemm(
     device_capacity,
     rounding_mode,
     sr_seed_mode,
-    has_trace_ptr,
 ):
     sm_to_cls = {
         9: GemmDefaultSm90,
         10: GemmDefaultSm100,
         11: GemmDefaultSm100,
@@ -111,10 +117,18 @@ def _compile_gemm(
         sr_seed=fake_scalar(sr_seed_mode, dtype=Int32),
     )
     scheduler_args = make_fake_scheduler_args(
-        (is_dynamic_persistent and device_capacity[0] == 9), has_batch_idx_permute, l
     )
     aidx_len = m if varlen_m else (k if varlen_k else None)
     varlen_args = make_fake_varlen_args(varlen_m, varlen_k, gather_A, aidx_len)
     return compile_gemm_kernel(
         GemmCls,
         a_dtype,
@@ -132,9 +146,12 @@ def _compile_gemm(
         epi_args,
         scheduler_args,
         varlen_args,
-        has_trace_ptr=has_trace_ptr,
         use_tma_gather=use_tma_gather,
         concat_layout=concat_layout or None,
     )
@@ -149,6 +166,8 @@ def gemm(
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
     pingpong: bool = False,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
@@ -166,18 +185,24 @@ def gemm(
     sr_seed: int | Tensor = 0,
     use_tma_gather: bool = False,
     concat_layout: dict | None = None,
-    trace_ptr=None,  # Optional Int64 from TraceSession.ptr
 ) -> None:
     varlen_m = cu_seqlens_m is not None
     varlen_k = cu_seqlens_k is not None
     varlen = varlen_m or varlen_k
     gather_A = A_idx is not None
     assert not (varlen_m and varlen_k), "Only one of cu_seqlens_m and cu_seqlens_k"
     if gather_A:
         assert varlen, "gather_A requires varlen"
         assert cluster_N == 1, "gather_A requires cluster_N=1"
-    if varlen:
-        assert persistent, "varlen requires persistent=True"
     if add_to_output:
         assert not varlen_m, "Add to output not supported with varlen_m"
     if varlen_m:
@@ -188,15 +213,35 @@ def gemm(
         assert B.stride(-2) == 1, "varlen_k requires B to be n-major"
     device_capacity = get_device_capacity(A.device)
-    assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported"
     if use_tma_gather:
         assert device_capacity[0] in [10, 11], "TMA gather currently requires SM100/SM110"
     if rounding_mode == RoundingMode.RS:
         assert device_capacity[0] == 10, "Stochastic rounding (RoundingMode.RS) requires SM100"
-    if is_dynamic_persistent and device_capacity[0] == 9:
         assert tile_count_semaphore is not None, (
-            "Dynamic persistent tile scheduler in SM90 requires a semaphore in GMEM"
         )
     A_p, B_p, D_p, C_p = perm3d(A, B, D, C, varlen_m=varlen_m, varlen_k=varlen_k)
     a_major, b_major, d_major, c_major = get_majors(A_p, B_p, D_p, C_p)
@@ -210,6 +255,7 @@ def gemm(
     sr_seed_mode = (
         2 if isinstance(sr_seed, Tensor) else (1 if rounding_mode == RoundingMode.RS else 0)
     )
     compiled_fn = _compile_gemm(
         a_dtype,
         b_dtype,
@@ -219,8 +265,8 @@ def gemm(
         b_major,
         d_major,
         c_major,
-        (tile_M, tile_N),
-        (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
         is_dynamic_persistent,
@@ -239,14 +285,11 @@ def gemm(
         device_capacity,
         rounding_mode,
         sr_seed_mode,
-        trace_ptr is not None,
     )
-    from .cache_utils import COMPILE_ONLY
-    if COMPILE_ONLY:
-        return
     def scalar_arg(scalar, mode, dtype=Float32):
         if mode == 0:
             return None
@@ -255,7 +298,10 @@ def gemm(
         else:
             return scalar.data_ptr()
-    max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     epi_args = GemmDefaultEpiMixin.EpilogueArguments(
         alpha=scalar_arg(alpha, alpha_mode),
@@ -269,14 +315,15 @@ def gemm(
     scheduler_args = make_scheduler_args(
         max_active_clusters,
         max_swizzle_size,
-        tile_count_semaphore,
         batch_idx_permute,
     )
     varlen_args = make_varlen_args(cu_seqlens_m, cu_seqlens_k, A_idx)
     if device_capacity[0] in [10, 11]:
-        compiled_fn(
-            A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, trace_ptr
-        )
     else:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, trace_ptr)

+# Copyright (c) 2025-2026, QuACK team.
 # GEMM compilation via TVM-FFI with fake tensors and NamedTuple args.
 from typing import Optional
 from cutlass import Int32, Float32
 from cutlass.cute.runtime import make_ptr
+from .cache import jit_cache
 from .compile_utils import make_fake_tensor as fake_tensor
 from .cute_dsl_utils import get_device_capacity, get_max_active_clusters, torch2cute_dtype_map
 from .gemm_default_epi import (
     GemmDefaultEpiMixin,
+    GemmDefaultSm80,
     GemmDefaultSm90,
     GemmDefaultSm100,
     GemmDefaultSm120,
     make_fake_scheduler_args,
     make_fake_varlen_args,
     make_fake_gemm_tensors,
+    make_fake_sf_tensor,
     compile_gemm_kernel,
+    validate_blockscaled_sf,
 )
     device_capacity,
     rounding_mode,
     sr_seed_mode,
+    num_warps,
+    sf_dtype=None,
+    sf_vec_size=None,
 ):
     sm_to_cls = {
+        8: GemmDefaultSm80,
         9: GemmDefaultSm90,
         10: GemmDefaultSm100,
         11: GemmDefaultSm100,
         sr_seed=fake_scalar(sr_seed_mode, dtype=Int32),
     )
     scheduler_args = make_fake_scheduler_args(
+        (is_dynamic_persistent and device_capacity[0] <= 9), has_batch_idx_permute, l
     )
     aidx_len = m if varlen_m else (k if varlen_k else None)
     varlen_args = make_fake_varlen_args(varlen_m, varlen_k, gather_A, aidx_len)
+    if sf_dtype is not None:
+        # Padded SF buffers have a static batch dim of exactly 1 (not l): SFA for
+        # varlen_m (M-padded) and varlen_k (K-padded); SFB is K-padded too for
+        # varlen_k but stays per-batch (l, rn, rk, ...) for varlen_m.
+        mSFA = make_fake_sf_tensor(sf_dtype, 1 if (varlen_m or varlen_k) else l)
+        mSFB = make_fake_sf_tensor(sf_dtype, 1 if varlen_k else l)
+    else:
+        mSFA, mSFB = None, None
     return compile_gemm_kernel(
         GemmCls,
         a_dtype,
         epi_args,
         scheduler_args,
         varlen_args,
+        mSFA=mSFA,
+        mSFB=mSFB,
         use_tma_gather=use_tma_gather,
         concat_layout=concat_layout or None,
+        num_warps=num_warps,
+        sf_vec_size=sf_vec_size,
     )
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
+    cluster_K: int = 1,
+    tile_K: int | None = None,
     pingpong: bool = False,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
     sr_seed: int | Tensor = 0,
     use_tma_gather: bool = False,
     concat_layout: dict | None = None,
+    num_warps: Optional[int] = None,
+    # SFA/SFB: (l, rm/rn, rk, 32, 4, 4) blocked scale factors. For varlen_m, SFA is
+    # M-padded (1, total_padded_rm, rk, 32, 4, 4) while SFB stays per-batch. For
+    # varlen_k, BOTH are K-padded (1, rm/rn, total_padded_rk, 32, 4, 4); pad bytes
+    # may be arbitrary (the kernel skips the MMA instructions covering them).
+    # See AI/varlen_blockscaled_sf_layout.md.
+    SFA: Optional[Tensor] = None,
+    SFB: Optional[Tensor] = None,
 ) -> None:
     varlen_m = cu_seqlens_m is not None
     varlen_k = cu_seqlens_k is not None
     varlen = varlen_m or varlen_k
     gather_A = A_idx is not None
+    blockscaled = SFA is not None
     assert not (varlen_m and varlen_k), "Only one of cu_seqlens_m and cu_seqlens_k"
     if gather_A:
         assert varlen, "gather_A requires varlen"
         assert cluster_N == 1, "gather_A requires cluster_N=1"
     if add_to_output:
         assert not varlen_m, "Add to output not supported with varlen_m"
     if varlen_m:
         assert B.stride(-2) == 1, "varlen_k requires B to be n-major"
     device_capacity = get_device_capacity(A.device)
+    assert device_capacity[0] in [8, 9, 10, 11, 12], (
+        "Only SM8x, SM90, SM100, SM110, and SM120 are supported"
+    )
+    sf_dtype, sf_vec_size = None, None
+    if blockscaled:
+        assert not gather_A, "Blockscaled GEMM does not support gather_A yet"
+        assert not concat_layout, "Blockscaled GEMM does not support concat_layout"
+        assert tile_K is None, "Blockscaled GEMM derives tile_K from the MMA instruction"
+        if varlen_m:
+            num_batches = cu_seqlens_m.shape[0] - 1
+        elif varlen_k:
+            num_batches = cu_seqlens_k.shape[0] - 1
+        else:
+            num_batches = None
+        sf_dtype, sf_vec_size = validate_blockscaled_sf(
+            A, B, SFA, SFB, device_capacity, num_batches=num_batches, varlen_k=varlen_k
+        )
     if use_tma_gather:
         assert device_capacity[0] in [10, 11], "TMA gather currently requires SM100/SM110"
     if rounding_mode == RoundingMode.RS:
         assert device_capacity[0] == 10, "Stochastic rounding (RoundingMode.RS) requires SM100"
+    if is_dynamic_persistent and device_capacity[0] <= 9:
         assert tile_count_semaphore is not None, (
+            "Dynamic persistent tile scheduler for SM8x and SM90 requires a semaphore in GMEM"
         )
+    if device_capacity[0] == 8:
+        if add_to_output:
+            C = D
+            add_to_output = False
     A_p, B_p, D_p, C_p = perm3d(A, B, D, C, varlen_m=varlen_m, varlen_k=varlen_k)
     a_major, b_major, d_major, c_major = get_majors(A_p, B_p, D_p, C_p)
     sr_seed_mode = (
         2 if isinstance(sr_seed, Tensor) else (1 if rounding_mode == RoundingMode.RS else 0)
     )
+    tile_shape_mnk = (tile_M, tile_N) if tile_K is None else (tile_M, tile_N, tile_K)
     compiled_fn = _compile_gemm(
         a_dtype,
         b_dtype,
         b_major,
         d_major,
         c_major,
+        tile_shape_mnk,
+        (cluster_M, cluster_N, cluster_K),
         pingpong,
         persistent,
         is_dynamic_persistent,
         device_capacity,
         rounding_mode,
         sr_seed_mode,
+        num_warps,
+        sf_dtype,
+        sf_vec_size,
     )
     def scalar_arg(scalar, mode, dtype=Float32):
         if mode == 0:
             return None
         else:
             return scalar.data_ptr()
+    cluster_size = cluster_M * cluster_N * cluster_K
+    max_active_clusters = (
+        get_max_active_clusters(cluster_size, device_capacity=device_capacity) if persistent else 0
+    )
     epi_args = GemmDefaultEpiMixin.EpilogueArguments(
         alpha=scalar_arg(alpha, alpha_mode),
     scheduler_args = make_scheduler_args(
         max_active_clusters,
         max_swizzle_size,
+        # Must mirror make_fake_scheduler_args in _compile_gemm: only the SM8x/SM90
+        # dynamic scheduler consumes the semaphore; SM100 uses CLC instead, and the
+        # compiled signature has None there.
+        tile_count_semaphore if (is_dynamic_persistent and device_capacity[0] <= 9) else None,
         batch_idx_permute,
     )
     varlen_args = make_varlen_args(cu_seqlens_m, cu_seqlens_k, A_idx)
     if device_capacity[0] in [10, 11]:
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, SFA, SFB)
     else:
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args)

build/torch-cuda/quack/gemm_act.py CHANGED Viewed

@@ -1,26 +1,27 @@
 # Copyright (c) 2025, Wentao Guo, Tri Dao.
 from __future__ import annotations
-from typing import NamedTuple, Tuple, Optional, Callable
-from functools import partial
 from torch import Tensor
 import cutlass
 import cutlass.cute as cute
-import cutlass.utils.hopper_helpers as sm90_utils_og
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass import Int32, Float32, const_expr
 from cutlass.cute.runtime import make_ptr
 from .compile_utils import make_fake_tensor as fake_tensor
 from .cute_dsl_utils import (
-    ParamsBase,
     mlir_namedtuple,
     get_device_capacity,
     get_max_active_clusters,
     torch2cute_dtype_map,
 )
-from .epi_ops import TileStore
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
@@ -34,23 +35,32 @@ from .gemm_tvm_ffi_utils import (
     make_fake_varlen_args,
     div_for_dtype,
     make_fake_gemm_tensors,
     compile_gemm_kernel,
 )
-from .cache_utils import jit_cache
 from . import layout_utils as layout_utils
 from .layout_utils import permute_gated_Cregs_b16
 from .activation import act_fn_map, gate_fn_map
-from .rounding import RoundingMode
-class GemmActMixin(GemmDefaultEpiMixin):
-    _epi_ops = (*GemmDefaultEpiMixin._epi_ops, TileStore("mPostAct"))
     _extra_param_fields = (("act_fn", cutlass.Constexpr, None),)
-    _epi_param_bases = (ParamsBase,)
     @mlir_namedtuple
     class EpilogueArguments(NamedTuple):
-        mPostAct: cute.Tensor
         act_fn: cutlass.Constexpr[Optional[Callable]] = None
         alpha: Optional[Float32 | cute.Tensor] = None
         beta: Optional[Float32 | cute.Tensor] = None
@@ -63,20 +73,39 @@ class GemmActMixin(GemmDefaultEpiMixin):
     def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None):
         self.rounding_mode = args.rounding_mode
-        self.postact_dtype = args.mPostAct.element_type
-        self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct)
-        self.cta_tile_shape_postact_mn = self.cta_tile_shape_mnk[:2]
         d = self._epi_ops_to_params_dict(args)
         d["act_fn"] = args.act_fn
         for key in ("mRowVecBroadcast", "mColVecBroadcast"):
-            if key in self.concat_layout and key in d and d[key] is not None:
                 d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
-    # epi_get_tma_atoms, epi_smem_bytes_per_stage, epi_get_smem_struct,
     # epi_get_smem_tensors are all inherited from ComposableEpiMixin via _epi_ops.
-    def epi_setup_postact(
         self,
         params,
         epi_smem_tensors,
@@ -86,61 +115,56 @@ class GemmActMixin(GemmDefaultEpiMixin):
         varlen_manager,
         tidx,
     ):
-        """Setup postact TMA copies and partitions before the epilogue loop."""
-        sPostAct = epi_smem_tensors[self._epi_smem_map["mPostAct"]]
-        get_smem_store_op = (
-            partial(sm100_utils.get_smem_store_op, tiled_tmem_load=tiled_copy_t2r)
-            if self.arch == 100
-            else sm90_utils_og.sm90_get_smem_store_op
-        )
-        copy_atom_postact_r2s = get_smem_store_op(
-            self.postact_layout, self.postact_dtype, self.acc_dtype
         )
-        tiled_copy_postact_r2s = cute.make_tiled_copy_S(copy_atom_postact_r2s, tiled_copy_r2s)
-        tRS_sPostAct = tiled_copy_postact_r2s.get_slice(tidx).partition_D(sPostAct)
         batch_idx = tile_coord_mnkl[3]
-        copy_postact, _, _ = self.epilog_gmem_copy_and_partition(
-            params.tma_atom_mPostAct,
-            varlen_manager.offset_batch_epi(params.mPostAct, batch_idx),
-            self.cta_tile_shape_postact_mn,
-            params.epi_tile_mPostAct,
-            sPostAct,
             tile_coord_mnkl,
         )
-        return tiled_copy_postact_r2s, tRS_sPostAct, copy_postact
     @cute.jit
-    def epi_convert_postact(
-        self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx
     ):
-        """Convert postact from acc_dtype to postact_dtype. Override for custom postprocessing."""
         if const_expr(
             self.rounding_mode == RoundingMode.RS
-            and tRS_rPostAct.element_type == cutlass.Float32
-            and self.postact_dtype == cutlass.BFloat16
         ):
-            from .rounding import convert_f32_to_bf16_sr
             from cutlass.cute.tensor import TensorSSA
-            # Salt with 0x9E3779B1 to avoid sharing entropy with the D output seed
-            seed = (
-                sr_seed
-                + 0x9E3779B1
-                + (
-                    tile_coord_mnkl[0] * 65537
-                    + tile_coord_mnkl[1] * 257
-                    + tile_coord_mnkl[3] * 17
-                    + (num_prev_subtiles + epi_idx) * 7
-                )
-            )
-            tRS_rPostAct_out = cute.make_rmem_tensor_like(tRS_rPostAct, self.postact_dtype)
-            src_vec = tRS_rPostAct.load()
             raw_vec = convert_f32_to_bf16_sr(src_vec, seed, tidx)
-            tRS_rPostAct_out.store(TensorSSA(raw_vec, src_vec.shape, self.postact_dtype))
         else:
-            tRS_rPostAct_out = cute.make_rmem_tensor_like(tRS_rPostAct, self.postact_dtype)
-            tRS_rPostAct_out.store(tRS_rPostAct.load().to(self.postact_dtype))
-        return tRS_rPostAct_out
     @cute.jit
     def epi_visit_subtile(
@@ -149,29 +173,28 @@ class GemmActMixin(GemmDefaultEpiMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
         GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC)
         # Apply activation function if provided
         # If we don't have .shape here, the compiler generates local stores and loads
         if const_expr(params.act_fn is not None):
-            tRS_rPostAct = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype)
-            if const_expr(self.arch < 100):
-                for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
-                    tRS_rPostAct[i] = params.act_fn(tRS_rD[i])
-            else:
-                for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True):
-                    tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn(
-                        (tRS_rD[2 * i], tRS_rD[2 * i + 1])
-                    )
         else:
-            tRS_rPostAct = tRS_rD
-        return tRS_rPostAct
 class GemmActSm90(GemmActMixin, GemmSm90):
     pass
 class GemmActSm100(GemmActMixin, GemmSm100):
     pass
@@ -189,33 +212,37 @@ def _gated_epi_tile_fn(gemm, epi_tile):
 class GemmGatedMixin(GemmActMixin):
     _epi_ops = (
-        *GemmDefaultEpiMixin._epi_ops,
-        TileStore("mPostAct", epi_tile_fn=_gated_epi_tile_fn),
     )
     def epi_to_underlying_arguments(
         self, args: GemmActMixin.EpilogueArguments, *, loc=None, ip=None
     ) -> GemmActMixin.EpilogueParams:
-        assert args.mPostAct.element_type.width == 16, (
             "GemmGated only supports 16bit postact for now"
         )
         assert self.d_layout is None or self.d_layout.is_n_major_c()
-        assert cutlass.utils.LayoutEnum.from_tensor(args.mPostAct).is_n_major_c()
         if self.arch == 90:
             assert self.cta_tile_shape_mnk[1] % 32 == 0, (
                 "GemmGatedSm90 requires tileN to be divisible by 32"
             )
         self.rounding_mode = args.rounding_mode
-        self.postact_dtype = args.mPostAct.element_type
-        self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct)
-        self.cta_tile_shape_postact_mn = (
             self.cta_tile_shape_mnk[0],
             self.cta_tile_shape_mnk[1] // 2,
         )
         d = self._epi_ops_to_params_dict(args)
         d["act_fn"] = args.act_fn
         for key in ("mRowVecBroadcast", "mColVecBroadcast"):
-            if key in self.concat_layout and key in d and d[key] is not None:
                 d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
@@ -226,43 +253,96 @@ class GemmGatedMixin(GemmActMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
         GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC)
-        tRS_rPostAct_layout = cute.recast_layout(2, 1, tRS_rD.layout)
         # If we don't have .shape here, the compiler generates local stores and loads
-        tRS_rPostAct = cute.make_rmem_tensor(tRS_rPostAct_layout.shape, self.acc_dtype)
-        if const_expr(self.arch < 100):
-            for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
-                tRS_rPostAct[i] = params.act_fn(tRS_rD[2 * i], tRS_rD[2 * i + 1])
-        else:
-            for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True):
-                tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn(
-                    (tRS_rD[4 * i], tRS_rD[4 * i + 2]), (tRS_rD[4 * i + 1], tRS_rD[4 * i + 3])
-                )
-        return tRS_rPostAct
     @cute.jit
-    def epi_convert_postact(
-        self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx
     ):
-        tRS_rPostAct_out = GemmActMixin.epi_convert_postact(
-            self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx
         )
-        if const_expr(self.arch == 90):
             # Only need this if we're using STSM
-            permute_gated_Cregs_b16(tRS_rPostAct_out)
-        return tRS_rPostAct_out
 class GemmGatedSm90(GemmGatedMixin, GemmSm90):
     pass
 class GemmGatedSm100(GemmGatedMixin, GemmSm100):
     pass
-class GemmGatedSm120(GemmGatedMixin, GemmSm120):
     pass
@@ -295,13 +375,25 @@ def _compile_gemm_act(
     rounding_mode=RoundingMode.RN,
     sr_seed_mode=0,
     use_tma_gather=False,
 ):
     sm_to_cls = {
-        "act": {9: GemmActSm90, 10: GemmActSm100, 11: GemmActSm100, 12: GemmActSm120},
-        "gated": {9: GemmGatedSm90, 10: GemmGatedSm100, 11: GemmGatedSm100, 12: GemmGatedSm120},
     }
-    if device_capacity[0] == 12 and gemm_cls_name == "act":
-        raise NotImplementedError("SM120 non-gated activation GEMM epilogue is not yet supported")
     GemmCls = sm_to_cls[gemm_cls_name][device_capacity[0]]
     pa_leading = 1 if postact_major == "n" else 0
     mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors(
@@ -320,7 +412,7 @@ def _compile_gemm_act(
     div_pa = div_for_dtype(postact_dtype)
     pa_leading_dim = 1 if gemm_cls_name == "gated" else pa_leading
     pa_shape = (m, pa_n) if varlen_m else (m, pa_n, l)
-    mPostAct = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa)
     mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4)
     if colvec_ndim == 2:
@@ -341,7 +433,7 @@ def _compile_gemm_act(
             return make_ptr(dtype, 0, cute.AddressSpace.gmem, assumed_align=4)
     epi_args = GemmCls.EpilogueArguments(
-        mPostAct,
         act_fn,
         mRowVecBroadcast=mRowVec,
         mColVecBroadcast=mColVec,
@@ -352,6 +444,11 @@ def _compile_gemm_act(
         (is_dynamic_persistent and device_capacity[0] == 9), False, l
     )
     varlen_args = make_fake_varlen_args(varlen_m, False, gather_A, m if varlen_m else None)
     return compile_gemm_kernel(
         GemmCls,
         a_dtype,
@@ -369,8 +466,11 @@ def _compile_gemm_act(
         epi_args,
         scheduler_args,
         varlen_args,
         use_tma_gather=use_tma_gather,
         concat_layout=concat_layout or None,
     )
@@ -386,6 +486,7 @@ def gemm_act(
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
     pingpong: bool = False,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
@@ -398,6 +499,8 @@ def gemm_act(
     sr_seed: int | Tensor = 0,
     use_tma_gather: bool = False,
     concat_layout: tuple | None = None,
 ) -> None:
     if activation in gate_fn_map:
         gemm_cls_name = "gated"
@@ -407,6 +510,7 @@ def gemm_act(
     varlen_m = cu_seqlens_m is not None
     gather_A = A_idx is not None
     if varlen_m:
         assert persistent, "varlen_m requires persistent=True"
         assert A.stride(-1) == 1, "varlen_m requires A to be k-major"
@@ -437,7 +541,16 @@ def gemm_act(
     colvec_ndim = colvec_bias.ndim if colvec_bias is not None else 0
     device_capacity = get_device_capacity(A.device)
-    assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported"
     if rounding_mode == RoundingMode.RS:
         assert device_capacity[0] == 10, "Stochastic rounding (RoundingMode.RS) requires SM100"
@@ -461,7 +574,7 @@ def gemm_act(
         d_major,
         c_major,
         postact_major,
-        (tile_M, tile_N),
         (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
@@ -478,13 +591,10 @@ def gemm_act(
         rounding_mode=rounding_mode,
         sr_seed_mode=sr_seed_mode,
         use_tma_gather=use_tma_gather,
     )
-    from .cache_utils import COMPILE_ONLY
-    if COMPILE_ONLY:
-        return
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     def scalar_arg(scalar, mode, dtype=Int32):
@@ -511,9 +621,9 @@ def gemm_act(
     varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx)
     if device_capacity[0] in [10, 11]:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, None)
     else:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None)
 gemm_gated = gemm_act

 # Copyright (c) 2025, Wentao Guo, Tri Dao.
 from __future__ import annotations
+import math
+from typing import NamedTuple, Tuple, Optional, Callable, Type
 from torch import Tensor
 import cutlass
 import cutlass.cute as cute
 import cutlass.utils.blackwell_helpers as sm100_utils
 from cutlass import Int32, Float32, const_expr
 from cutlass.cute.runtime import make_ptr
+from cutlass.cute.nvgpu import warp
 from .compile_utils import make_fake_tensor as fake_tensor
 from .cute_dsl_utils import (
     mlir_namedtuple,
     get_device_capacity,
     get_max_active_clusters,
     torch2cute_dtype_map,
 )
+from .epi_composable import ComposableEpiMixin
+from .epi_ops import ColVecLoad, RowVecLoad, Scalar, TileStore
+from .gemm_sm80 import GemmSm80
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
     make_fake_varlen_args,
     div_for_dtype,
     make_fake_gemm_tensors,
+    make_fake_sf_tensor,
     compile_gemm_kernel,
+    validate_blockscaled_sf,
 )
+from .cache import jit_cache
 from . import layout_utils as layout_utils
+from . import copy_utils as copy_utils
 from .layout_utils import permute_gated_Cregs_b16
 from .activation import act_fn_map, gate_fn_map
+from .rounding import RoundingMode, convert_f32_to_bf16_sr, epilogue_aux_out_sr_seed
+class GemmActMixin(ComposableEpiMixin):
+    _epi_ops = (
+        Scalar("alpha"),
+        Scalar("beta"),
+        Scalar("sr_seed", dtype=Int32),
+        RowVecLoad("mRowVecBroadcast"),
+        ColVecLoad("mColVecBroadcast"),
+        TileStore("mAuxOut"),
+    )
     _extra_param_fields = (("act_fn", cutlass.Constexpr, None),)
     @mlir_namedtuple
     class EpilogueArguments(NamedTuple):
+        mAuxOut: cute.Tensor
         act_fn: cutlass.Constexpr[Optional[Callable]] = None
         alpha: Optional[Float32 | cute.Tensor] = None
         beta: Optional[Float32 | cute.Tensor] = None
     def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=None):
         self.rounding_mode = args.rounding_mode
+        self.aux_out_dtype = args.mAuxOut.element_type
+        self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut)
+        self.cta_tile_shape_aux_out_mn = self.cta_tile_shape_mnk[:2]
         d = self._epi_ops_to_params_dict(args)
         d["act_fn"] = args.act_fn
         for key in ("mRowVecBroadcast", "mColVecBroadcast"):
+            if key in self.concat_layout and key in d:
                 d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
+    # epi_get_tma_atoms, epi_smem_bytes, epi_get_smem_struct,
     # epi_get_smem_tensors are all inherited from ComposableEpiMixin via _epi_ops.
+    def epi_make_aux_out_copy_atom_r2s(self, params, tiled_copy_t2r):
+        """Build the register-to-shared copy atom used by aux outputs."""
+        if self.arch == 100:
+            return sm100_utils.get_smem_store_op(
+                self.aux_out_layout, self.aux_out_dtype, self.acc_dtype, tiled_copy_t2r
+            )
+        else:
+            return copy_utils.get_smem_store_atom(
+                self.aux_out_dtype,
+                transpose=self.aux_out_layout != cutlass.utils.LayoutEnum.ROW_MAJOR,
+                major_mode_size=cute.size(params.epi_tile_mAuxOut, mode=[1])
+                // self.atom_layout_mnk[1],
+            )
+    def epi_make_aux_out_tiled_copy_r2s(self, params, tiled_copy_r2s, tiled_copy_t2r):
+        """Build the register-to-shared tiled copy used by aux outputs."""
+        copy_atom_aux_out_r2s = self.epi_make_aux_out_copy_atom_r2s(params, tiled_copy_t2r)
+        return cute.make_tiled_copy_S(copy_atom_aux_out_r2s, tiled_copy_r2s)
+    def epi_setup_aux_out(
         self,
         params,
         epi_smem_tensors,
         varlen_manager,
         tidx,
     ):
+        """Setup aux output TMA copies and partitions before the epilogue loop.
+        Returns an empty tuple when mAuxOut wasn't supplied so the framework
+        skips the aux-out path.
+        """
+        if getattr(params, "mAuxOut", None) is None:
+            return ()
+        sAuxOut = epi_smem_tensors["mAuxOut"]
+        tiled_copy_aux_out_r2s = self.epi_make_aux_out_tiled_copy_r2s(
+            params, tiled_copy_r2s, tiled_copy_t2r
         )
+        tRS_sAuxOut = tiled_copy_aux_out_r2s.get_slice(tidx).partition_D(sAuxOut)
         batch_idx = tile_coord_mnkl[3]
+        copy_aux_out, _, _ = self.epilog_gmem_copy_and_partition(
+            params.tma_atom_mAuxOut,
+            varlen_manager.offset_batch_epi(params.mAuxOut, batch_idx),
+            self.cta_tile_shape_aux_out_mn,
+            params.epi_tile_mAuxOut,
+            sAuxOut,
             tile_coord_mnkl,
         )
+        return ((tiled_copy_aux_out_r2s, tRS_sAuxOut, copy_aux_out),)
     @cute.jit
+    def epi_convert_aux_out(
+        self,
+        output_idx: cutlass.Constexpr[int],
+        tRS_rAuxOut,
+        sr_seed,
+        tidx,
+        tile_coord_mnkl,
+        num_prev_subtiles,
+        epi_idx,
     ):
+        """Convert aux output from acc_dtype to aux_out_dtype. Override for custom postprocessing."""
         if const_expr(
             self.rounding_mode == RoundingMode.RS
+            and tRS_rAuxOut.element_type == cutlass.Float32
+            and self.aux_out_dtype == cutlass.BFloat16
         ):
             from cutlass.cute.tensor import TensorSSA
+            seed = epilogue_aux_out_sr_seed(sr_seed, tile_coord_mnkl, num_prev_subtiles + epi_idx)
+            tRS_rAuxOut_out = cute.make_rmem_tensor_like(tRS_rAuxOut, self.aux_out_dtype)
+            src_vec = tRS_rAuxOut.load()
             raw_vec = convert_f32_to_bf16_sr(src_vec, seed, tidx)
+            tRS_rAuxOut_out.store(TensorSSA(raw_vec, src_vec.shape, self.aux_out_dtype))
         else:
+            tRS_rAuxOut_out = tRS_rAuxOut.to(self.aux_out_dtype)
+        return tRS_rAuxOut_out
     @cute.jit
     def epi_visit_subtile(
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
         GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC)
         # Apply activation function if provided
         # If we don't have .shape here, the compiler generates local stores and loads
         if const_expr(params.act_fn is not None):
+            tRS_rAuxOut = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype)
+            vectorize = const_expr(self.arch == 100)
+            for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize):
+                tRS_rAuxOut[i] = params.act_fn(tRS_rD[i])
         else:
+            tRS_rAuxOut = tRS_rD
+        return (tRS_rAuxOut,)
 class GemmActSm90(GemmActMixin, GemmSm90):
     pass
+class GemmActSm80(GemmActMixin, GemmSm80):
+    pass
 class GemmActSm100(GemmActMixin, GemmSm100):
     pass
 class GemmGatedMixin(GemmActMixin):
     _epi_ops = (
+        Scalar("alpha"),
+        Scalar("beta"),
+        Scalar("sr_seed", dtype=Int32),
+        RowVecLoad("mRowVecBroadcast"),
+        ColVecLoad("mColVecBroadcast"),
+        TileStore("mAuxOut", epi_tile_fn=_gated_epi_tile_fn),
     )
     def epi_to_underlying_arguments(
         self, args: GemmActMixin.EpilogueArguments, *, loc=None, ip=None
     ) -> GemmActMixin.EpilogueParams:
+        assert args.mAuxOut.element_type.width == 16, (
             "GemmGated only supports 16bit postact for now"
         )
         assert self.d_layout is None or self.d_layout.is_n_major_c()
+        assert cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut).is_n_major_c()
         if self.arch == 90:
             assert self.cta_tile_shape_mnk[1] % 32 == 0, (
                 "GemmGatedSm90 requires tileN to be divisible by 32"
             )
         self.rounding_mode = args.rounding_mode
+        self.aux_out_dtype = args.mAuxOut.element_type
+        self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut)
+        self.cta_tile_shape_aux_out_mn = (
             self.cta_tile_shape_mnk[0],
             self.cta_tile_shape_mnk[1] // 2,
         )
         d = self._epi_ops_to_params_dict(args)
         d["act_fn"] = args.act_fn
         for key in ("mRowVecBroadcast", "mColVecBroadcast"):
+            if key in self.concat_layout and key in d:
                 d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
         GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC)
+        tRS_rAuxOut_layout = cute.recast_layout(2, 1, tRS_rD.layout)
         # If we don't have .shape here, the compiler generates local stores and loads
+        tRS_rAuxOut = cute.make_rmem_tensor(tRS_rAuxOut_layout.shape, self.acc_dtype)
+        tRS_rD_pair = cute.flat_divide(tRS_rD, cute.make_layout(2))
+        tRS_rGate = tRS_rD_pair[0, ...]
+        tRS_rUp = tRS_rD_pair[1, ...]
+        vectorize = const_expr(self.arch == 100)
+        for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize):
+            tRS_rAuxOut[i] = params.act_fn(tRS_rGate[i], tRS_rUp[i])
+        return (tRS_rAuxOut,)
     @cute.jit
+    def epi_convert_aux_out(
+        self,
+        output_idx: cutlass.Constexpr[int],
+        tRS_rAuxOut,
+        sr_seed,
+        tidx,
+        tile_coord_mnkl,
+        num_prev_subtiles,
+        epi_idx,
     ):
+        tRS_rAuxOut_out = GemmActMixin.epi_convert_aux_out(
+            self,
+            output_idx,
+            tRS_rAuxOut,
+            sr_seed,
+            tidx,
+            tile_coord_mnkl,
+            num_prev_subtiles,
+            epi_idx,
         )
+        if const_expr(self.arch in (90, 120)):
             # Only need this if we're using STSM
+            permute_gated_Cregs_b16(tRS_rAuxOut_out)
+        return tRS_rAuxOut_out
 class GemmGatedSm90(GemmGatedMixin, GemmSm90):
     pass
+class GemmGatedSm80(GemmGatedMixin, GemmSm80):
+    pass
 class GemmGatedSm100(GemmGatedMixin, GemmSm100):
     pass
+class GemmGatedSm120Mixin:
+    @staticmethod
+    def _compute_tile_shape_or_override(
+        cta_tile_shape_mnk: Tuple[int, int, int],
+        atom_layout_mnk: Tuple[int, int, int],
+        element_type: Optional[Type[cutlass.Numeric]] = None,
+        epi_tile_override: Tuple[int, int] | None = None,
+    ) -> Tuple[int, int]:
+        if epi_tile_override is not None:
+            return epi_tile_override
+        # Typically epi_tile is (64, 32) but since we want tile_n = 64 (see below), we might set
+        # tile_m = 32 if there's only 2 warps along the M direction.
+        tile_m = math.gcd(atom_layout_mnk[0] * 16, cute.size(cta_tile_shape_mnk, mode=[0]))
+        atom_n = atom_layout_mnk[1]
+        # E.g. if we have 2 warps along N direction, we want each warp to have 32 elems so that
+        # postact has 16 elements, which means tile_n should be 64.
+        tile_n = math.gcd(atom_n * 8 * 4, cute.size(cta_tile_shape_mnk, mode=[1]))
+        return (tile_m, tile_n)
+    def epi_make_aux_out_tiled_copy_r2s(self, params, tiled_copy_r2s, tiled_copy_t2r):
+        copy_atom_aux_out_r2s = self.epi_make_aux_out_copy_atom_r2s(params, tiled_copy_t2r)
+        copy_atom_postact_c = self.epi_make_aux_out_copy_atom_r2s(params, cutlass.Float16)
+        op = warp.MmaF16BF16Op(self.a_dtype, self.acc_dtype, self.mma_inst_mnk)
+        tC = cute.make_layout(self.atom_layout_mnk)
+        atom_m, atom_n, atom_k = self.atom_layout_mnk
+        permutation_mnk = (
+            self.mma_inst_mnk[0] * atom_m,
+            self.mma_inst_mnk[1] * atom_n * 2,
+            self.mma_inst_mnk[2] * atom_k,
+        )
+        tiled_mma_gated_postact = cute.make_tiled_mma(op, tC, permutation_mnk=permutation_mnk)
+        tiled_copy_aux_out_c_atom = cute.make_tiled_copy_C_atom(
+            copy_atom_postact_c, tiled_mma_gated_postact
+        )
+        return cute.make_tiled_copy_S(copy_atom_aux_out_r2s, tiled_copy_aux_out_c_atom)
+class GemmGatedSm120(GemmGatedSm120Mixin, GemmGatedMixin, GemmSm120):
     pass
     rounding_mode=RoundingMode.RN,
     sr_seed_mode=0,
     use_tma_gather=False,
+    sf_dtype=None,
+    sf_vec_size=None,
 ):
     sm_to_cls = {
+        "act": {
+            8: GemmActSm80,
+            9: GemmActSm90,
+            10: GemmActSm100,
+            11: GemmActSm100,
+            12: GemmActSm120,
+        },
+        "gated": {
+            8: GemmGatedSm80,
+            9: GemmGatedSm90,
+            10: GemmGatedSm100,
+            11: GemmGatedSm100,
+            12: GemmGatedSm120,
+        },
     }
     GemmCls = sm_to_cls[gemm_cls_name][device_capacity[0]]
     pa_leading = 1 if postact_major == "n" else 0
     mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors(
     div_pa = div_for_dtype(postact_dtype)
     pa_leading_dim = 1 if gemm_cls_name == "gated" else pa_leading
     pa_shape = (m, pa_n) if varlen_m else (m, pa_n, l)
+    mAuxOut = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa)
     mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4)
     if colvec_ndim == 2:
             return make_ptr(dtype, 0, cute.AddressSpace.gmem, assumed_align=4)
     epi_args = GemmCls.EpilogueArguments(
+        mAuxOut,
         act_fn,
         mRowVecBroadcast=mRowVec,
         mColVecBroadcast=mColVec,
         (is_dynamic_persistent and device_capacity[0] == 9), False, l
     )
     varlen_args = make_fake_varlen_args(varlen_m, False, gather_A, m if varlen_m else None)
+    if sf_dtype is not None:
+        mSFA = make_fake_sf_tensor(sf_dtype, l)
+        mSFB = make_fake_sf_tensor(sf_dtype, l)
+    else:
+        mSFA, mSFB = None, None
     return compile_gemm_kernel(
         GemmCls,
         a_dtype,
         epi_args,
         scheduler_args,
         varlen_args,
+        mSFA=mSFA,
+        mSFB=mSFB,
         use_tma_gather=use_tma_gather,
         concat_layout=concat_layout or None,
+        sf_vec_size=sf_vec_size,
     )
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
+    tile_K: int | None = None,
     pingpong: bool = False,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
     sr_seed: int | Tensor = 0,
     use_tma_gather: bool = False,
     concat_layout: tuple | None = None,
+    SFA: Optional[Tensor] = None,  # (l, rm, rk, 32, 4, 4) blocked scale factors
+    SFB: Optional[Tensor] = None,  # (l, rn, rk, 32, 4, 4)
 ) -> None:
     if activation in gate_fn_map:
         gemm_cls_name = "gated"
     varlen_m = cu_seqlens_m is not None
     gather_A = A_idx is not None
+    blockscaled = SFA is not None
     if varlen_m:
         assert persistent, "varlen_m requires persistent=True"
         assert A.stride(-1) == 1, "varlen_m requires A to be k-major"
     colvec_ndim = colvec_bias.ndim if colvec_bias is not None else 0
     device_capacity = get_device_capacity(A.device)
+    assert device_capacity[0] in [8, 9, 10, 11, 12], (
+        "Only SM8x, SM90, SM100, SM110, and SM120 are supported"
+    )
+    sf_dtype, sf_vec_size = None, None
+    if blockscaled:
+        assert not varlen_m and not gather_A, "Blockscaled GEMM does not support varlen/gather yet"
+        assert not concat_layout, "Blockscaled GEMM does not support concat_layout"
+        assert tile_K is None, "Blockscaled GEMM derives tile_K from the MMA instruction"
+        # A / B are still (l, m, k) / (l, n, k) here (perm3d only made views).
+        sf_dtype, sf_vec_size = validate_blockscaled_sf(A, B, SFA, SFB, device_capacity)
     if rounding_mode == RoundingMode.RS:
         assert device_capacity[0] == 10, "Stochastic rounding (RoundingMode.RS) requires SM100"
         d_major,
         c_major,
         postact_major,
+        (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N),
         (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
         rounding_mode=rounding_mode,
         sr_seed_mode=sr_seed_mode,
         use_tma_gather=use_tma_gather,
+        sf_dtype=sf_dtype,
+        sf_vec_size=sf_vec_size,
     )
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     def scalar_arg(scalar, mode, dtype=Int32):
     varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx)
     if device_capacity[0] in [10, 11]:
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, SFA, SFB)
     else:
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args)
 gemm_gated = gemm_act

build/torch-cuda/quack/gemm_base.py ADDED Viewed

	@@ -0,0 +1,731 @@

+# Copyright (c) 2026, Tri Dao.
+import enum
+import math
+from dataclasses import dataclass
+from typing import Callable, Dict, Literal, Optional, Sequence, Tuple
+import cutlass
+import cutlass.cute as cute
+import cutlass.pipeline as pipeline
+from cutlass import Boolean, Int32, const_expr
+from cutlass.cute.nvgpu import cpasync
+from cutlass.utils import LayoutEnum
+from . import copy_utils as copy_utils
+from .cute_dsl_utils import ParamsBase
+from .epi_ops import EpiSmemBytes
+from .pipeline import PipelineTmaAsync, PipelineTmaCpAsync
+from .rounding import RoundingMode, epilogue_sr_seed
+from .tile_scheduler import (
+    PersistenceMode,
+    TileScheduler,
+    TileSchedulerArguments,
+    VarlenMTileScheduler,
+    VarlenMTileSchedulerArguments,
+)
+from .varlen_utils import VarlenManager
+class NamedBarrierGemm(enum.IntEnum):
+    Epilogue = enum.auto()  # starts from 1 as barrier 0 is reserved for sync_threads()
+    # For mainloop load warps to signal that the epilogue load warp can start.
+    # This is to avoid loading C too early, interfering with loading A and B.
+    EpilogueLoad = enum.auto()
+    MmaWG0 = enum.auto()
+    MmaWG1 = enum.auto()
+    EpiWG0 = enum.auto()
+    EpiWG1 = enum.auto()
+    TmemPtr = enum.auto()
+    # CLC-multicast throttle: CTA0 load warp arrives once per tile started,
+    # CTA0 scheduler warp syncs once per CLC query (2 warps, 64 threads).
+    ClcThrottle = enum.auto()
+class GemmBase:
+    """Common non-mainloop pieces shared by GEMM architectures."""
+    arch = 0
+    @dataclass
+    class EpilogueArguments:
+        pass
+    EpilogueParams = ParamsBase
+    def epi_smem_warp_shape_mnk(self):
+        return (self.num_epi_warps, 1, 1)
+    @cute.jit
+    def epilogue(
+        self,
+        params: EpilogueParams,
+        epi_smem_tensors: Dict[str, cute.Tensor],
+        epi_pipeline: Optional[cutlass.pipeline.PipelineAsync],
+        epi_store_pipeline: Optional[cutlass.pipeline.PipelineAsync],
+        epi_read_state: Optional[cutlass.pipeline.PipelineState],
+        epi_producer_state: Optional[cutlass.pipeline.PipelineState],
+        epi_tile: cute.Tile,
+        load_acc_subtile: Callable,
+        tRS_rD: cute.Tensor,
+        tRS_rC: Optional[cute.Tensor],
+        tiled_copy_t2r: Optional[cute.TiledCopy],  # Only for Sm100
+        tiled_copy_r2s: cute.TiledCopy,
+        tRS_sD: cute.Tensor,
+        tiled_copy_s2r: Optional[cute.ThrCopy],
+        tSR_rC: Optional[cute.Tensor],
+        tSR_sC: Optional[cute.Tensor],
+        copy_D: Optional[Callable],
+        copy_C: Optional[Callable],
+        tile_coord_mnkl: cute.Coord,
+        varlen_manager: VarlenManager,
+        epilogue_barrier: cutlass.pipeline.NamedBarrier,
+        tile_scheduler,
+        tidx: Int32,
+        is_tma_warp: cutlass.Boolean,
+    ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState]:
+        has_C = const_expr(tRS_rC is not None)
+        has_epi_load = const_expr(self.epi_c_stage > 0)
+        has_D = const_expr(copy_D is not None)
+        use_tma_epi = const_expr(epi_store_pipeline is not None)
+        use_tma_c = const_expr(epi_pipeline is not None)
+        inline_epi_load = const_expr(copy_C is not None)
+        use_stochastic_rounding = const_expr(
+            self.rounding_mode == RoundingMode.RS
+            and self.acc_dtype == cutlass.Float32
+            and self.d_dtype == cutlass.BFloat16
+        )
+        # Setup aux outputs. Returns a tuple of ``(tiled_copy_r2s,
+        # tRS_sAuxOut, copy_aux_out)`` triples — empty for the default
+        # epilogue, one entry for the standard ``GemmAct``/``GemmGated``
+        # single-output mixins, multiple entries for multi-output mixins
+        # (e.g. ``T*tanh`` + ``1-tanh^2`` from one GEMM).
+        aux_out_ctxs = self.epi_setup_aux_out(
+            params,
+            epi_smem_tensors,
+            tiled_copy_r2s,
+            tiled_copy_t2r,
+            tile_coord_mnkl,
+            varlen_manager,
+            tidx,
+        )
+        epi_tile_shape = cute.zipped_divide(
+            cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile
+        ).shape[1]
+        epi_tile_layout = cute.make_ordered_layout(
+            epi_tile_shape, order=(0, 1) if const_expr(self.epi_m_major) else (1, 0)
+        )
+        epi_tile_num = cute.size(epi_tile_shape)
+        num_prev_subtiles = tile_scheduler.num_tiles_executed * epi_tile_num
+        epi_tensors = self.epi_begin(
+            params,
+            epi_smem_tensors,
+            epi_tile,
+            tiled_copy_t2r,
+            tiled_copy_r2s,
+            tile_coord_mnkl,
+            varlen_manager,
+            epilogue_barrier,
+            tidx,
+            tRS_rD.layout,
+        )
+        if const_expr(inline_epi_load):
+            for epi_idx in cutlass.range(min(epi_tile_num, self.epi_c_stage), unroll=1):
+                epi_coord_C = epi_tile_layout.get_hier_coord(epi_idx)
+                if const_expr(use_tma_c):
+                    if is_tma_warp:
+                        epi_pipeline.producer_acquire(epi_producer_state)
+                        copy_C(src_idx=epi_coord_C, producer_state=epi_producer_state)
+                        epi_pipeline.producer_commit(epi_producer_state)
+                    epi_producer_state.advance()
+                else:
+                    # TODO: turn this to cp.async instead of direct G2R copy
+                    copy_C(src_idx=epi_coord_C, dst_idx=epi_idx % self.epi_c_stage)
+            if const_expr(use_tma_c):
+                epilogue_barrier.arrive_and_wait()
+        for epi_idx in cutlass.range_constexpr(epi_tile_num):
+            epi_coord = epi_tile_layout.get_hier_coord(epi_idx)  # (epi_m, epi_n)
+            # Copy from acc to D registers
+            load_acc_subtile(tRS_rD, epi_coord)
+            if const_expr(has_epi_load):
+                if const_expr(use_tma_c):
+                    epi_pipeline.consumer_wait(epi_read_state)
+                    if const_expr(has_C):
+                        cute.copy(
+                            tiled_copy_s2r, tSR_sC[None, None, None, epi_read_state.index], tSR_rC
+                        )
+                    self.epi_tile_load_s2r(params, epi_tensors, epi_read_state.index)
+                    cute.arch.fence_view_async_shared()
+                    epi_pipeline.consumer_release(epi_read_state)
+                    epi_read_state.advance()
+                else:
+                    c_buffer = epi_idx % self.epi_c_stage
+                    cute.copy(tiled_copy_s2r, tSR_sC[None, None, None, c_buffer], tSR_rC)
+                    # TODO: cp.async wait once we switch to cp.async
+                    epilogue_barrier.arrive_and_wait()
+            epi_loop_tensors = self.epi_begin_loop(params, epi_tensors, epi_coord)
+            if const_expr(inline_epi_load and epi_idx + self.epi_c_stage < epi_tile_num):
+                epi_coord_C = epi_tile_layout.get_hier_coord(epi_idx + self.epi_c_stage)
+                if const_expr(use_tma_c):
+                    if is_tma_warp:
+                        epi_pipeline.producer_acquire(epi_producer_state)
+                        copy_C(src_idx=epi_coord_C, producer_state=epi_producer_state)
+                        epi_pipeline.producer_commit(epi_producer_state)
+                    epi_producer_state.advance()
+                else:
+                    epilogue_barrier.arrive_and_wait()
+                    copy_C(
+                        src_idx=epi_coord_C,
+                        dst_idx=(epi_idx + self.epi_c_stage) % self.epi_c_stage,
+                    )
+            # Returns a tuple of register tensors — one per aux output.
+            # Length matches ``aux_out_ctxs``. ``()`` for the default
+            # epilogue (no aux output).
+            tRS_rAuxOuts = self.epi_visit_subtile(params, epi_loop_tensors, tRS_rD, tRS_rC)
+            self.epi_end_loop(
+                params,
+                epi_tensors,
+                epi_coord,
+                epi_tile,
+                tiled_copy_t2r,
+                tiled_copy_r2s,
+                tile_coord_mnkl,
+                varlen_manager,
+                tidx,
+            )
+            # Convert each output to its storage dtype.
+            tRS_rAuxOuts_out = tuple(
+                self.epi_convert_aux_out(
+                    i,
+                    tRS_rAuxOuts[i],
+                    epi_loop_tensors.get("sr_seed"),
+                    tidx,
+                    tile_coord_mnkl,
+                    num_prev_subtiles,
+                    epi_idx,
+                )
+                for i in range(len(aux_out_ctxs))
+            )
+            if const_expr(use_tma_epi):
+                if is_tma_warp:
+                    epi_store_pipeline.producer_acquire()
+            else:
+                epilogue_barrier.arrive_and_wait()
+            if const_expr(use_tma_epi):
+                epilogue_barrier.arrive_and_wait()
+            epi_buffer = (num_prev_subtiles + epi_idx) % self.epi_stage
+            if const_expr(has_D):
+                tRS_sD_cur = tRS_sD[None, None, None, epi_buffer]
+                if const_expr(use_stochastic_rounding):
+                    seed = epilogue_sr_seed(
+                        epi_loop_tensors.get("sr_seed"),
+                        tile_coord_mnkl,
+                        num_prev_subtiles + epi_idx,
+                    )
+                    copy_utils.sr_cvt_copy(tiled_copy_r2s, tRS_rD, tRS_sD_cur, seed, tidx)
+                else:
+                    copy_utils.cvt_copy(tiled_copy_r2s, tRS_rD, tRS_sD_cur)
+            # Copy each aux output from registers to shared memory. All share
+            # the same ``epi_buffer`` index so the s2g TMA stores below happen
+            # in lockstep after the fence.
+            for i in cutlass.range_constexpr(len(aux_out_ctxs)):
+                tiled_copy_aux_out_r2s, tRS_sAuxOut, _ = aux_out_ctxs[i]
+                cute.copy(
+                    tiled_copy_aux_out_r2s,
+                    # Need contiguous for Sm80 and Sm120 where acc layout is ((2, 2), MMA_M, MMA_N)
+                    tiled_copy_aux_out_r2s.retile(tRS_rAuxOuts_out[i]).contiguous(),
+                    tRS_sAuxOut[None, None, None, epi_buffer],
+                )
+            if const_expr(use_tma_epi):
+                cute.arch.fence_view_async_shared()
+                epilogue_barrier.arrive_and_wait()
+                if is_tma_warp:
+                    if const_expr(has_D):
+                        copy_D(src_idx=epi_buffer, dst_idx=epi_coord)
+                    for i in cutlass.range_constexpr(len(aux_out_ctxs)):
+                        _, _, copy_aux_out = aux_out_ctxs[i]
+                        copy_aux_out(src_idx=epi_buffer, dst_idx=epi_coord)
+                    epi_store_pipeline.producer_commit()
+            else:
+                epilogue_barrier.arrive_and_wait()
+                if const_expr(has_D):
+                    copy_D(src_idx=epi_buffer, dst_idx=epi_coord)
+                for i in cutlass.range_constexpr(len(aux_out_ctxs)):
+                    _, _, copy_aux_out = aux_out_ctxs[i]
+                    copy_aux_out(src_idx=epi_buffer, dst_idx=epi_coord)
+                epilogue_barrier.arrive_and_wait()
+        self.epi_end(
+            params,
+            epi_tensors,
+            epi_tile,
+            tiled_copy_t2r,
+            tiled_copy_r2s,
+            tile_coord_mnkl,
+            varlen_manager,
+            tidx,
+        )
+        return epi_read_state, epi_producer_state
+    def get_scheduler_class(self, varlen_m: bool = False):
+        """Return the scheduler class to use. Override in subclasses for custom schedulers."""
+        return TileScheduler if not varlen_m else VarlenMTileScheduler
+    def resolve_epi_m_major(self, epilogue_args: EpilogueArguments):
+        return True
+    def get_scheduler_arguments(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        mD: Optional[cute.Tensor],
+        scheduler_args,
+        varlen_args,
+        epilogue_args,
+    ):
+        """Create scheduler arguments. Override in subclasses for custom schedulers."""
+        if const_expr(not self.is_persistent):
+            persistence_mode = PersistenceMode.NONE
+        else:
+            if const_expr(self.arch >= 100 and self.use_clc_persistence):
+                persistence_mode = PersistenceMode.CLC
+            elif const_expr(scheduler_args.tile_count_semaphore is not None):
+                persistence_mode = PersistenceMode.DYNAMIC
+            else:
+                persistence_mode = PersistenceMode.STATIC
+        if const_expr(varlen_args.mCuSeqlensM is None):
+            num_problems = (
+                mD.shape[2]
+                if mD is not None
+                else (
+                    mB.shape[2]
+                    if varlen_args.mCuSeqlensK is None
+                    else varlen_args.mCuSeqlensK.shape[0] - 1
+                )
+            )
+            problem_shape_ntile_mnl = (
+                cute.ceil_div(cute.size(mA, mode=[0]), self.cta_tile_shape_mnk[0]),
+                cute.ceil_div(cute.size(mB, mode=[0]), self.cta_tile_shape_mnk[1]),
+                num_problems,
+            )
+            tile_sched_args = TileSchedulerArguments(
+                problem_shape_ntile_mnl=problem_shape_ntile_mnl,
+                raster_order=scheduler_args.raster_order,
+                group_size=scheduler_args.max_swizzle_size,
+                cluster_shape_mnk=self.cluster_shape_mnk,
+                tile_count_semaphore=scheduler_args.tile_count_semaphore,
+                batch_idx_permute=scheduler_args.batch_idx_permute,
+                persistence_mode=persistence_mode,
+            )
+        else:
+            assert (mD is not None) or (epilogue_args.mAuxOut is not None) or (not self.gather_A)
+            problem_shape_ntile_mnl = (
+                None,
+                cute.ceil_div(cute.size(mB, mode=[0]), self.cta_tile_shape_mnk[1]),
+                varlen_args.mCuSeqlensM.shape[0] - 1,
+            )
+            tile_sched_args = VarlenMTileSchedulerArguments(
+                problem_shape_ntile_mnl=problem_shape_ntile_mnl,
+                total_m=(
+                    mD.shape[0]
+                    if mD is not None
+                    else (
+                        varlen_args.mAIdx.shape[0]
+                        if varlen_args.mAIdx is not None
+                        else cute.size(mA, mode=[0])
+                    )
+                ),
+                cu_seqlens_m=varlen_args.mCuSeqlensM,
+                max_active_clusters=scheduler_args.max_active_clusters,
+                raster_order=scheduler_args.raster_order,
+                group_size=scheduler_args.max_swizzle_size,
+                tile_shape_mn=self.cta_tile_shape_mnk[:2],
+                cluster_shape_mnk=self.cluster_shape_mnk,
+                tile_count_semaphore=scheduler_args.tile_count_semaphore,
+                persistence_mode=persistence_mode,
+            )
+        return tile_sched_args
+    @cute.jit
+    def epi_load_acc_subtile(
+        self,
+        tRS_rAcc: cute.Tensor,
+        tRS_rD: cute.Tensor,
+        epi_coord,  # (int, int)
+    ):
+        cute.autovec_copy(tRS_rAcc[None, None, None, epi_coord], tRS_rD)
+    @cute.jit
+    def epi_begin(
+        self,
+        params: EpilogueParams,
+        epi_smem_tensors: Dict[str, cute.Tensor],
+        epi_tile: cute.Tile,
+        tiled_copy_t2r: Optional[cute.TiledCopy],
+        tiled_copy_r2s: cute.TiledCopy,
+        tile_coord_mnkl: cute.Coord,
+        varlen_manager: VarlenManager,
+        epilogue_barrier: cutlass.pipeline.NamedBarrier,
+        tidx: Int32,
+        tRS_rD_layout=None,
+    ) -> Tuple[cute.Tensor, ...]:
+        return ()
+    def epi_begin_loop(
+        self, params: EpilogueParams, epi_tensors: Tuple[cute.Tensor, ...], epi_coord: cute.Coord
+    ) -> Tuple[cute.Tensor, ...]:
+        return ()
+    def epi_visit_subtile(
+        self,
+        params: EpilogueParams,
+        epi_loop_tensors: Tuple[cute.Tensor, ...],
+        tRS_rD: cute.Tensor,
+        tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
+        return ()
+    def epi_visit_acc(
+        self,
+        params: EpilogueParams,
+        acc: cute.Tensor,
+        tiled_mma: cute.TiledMma,
+        tile_coord_mnkl: cute.Coord,
+        tidx: Int32,
+    ) -> None:
+        pass
+    @cute.jit
+    def epi_end_loop(
+        self,
+        params: EpilogueParams,
+        epi_tensors: Tuple[cute.Tensor, ...],
+        epi_coord: cute.Coord,
+        epi_tile: cute.Tile,
+        tiled_copy_t2r: Optional[cute.TiledCopy],
+        tiled_copy_r2s: cute.TiledCopy,
+        tile_coord_mnkl: cute.Coord,
+        varlen_manager,
+        tidx,
+    ) -> None:
+        pass
+    @cute.jit
+    def epi_end(
+        self,
+        params: EpilogueParams,
+        epi_tensors: Tuple[cute.Tensor, ...],
+        epi_tile: cute.Tile,
+        tiled_copy_t2r: Optional[cute.TiledCopy],
+        tiled_copy_r2s: cute.TiledCopy,
+        tile_coord_mnkl: cute.Coord,
+        varlen_manager,
+        tidx,
+    ) -> None:
+        pass
+    def epi_to_underlying_arguments(
+        self, args: EpilogueArguments, *, loc=None, ip=None
+    ) -> EpilogueParams:
+        return self.EpilogueParams()
+    def epi_get_tma_atoms(
+        self, params: EpilogueParams, *, loc=None, ip=None
+    ) -> list[cute.CopyAtom]:
+        """Subclasses can override this."""
+        return []
+    def epi_tile_load_g2s_copy_fns(
+        self,
+        params,
+        epi_smem_tensors,
+        tile_coord_mnkl,
+        varlen_manager,
+        epi_pipeline,
+    ):
+        return ()
+    @cute.jit
+    def epi_tile_load_s2r(self, params, epi_tensors, stage_idx):
+        pass
+    @staticmethod
+    def epi_smem_bytes(
+        args: Optional[EpilogueArguments],
+        cta_tile_shape_mnk: Tuple[int, int, int],
+        epi_tile: cute.Tile,
+        warp_shape_mnk: Tuple[int, int, int] | None = None,
+    ) -> EpiSmemBytes:
+        return EpiSmemBytes()
+    def epi_get_smem_struct(self, params: EpilogueParams):
+        return cute.struct.MemRange[Int32, 0]  # Dummy struct
+    def epi_get_smem_tensors(self, params: EpilogueParams, storage) -> Dict[str, cute.Tensor]:
+        return {}
+    def epi_setup_aux_out(
+        self,
+        params,
+        epi_smem_tensors,
+        tiled_copy_r2s,
+        tiled_copy_t2r,
+        tile_coord_mnkl,
+        varlen_manager,
+        tidx,
+    ):
+        """Return a tuple of ``(tiled_copy_r2s, tRS_sAuxOut, copy_aux_out)``
+        triples — one per aux output. The default epilogue has no aux output,
+        so the tuple is empty.
+        """
+        return ()
+    @cute.jit
+    def epi_convert_aux_out(
+        self,
+        output_idx: cutlass.Constexpr[int],
+        tRS_rAuxOut,
+        sr_seed,
+        tidx,
+        tile_coord_mnkl,
+        num_prev_subtiles,
+        epi_idx,
+    ):
+        """Convert one aux output register tensor from acc_dtype to its storage
+        dtype. ``output_idx`` selects which aux output this call is for
+        (single-output mixins can ignore it).
+        """
+        return tRS_rAuxOut
+class GemmTmaBase(GemmBase):
+    """Common TMA descriptor and pipeline helpers for SM90+ GEMM paths."""
+    @cute.jit
+    def load_tma(
+        self,
+        pipeline: cutlass.pipeline.PipelineAsync,
+        producer_state: cutlass.pipeline.PipelineState,
+        copy_fns: Sequence[Optional[Callable]],
+        k_tile_cnt: Int32,
+    ) -> cutlass.pipeline.PipelineState:
+        # Peek (try_wait) AB buffer empty for k_block = prefetch_k_tile_cnt.
+        peek_empty_status = Boolean(True)
+        if 0 < k_tile_cnt:
+            peek_empty_status = pipeline.producer_try_acquire(producer_state)
+        # TMA load
+        for k_tile in cutlass.range(k_tile_cnt, unroll=1):
+            # Wait for A/B buffers to be empty before loading into them.
+            # Also sets the transaction barrier for the A/B buffers.
+            pipeline.producer_acquire(producer_state, peek_empty_status)
+            tma_bar_ptr = pipeline.producer_get_barrier(producer_state)
+            smem_idx = producer_state.index
+            for copy_fn in copy_fns:
+                if const_expr(copy_fn is not None):
+                    copy_fn(k_tile, smem_idx, tma_bar_ptr=tma_bar_ptr)
+            # Mainloop pipeline's producer commit is a NOP for TMA pipelines.
+            pipeline.producer_commit(producer_state)
+            producer_state.advance()
+            peek_empty_status = Boolean(True)
+            if k_tile + 1 < k_tile_cnt:
+                peek_empty_status = pipeline.producer_try_acquire(producer_state)
+        return producer_state
+    def _make_gmem_tiled_copy_A(self, dtype, major_mode, num_threads, copy_bits=128):
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=copy_bits,
+        )
+        copy_elems = copy_bits // dtype.width
+        loads_per_cache_line = 128 * 8 // copy_bits  # 128 bytes per cache line
+        shape_dim_1 = cute.size(self.cta_tile_shape_mnk[2]) // copy_elems
+        if shape_dim_1 > loads_per_cache_line:
+            shape_dim_1 = math.gcd(shape_dim_1, loads_per_cache_line)
+        # thread layout for copy
+        thread_layout = cute.make_layout(
+            (num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1)
+        )
+        if major_mode != LayoutEnum.ROW_MAJOR:
+            shape_dim_0 = cute.size(self.cta_tile_shape_mnk[0]) // copy_elems
+            if shape_dim_0 > loads_per_cache_line:
+                shape_dim_0 = math.gcd(shape_dim_0, loads_per_cache_line)
+            thread_layout = cute.make_layout(
+                (shape_dim_0, num_threads // shape_dim_0), stride=(1, shape_dim_0)
+            )
+        # Value layout for copy
+        value_layout = (
+            cute.make_layout((1, copy_elems))
+            if major_mode == LayoutEnum.ROW_MAJOR
+            else cute.make_layout((copy_elems, 1))
+        )
+        return cute.make_tiled_copy_tv(atom_async_copy, thread_layout, value_layout)
+    def make_tma_load_atoms_and_tensors(
+        self,
+        mA: cute.Tensor,
+        mB: cute.Tensor,
+        a_smem_layout: cute.ComposedLayout,
+        b_smem_layout: cute.ComposedLayout,
+        varlen_k: bool,
+    ):
+        tma_atom_a, tma_tensor_a = None, None
+        if const_expr(not self.gather_A):
+            tma_atom_a, tma_tensor_a = self._make_tma_atoms_and_tensors(
+                copy_utils.create_ragged_tensor_for_tma(mA, ragged_dim=1)
+                if varlen_k and not self.gather_A
+                else mA,
+                a_smem_layout,
+                (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[2]),
+                self.cluster_shape_mnk[1],
+            )
+        tma_atom_b, tma_tensor_b = self._make_tma_atoms_and_tensors(
+            copy_utils.create_ragged_tensor_for_tma(mB, ragged_dim=1) if varlen_k else mB,
+            b_smem_layout,
+            (self.cta_tile_shape_mnk[1], self.cta_tile_shape_mnk[2]),
+            self.cluster_shape_mnk[0],
+        )
+        return tma_atom_a, tma_tensor_a, tma_atom_b, tma_tensor_b
+    def make_tma_epilogue_atoms_and_tensors(
+        self,
+        mD: Optional[cute.Tensor],
+        mC: Optional[cute.Tensor],
+        epilogue_args,
+        varlen_m: bool,
+    ):
+        tma_atom_d, tma_tensor_d = None, None
+        if const_expr(mD is not None):
+            tma_atom_d, tma_tensor_d = self._make_tma_epi_atoms_and_tensors(
+                copy_utils.create_ragged_tensor_for_tma(mD, ragged_dim=0, ptr_shift=True)
+                if varlen_m
+                else mD,
+                self.epi_smem_layout_staged,
+                self.epi_tile,
+                op_type="store"
+                if not (hasattr(epilogue_args, "add_to_output") and epilogue_args.add_to_output)
+                else "add",
+            )
+        tma_atom_c, tma_tensor_c = None, None
+        if const_expr(mC is not None):
+            tma_atom_c, tma_tensor_c = self._make_tma_epi_atoms_and_tensors(
+                mC, self.epi_c_smem_layout_staged, self.epi_tile, op_type="load"
+            )
+        return tma_atom_d, tma_tensor_d, tma_atom_c, tma_tensor_c
+    def epilog_gmem_copy_and_partition(
+        self,
+        atom: cute.CopyAtom | cute.TiledCopy,
+        mD_mn: cute.Tensor,
+        tile_shape_mn: cute.Tile,
+        epi_tile: cute.Tile,
+        sD: cute.Tensor,
+        tile_coord_mnkl: cute.Coord,
+    ) -> Tuple[cute.Tensor, cute.Tensor]:
+        gD = cute.local_tile(mD_mn, tile_shape_mn, tile_coord_mnkl[:2])  # (bM, bN)
+        tDgD_for_tma_partition = cute.zipped_divide(gD, epi_tile)
+        is_s2g = isinstance(
+            atom.op, (cpasync.CopyBulkTensorTileS2GOp, cpasync.CopyReduceBulkTensorTileS2GOp)
+        )
+        src_tensor, dst_tensor = (
+            (sD, tDgD_for_tma_partition) if is_s2g else (tDgD_for_tma_partition, sD)
+        )
+        return copy_utils.tma_get_copy_fn(
+            atom,
+            cta_coord=0,
+            cta_layout=cute.make_layout(1),
+            src_tensor=src_tensor,
+            dst_tensor=dst_tensor,
+        )
+    def make_ab_pipeline(
+        self,
+        tiled_mma: cute.TiledMma,
+        cluster_layout_vmnk: cute.Layout,
+    ):
+        # Threads/warps participating in this pipeline
+        producer_cnt = 1 if const_expr(not self.gather_A) else 1 + self.num_ab_load_warps * 32
+        ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt)
+        # Each warp will contribute to the arrive count with the number of mcast size
+        mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
+        consumer_arrive_cnt = mcast_size * tiled_mma.size // cute.arch.WARP_SIZE
+        ab_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, consumer_arrive_cnt
+        )
+        pipeline_cls = pipeline.PipelineTmaAsync if not self.gather_A else PipelineTmaCpAsync
+        return pipeline_cls.create(
+            num_stages=self.ab_stage,
+            producer_group=ab_pipeline_producer_group,
+            consumer_group=ab_pipeline_consumer_group,
+            tx_count=self.num_tma_load_bytes,
+            cta_layout_vmnk=cluster_layout_vmnk,
+            defer_sync=True,
+        )
+    def make_epi_pipeline(
+        self,
+        tx_count: int,
+    ):
+        epi_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
+        # Each warp will contribute 1 to the arrive count
+        consumer_arrive_cnt = self.num_epi_warps
+        epi_pipeline_consumer_group = pipeline.CooperativeGroup(
+            pipeline.Agent.Thread, consumer_arrive_cnt
+        )
+        return PipelineTmaAsync.create(
+            num_stages=self.epi_c_stage,
+            producer_group=epi_pipeline_producer_group,
+            consumer_group=epi_pipeline_consumer_group,
+            tx_count=tx_count,
+            defer_sync=True,
+            elect_one_release=True,
+            syncwarp_before_release=True,
+        )
+    def make_epi_store_pipeline(self):
+        num_epi_threads = self.num_epi_warps * cute.arch.WARP_SIZE
+        epi_store_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_epi_threads)
+        return pipeline.PipelineTmaStore.create(
+            num_stages=self.epi_stage, producer_group=epi_store_producer_group
+        )
+    @staticmethod
+    def _make_tma_epi_atoms_and_tensors(
+        tensor_d: cute.Tensor,
+        epi_smem_layout_staged: cute.ComposedLayout,
+        epi_tile: Tuple[int, int],
+        op_type: Literal["store", "load", "add"],
+    ) -> Tuple[cute.CopyAtom, cute.Tensor]:
+        """Create TMA atoms and tensors for storing D or loading C."""
+        assert op_type in ["load", "store", "add"]
+        epi_smem_layout = cute.slice_(epi_smem_layout_staged, (None, None, 0))
+        d_cta_v_layout = cute.composition(cute.make_identity_layout(tensor_d.shape), epi_tile)
+        op = {
+            "load": cpasync.CopyBulkTensorTileG2SOp(),
+            "store": cpasync.CopyBulkTensorTileS2GOp(),
+            "add": cpasync.CopyReduceBulkTensorTileS2GOp(cpasync.ReductionOp.ADD),
+        }[op_type]
+        tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom(
+            op, tensor_d, epi_smem_layout, d_cta_v_layout
+        )
+        return tma_atom_d, tma_tensor_d
+    @staticmethod
+    def _make_tma_atoms_and_tensors(
+        tensor: cute.Tensor,
+        smem_layout: cute.ComposedLayout,
+        smem_tile: Tuple[int, int],
+        mcast_dim: int,
+    ) -> Tuple[cute.CopyAtom, cute.Tensor]:
+        """Create TMA atoms and tensors for input tensors."""
+        # block_copy takes compiler-driven multicast metadata at the copy site,
+        # so the TMA atom itself must stay the non-multicast variant here.
+        op = cpasync.CopyBulkTensorTileG2SOp()
+        tma_atom, tma_tensor = cpasync.make_tiled_tma_atom(op, tensor, smem_layout, smem_tile)
+        return tma_atom, tma_tensor

build/torch-cuda/quack/gemm_blockscaled_interface.py DELETED Viewed

@@ -1,326 +0,0 @@
-# Copyright (c) 2026, Tri Dao.
-"""PyTorch-friendly interface for the SM100 MXFP8 blockscaled GEMM.
-Shape / layout conventions (matches torch.matmul, torch._scaled_mm, cuBLAS):
-  A:       (M, K)     or (L, M, K)       dtype float8_e4m3fn, K-contiguous (row-major)
-  B:       (K, N)     or (L, K, N)       dtype float8_e4m3fn, K-contiguous (col-major)
-  A_scale: (M, K/32)  or (L, M, K/32)    dtype float8_e8m0fnu, K-contiguous
-  B_scale: (K/32, N)  or (L, K/32, N)    dtype float8_e8m0fnu, K-contiguous
-  out:     (M, N)     or (L, M, N)       dtype bfloat16/float16, contiguous
-"K-contiguous" means stride 1 on the K axis. This matches how torchao/cuBLAS
-use `torch._scaled_mm(a, b.t(), ...)`:
-  - you store a weight as nn.Linear-style `W` of shape `(N, K)` row-major
-  - you pass `W.mT` (a zero-copy view of shape (K, N) with K-contig) as B
-The interface applies `.mT` internally to reach the `(N, K) K-major` layout
-the quack kernel consumes. No data is copied.
-"""
-from functools import lru_cache
-from typing import Optional, Tuple
-import torch
-from torch import Tensor
-import cutlass
-from .blockscaled_gemm_utils import (
-    ceil_div,
-    compile_blockscaled_gemm_tvm_ffi,
-    pack_scale_2d_to_blocked_contig,
-    scale_blocked_for_cublas,
-    scale_view_for_kernel,
-)
-from .gemm_default_epi import GemmDefaultSm100
-from .mx_utils import to_mx
-_SF_VEC_SIZE = 32
-_TORCH_TO_CUTLASS_D = {
-    torch.bfloat16: cutlass.BFloat16,
-    torch.float16: cutlass.Float16,
-    torch.float32: cutlass.Float32,
-}
-def _default_tiler_cluster(m: int, n: int) -> Tuple[Tuple[int, int], Tuple[int, int]]:
-    """Pick a reasonable default (mma_tiler_mn, cluster_shape_mn)."""
-    if m >= 512 and n >= 128:
-        return (256, 128), (2, 1)
-    return (128, 128), (1, 1)
-@lru_cache(maxsize=64)
-def _compile_cached(
-    m: int,
-    n: int,
-    k: int,
-    l: int,
-    mma_tiler_mn: Tuple[int, int],
-    cluster_shape_mn: Tuple[int, int],
-    out_torch_dtype,
-    ab_dtype_cutlass,
-    sf_dtype_cutlass,
-):
-    """Compile kernel for a given (shape, dtype, tiler, cluster) and cache it."""
-    dev = torch.device("cuda")
-    rm = ceil_div(m, 128)
-    rn = ceil_div(n, 128)
-    rk = ceil_div(k // _SF_VEC_SIZE, 4)
-    # K-major: (l, m, k) contiguous, viewed as (m, k, l) strides (k, 1, m*k)
-    fake_mA = torch.empty(l, m, k, dtype=torch.float8_e4m3fn, device=dev).permute(1, 2, 0)
-    fake_mB = torch.empty(l, n, k, dtype=torch.float8_e4m3fn, device=dev).permute(1, 2, 0)
-    # N-major: (l, m, n) contiguous, viewed as (m, n, l) strides (n, 1, m*n)
-    fake_mD = torch.empty(l, m, n, dtype=out_torch_dtype, device=dev).permute(1, 2, 0)
-    fake_sc_A = torch.empty(l, rm, rk, 512, dtype=torch.float8_e8m0fnu, device=dev)
-    fake_sc_B = torch.empty(l, rn, rk, 512, dtype=torch.float8_e8m0fnu, device=dev)
-    fake_mSFA = scale_view_for_kernel(fake_sc_A, m, k // _SF_VEC_SIZE, l)
-    fake_mSFB = scale_view_for_kernel(fake_sc_B, n, k // _SF_VEC_SIZE, l)
-    return compile_blockscaled_gemm_tvm_ffi(
-        ab_dtype_cutlass,
-        sf_dtype_cutlass,
-        _SF_VEC_SIZE,
-        _TORCH_TO_CUTLASS_D[out_torch_dtype],
-        mma_tiler_mn,
-        cluster_shape_mn,
-        fake_mA,
-        fake_mB,
-        fake_mD,
-        fake_mSFA,
-        fake_mSFB,
-    )
-def _as_3d(x: Tensor, ndim_in: int) -> Tensor:
-    """Add a leading batch dim if input is 2D. Returns a view."""
-    if ndim_in == 2:
-        return x.unsqueeze(0)
-    return x
-def _to_kernel_layout(
-    A: Tensor,
-    B: Tensor,
-    A_scale: Tensor,
-    B_scale: Tensor,
-) -> Tuple[int, int, int, int, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, bool]:
-    """Normalize shapes/strides, validate, and repack scales. Returns
-    (m, n, k, l, mA_mkl, mB_nkl, sc_contig_A, sc_contig_B, sfa_view, sfb_view, was_2d).
-    A: (M,K) or (L,M,K) K-contig.  B: (K,N) or (L,K,N) K-contig.
-    A_scale: (M,K/32) or (L,M,K/32) K-contig.  B_scale: (K/32,N) or (L,K/32,N) K-contig.
-    """
-    assert A.dtype == torch.float8_e4m3fn, f"A dtype must be float8_e4m3fn, got {A.dtype}"
-    assert B.dtype == torch.float8_e4m3fn, f"B dtype must be float8_e4m3fn, got {B.dtype}"
-    assert A_scale.dtype == torch.float8_e8m0fnu
-    assert B_scale.dtype == torch.float8_e8m0fnu
-    was_2d = A.dim() == 2
-    # Flip B from (K,N) to (N,K) via .mT (zero-copy). User's B K-contig → .mT K-contig.
-    A3 = _as_3d(A, A.dim())  # (l, m, k) K-contig row-major expected
-    B3 = _as_3d(B, B.dim()).mT  # (l, n, k) K-contig (view) from (l, k, n)
-    l, m, k = A3.shape
-    l2, n, k2 = B3.shape
-    assert l == l2, f"batch mismatch: A={l}, B={l2}"
-    assert k == k2, f"K mismatch: A K={k}, B K={k2}"
-    assert k % _SF_VEC_SIZE == 0, f"K ({k}) must be divisible by {_SF_VEC_SIZE}"
-    assert A3.stride(-1) == 1, "A must be K-contiguous (stride 1 on K)"
-    assert B3.stride(-1) == 1, (
-        "B must be K-contiguous on its K axis (pass .mT of an (N,K) row-major tensor)"
-    )
-    sf_k = k // _SF_VEC_SIZE
-    as3 = _as_3d(A_scale, A_scale.dim())  # expected (l, m, sf_k) K-contig row-major
-    bs3 = _as_3d(B_scale, B_scale.dim()).mT  # (l, n, sf_k) K-contig (view) from (l, sf_k, n)
-    assert as3.stride(-1) == 1, "A_scale must be K-contiguous"
-    assert bs3.stride(-1) == 1, (
-        "B_scale must be K-contiguous on its K axis (pass .mT of an (N, K/32) row-major tensor)"
-    )
-    assert as3.shape == (l, m, sf_k), (
-        f"A_scale shape: expected (l={l},m={m},sf_k={sf_k}) K-contig, got {tuple(as3.shape)}"
-    )
-    assert bs3.shape == (l, n, sf_k), (
-        f"B_scale shape: expected .mT of (l={l},sf_k={sf_k},n={n}) -> ({l},{n},{sf_k}), got {tuple(bs3.shape)}"
-    )
-    # Force row-major contiguous for packer/kernel consumption.
-    # A3 / B3 are views — .contiguous() materializes (l,m,k) / (l,n,k) row-major.
-    A3_c = A3.contiguous()
-    B3_c = B3.contiguous()
-    # (l, m, k) -> (m, k, l) K-major view (no copy; strides (k, 1, m*k))
-    mA_mkl = A3_c.permute(1, 2, 0)
-    mB_nkl = B3_c.permute(1, 2, 0)
-    sc_contig_A = pack_scale_2d_to_blocked_contig(as3.contiguous())
-    sc_contig_B = pack_scale_2d_to_blocked_contig(bs3.contiguous())
-    sfa_view = scale_view_for_kernel(sc_contig_A, m, sf_k, l)
-    sfb_view = scale_view_for_kernel(sc_contig_B, n, sf_k, l)
-    return m, n, k, l, mA_mkl, mB_nkl, sc_contig_A, sc_contig_B, sfa_view, sfb_view, was_2d
-def mxfp8_gemm_out(
-    A: Tensor,
-    B: Tensor,
-    A_scale: Tensor,
-    B_scale: Tensor,
-    out: Tensor,
-    *,
-    mma_tiler_mn: Optional[Tuple[int, int]] = None,
-    cluster_shape_mn: Optional[Tuple[int, int]] = None,
-) -> None:
-    """MXFP8 blockscaled GEMM with pre-allocated output. See module doc for shape conventions."""
-    m, n, k, l, mA, mB, _scA, _scB, sfa, sfb, was_2d = _to_kernel_layout(A, B, A_scale, B_scale)
-    out_dtype = out.dtype
-    assert out_dtype in _TORCH_TO_CUTLASS_D, f"unsupported out dtype: {out_dtype}"
-    expected_out_shape = (m, n) if was_2d else (l, m, n)
-    assert tuple(out.shape) == expected_out_shape, (
-        f"out shape {tuple(out.shape)} != expected {expected_out_shape}"
-    )
-    assert out.is_contiguous(), "out must be contiguous"
-    # View caller's contiguous (M,N) or (L,M,N) as (M,N,L) N-major strided view, no copy.
-    out_3d = out.unsqueeze(0) if was_2d else out  # (l, m, n)
-    mD = out_3d.permute(1, 2, 0)  # (m, n, l), strides (n, 1, m*n)
-    if mma_tiler_mn is None or cluster_shape_mn is None:
-        tlr, clu = _default_tiler_cluster(m, n)
-        mma_tiler_mn = mma_tiler_mn or tlr
-        cluster_shape_mn = cluster_shape_mn or clu
-    if not GemmDefaultSm100.can_implement_blockscaled(
-        cutlass.Float8E4M3FN,
-        cutlass.Float8E8M0FNU,
-        _SF_VEC_SIZE,
-        _TORCH_TO_CUTLASS_D[out_dtype],
-        mma_tiler_mn,
-        cluster_shape_mn,
-        m,
-        n,
-        k,
-        l,
-        "k",
-        "k",
-        "n",
-    ):
-        raise ValueError(
-            f"unsupported config: m={m}, n={n}, k={k}, l={l}, "
-            f"tiler={mma_tiler_mn}, cluster={cluster_shape_mn}"
-        )
-    runner = _compile_cached(
-        m,
-        n,
-        k,
-        l,
-        mma_tiler_mn,
-        cluster_shape_mn,
-        out_dtype,
-        cutlass.Float8E4M3FN,
-        cutlass.Float8E8M0FNU,
-    )
-    runner(mA, mB, mD, sfa, sfb)
-def mxfp8_gemm(
-    A: Tensor,
-    B: Tensor,
-    A_scale: Tensor,
-    B_scale: Tensor,
-    out: Optional[Tensor] = None,
-    out_dtype: torch.dtype = torch.bfloat16,
-    *,
-    mma_tiler_mn: Optional[Tuple[int, int]] = None,
-    cluster_shape_mn: Optional[Tuple[int, int]] = None,
-) -> Tensor:
-    """MXFP8 blockscaled GEMM. Allocates output if not provided."""
-    if out is None:
-        # A: (M,K) or (L,M,K); B: (K,N) or (L,K,N); out: (M,N) or (L,M,N)
-        if A.dim() == 2:
-            out_shape = (A.shape[0], B.shape[1])
-        else:
-            out_shape = (A.shape[0], A.shape[1], B.shape[2])
-        out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
-    mxfp8_gemm_out(
-        A,
-        B,
-        A_scale,
-        B_scale,
-        out,
-        mma_tiler_mn=mma_tiler_mn,
-        cluster_shape_mn=cluster_shape_mn,
-    )
-    return out
-def mxfp8_quantize(x: Tensor) -> Tuple[Tensor, Tensor]:
-    """Quantize a (..., K) bf16/fp32 tensor to MXFP8. Returns (qdata, scale_2d)
-    in torchao-convention layout. Last dim (K) must be divisible by 32."""
-    assert x.shape[-1] % _SF_VEC_SIZE == 0, (
-        f"last dim ({x.shape[-1]}) must be divisible by {_SF_VEC_SIZE}"
-    )
-    return to_mx(x.contiguous(), _SF_VEC_SIZE)
-def mxfp8_gemm_quantize(
-    A: Tensor,
-    B: Tensor,
-    out: Optional[Tensor] = None,
-    out_dtype: torch.dtype = torch.bfloat16,
-    *,
-    mma_tiler_mn: Optional[Tuple[int, int]] = None,
-    cluster_shape_mn: Optional[Tuple[int, int]] = None,
-) -> Tensor:
-    """High-level: quantize bf16 A, B_as_NK to MXFP8, then run C = A @ B_as_NK.mT.
-    Inputs: A=(M,K)/(L,M,K), B_as_NK=(N,K)/(L,N,K) bf16/fp32. Quantization
-    scales along the last (K) dim. Returned output has shape (M,N)/(L,M,N)."""
-    A_q, A_sc = mxfp8_quantize(A)
-    B_q, B_sc = mxfp8_quantize(B)
-    # B_q, B_sc are (..., N, K) / (..., N, K/32). Flip to (..., K, N) / (..., K/32, N)
-    # K-contig zero-copy views to match the interface convention.
-    return mxfp8_gemm(
-        A_q,
-        B_q.mT,
-        A_sc,
-        B_sc.mT,
-        out=out,
-        out_dtype=out_dtype,
-        mma_tiler_mn=mma_tiler_mn,
-        cluster_shape_mn=cluster_shape_mn,
-    )
-def mxfp8_gemm_cublas(
-    A: Tensor,
-    B: Tensor,
-    A_scale: Tensor,
-    B_scale: Tensor,
-    out_dtype: torch.dtype = torch.bfloat16,
-) -> Tensor:
-    """Reference path via torch._scaled_mm. Requires l=1 (or 2D inputs)."""
-    m, n, k, l, _mA, _mB, sc_A, sc_B, _sfa, _sfb, was_2d = _to_kernel_layout(A, B, A_scale, B_scale)
-    assert l == 1, "torch._scaled_mm MXFP8 path is 2D only; pass 2D inputs or l=1"
-    # torch._scaled_mm: A=(M,K) row-major, B=(K,N) col-major (both K-contig) -- same layout user gave us.
-    a2d = A if A.dim() == 2 else A.squeeze(0)
-    b2d = B if B.dim() == 2 else B.squeeze(0)
-    sca = scale_blocked_for_cublas(sc_A, m, k // _SF_VEC_SIZE, 0)
-    scb = scale_blocked_for_cublas(sc_B, n, k // _SF_VEC_SIZE, 0)
-    out = torch._scaled_mm(
-        a2d,
-        b2d,
-        scale_a=sca,
-        scale_b=scb,
-        out_dtype=out_dtype,
-    )
-    return out if was_2d else out.unsqueeze(0)
-def mxfp8_gemm_ref(
-    A: Tensor,
-    B: Tensor,
-    A_scale: Tensor,
-    B_scale: Tensor,
-    out_dtype: torch.dtype = torch.bfloat16,
-) -> Tensor:
-    """Dequantize + plain matmul reference. A=(M,K), B=(K,N)."""
-    was_2d = A.dim() == 2
-    # (l, m, k)
-    A3 = _as_3d(A, A.dim()).float()
-    # B is (K, N)/(L, K, N); flip to (l, n, k) for dequant by last-dim
-    B3 = _as_3d(B, B.dim()).mT.contiguous().float()
-    as3 = _as_3d(A_scale, A_scale.dim()).float()
-    bs3 = _as_3d(B_scale, B_scale.dim()).mT.contiguous().float()
-    a_dq = A3 * as3.repeat_interleave(_SF_VEC_SIZE, dim=-1)
-    b_dq = B3 * bs3.repeat_interleave(_SF_VEC_SIZE, dim=-1)
-    out3 = torch.einsum("lmk,lnk->lmn", a_dq, b_dq).to(out_dtype)
-    return out3.squeeze(0) if was_2d else out3

build/torch-cuda/quack/gemm_config.py CHANGED Viewed

@@ -1,4 +1,4 @@
-# Copyright (C) 2025, Fri Dao.
 import itertools
 from typing import Optional, List
 from functools import partial
@@ -9,11 +9,14 @@ from dataclasses import dataclass
 class GemmConfig:
     tile_m: int = 128
     tile_n: int = 192
     pingpong: bool = True
     # by default, we use dynamic persistent tile scheduler on SM100 but not on SM90
     is_dynamic_persistent: bool = True
     cluster_m: int = 2
     cluster_n: int = 1
     swap_ab: bool = False
     # raster_order: int = 1
     max_swizzle_size: int = 8
@@ -70,6 +73,39 @@ def _get_sm90_configs(
     ]
 def _get_sm100_configs(
     epilogue: Optional[str] = None,
 ) -> List[GemmConfig]:
@@ -141,14 +177,15 @@ def get_all_configs(
     epilogue: Optional[str] = None,
     tune_coop: bool = True,
 ) -> List[GemmConfig]:
-    """Return autotuning configs for all supported device capabilities (sm90 + sm100 + sm120).
     Each GemmConfig is tagged with its target device_capacity, so the caller can
     filter at runtime based on the actual device. This avoids querying the device
     (and initializing a CUDA context) at import time.
     """
     return (
-        _get_sm90_configs(epilogue, tune_coop)
         + _get_sm100_configs(epilogue)
         + _get_sm120_configs(epilogue, tune_coop)
     )

+# Copyright (C) 2025, Tri Dao.
 import itertools
 from typing import Optional, List
 from functools import partial
 class GemmConfig:
     tile_m: int = 128
     tile_n: int = 192
+    tile_k: int | None = None
+    num_warps: int | None = None
     pingpong: bool = True
     # by default, we use dynamic persistent tile scheduler on SM100 but not on SM90
     is_dynamic_persistent: bool = True
     cluster_m: int = 2
     cluster_n: int = 1
+    cluster_k: int = 1
     swap_ab: bool = False
     # raster_order: int = 1
     max_swizzle_size: int = 8
     ]
+def _get_sm80_configs() -> List[GemmConfig]:
+    tile_mn_warps_vals = [
+        (128, 128, 4),
+        (128, 128, 8),
+        (128, 160, 4),
+        # TODO: Make 128x160 work with 8 warps. It currently makes the accumulator
+        # N layout odd and fails epilogue retile.
+        (128, 192, 4),
+        (128, 192, 8),
+        (128, 256, 8),
+        (128, 64, 4),
+        (64, 128, 4),
+    ]
+    return [
+        GemmConfig(
+            tile_m=tile_m,
+            tile_n=tile_n,
+            tile_k=tile_k,
+            num_warps=num_warps,
+            pingpong=False,
+            cluster_m=1,
+            cluster_n=1,
+            swap_ab=swap_ab,
+            device_capacity=8,
+            is_dynamic_persistent=False,
+            use_tma_gather=False,
+        )
+        for (tile_m, tile_n, num_warps), tile_k, swap_ab in itertools.product(
+            tile_mn_warps_vals, [32, 64], [False, True]
+        )
+    ]
 def _get_sm100_configs(
     epilogue: Optional[str] = None,
 ) -> List[GemmConfig]:
     epilogue: Optional[str] = None,
     tune_coop: bool = True,
 ) -> List[GemmConfig]:
+    """Return autotuning configs for all supported device capabilities.
     Each GemmConfig is tagged with its target device_capacity, so the caller can
     filter at runtime based on the actual device. This avoids querying the device
     (and initializing a CUDA context) at import time.
     """
     return (
+        _get_sm80_configs()
+        + _get_sm90_configs(epilogue, tune_coop)
         + _get_sm100_configs(epilogue)
         + _get_sm120_configs(epilogue, tune_coop)
     )

build/torch-cuda/quack/gemm_dact.py CHANGED Viewed

@@ -8,15 +8,15 @@ from torch import Tensor
 import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, Float32, const_expr
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
 from .gemm_default_epi import GemmDefaultEpiMixin
 from .gemm_act import GemmActMixin
-from .epi_ops import ColVecReduce, colvec_reduce_accumulate
 from .compile_utils import make_fake_tensor as fake_tensor
 from .cute_dsl_utils import (
-    ParamsBase,
     mlir_namedtuple,
     torch2cute_dtype_map,
     get_device_capacity,
@@ -33,10 +33,10 @@ from .gemm_tvm_ffi_utils import (
     make_fake_gemm_tensors,
     compile_gemm_kernel,
 )
-from .cache_utils import jit_cache
 from .rounding import RoundingMode
-from . import layout_utils as layout_utils
 from .activation import dact_fn_map, dgate_fn_map
 class GemmDActMixin(GemmActMixin):
@@ -51,36 +51,30 @@ class GemmDActMixin(GemmActMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
         assert tRS_rC is not None
         # We don't add C to the accumulator
         GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC=None)
-        tRS_rC_acc = cute.make_rmem_tensor_like(tRS_rC, self.acc_dtype)
-        tRS_rC_acc.store(tRS_rC.load().to(self.acc_dtype))
         # If we don't have .shape here, the compiler generates local stores and loads
         if const_expr(params.act_fn is not None):
-            tRS_rPostAct = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype)
-            if const_expr(self.arch < 100):
-                for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
-                    tRS_rD[i], tRS_rPostAct[i] = params.act_fn(tRS_rC_acc[i], tRS_rD[i])
-            else:
-                for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True):
-                    (
-                        (tRS_rD[2 * i], tRS_rD[2 * i + 1]),
-                        (tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1]),
-                    ) = params.act_fn(
-                        (tRS_rC_acc[2 * i], tRS_rC_acc[2 * i + 1]),
-                        (tRS_rD[2 * i], tRS_rD[2 * i + 1]),
-                    )
         else:
-            tRS_rPostAct = tRS_rC_acc
-        return tRS_rPostAct
 class GemmDActSm90(GemmDActMixin, GemmSm90):
     pass
 class GemmDActSm100(GemmDActMixin, GemmSm100):
     pass
@@ -92,17 +86,18 @@ class GemmDActSm120(GemmDActMixin, GemmSm120):
 class GemmDGatedMixin(GemmActMixin):
     # Different from GemmActMixin, here act_bwd_fn must take in 3 arguments (x, y, dout)
     # and return 3 arguments (dx, dy, out)
-    _epi_ops = (*GemmActMixin._epi_ops, ColVecReduce("mColVecReduce"))
     _extra_param_fields = (("act_bwd_fn", cutlass.Constexpr, None),)
-    _epi_param_bases = (ParamsBase,)
     @mlir_namedtuple
     class EpilogueArguments(NamedTuple):
-        mPostAct: cute.Tensor
         act_bwd_fn: cutlass.Constexpr[Callable] = None
-        alpha: Optional[Float32 | cute.Tensor] = None
-        beta: Optional[Float32 | cute.Tensor] = None
-        mRowVecBroadcast: Optional[cute.Tensor] = None
         mColVecBroadcast: Optional[cute.Tensor] = None
         mColVecReduce: Optional[cute.Tensor] = None
         rounding_mode: cutlass.Constexpr[int] = RoundingMode.RN
@@ -117,9 +112,9 @@ class GemmDGatedMixin(GemmActMixin):
         assert self.d_dtype.width == 32, "D storage type must be 32 bit"
         assert self.c_dtype.width == 32, "C storage type must be 32 bit"
         self.rounding_mode = args.rounding_mode
-        self.postact_dtype = args.mPostAct.element_type
-        self.postact_layout = cutlass.utils.LayoutEnum.from_tensor(args.mPostAct)
-        self.cta_tile_shape_postact_mn = self.cta_tile_shape_mnk[:2]
         d = self._epi_ops_to_params_dict(args)
         d["act_bwd_fn"] = args.act_bwd_fn
         return self.EpilogueParams(**d)
@@ -133,13 +128,9 @@ class GemmDGatedMixin(GemmActMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
-        alpha = epi_loop_tensors["alpha"]
-        beta = epi_loop_tensors["beta"]
-        tDrRowVec = epi_loop_tensors["mRowVecBroadcast"]
-        tDrColVec = epi_loop_tensors["mColVecBroadcast"]
-        tDrColVecReduce = epi_loop_tensors["mColVecReduce"]
-        assert alpha is None and beta is None and tDrRowVec is None  # We don't use these for now
         assert tRS_rC is not None
         implicit_dtype = self.implicit_dtype
         assert implicit_dtype.width == 16, "GemmDGatedMixin only supports 16bit for now"
@@ -150,7 +141,7 @@ class GemmDGatedMixin(GemmActMixin):
         tRS_rOut = cute.make_rmem_tensor_like(tRS_rD, Float32)
         tRS_rD_scaled = cute.make_rmem_tensor_like(tRS_rD)
         if const_expr(tDrColVec is not None):  # Scale D by colvec
-            if const_expr(self.arch < 100):
                 tRS_rD_scaled.store(tRS_rD.load() * tDrColVec.load().to(tRS_rD.element_type))
             else:
                 tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout)
@@ -159,63 +150,45 @@ class GemmDGatedMixin(GemmActMixin):
                     tRS_rD_scaled, tDrColVec.layout
                 )
                 for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True):
                     for n in cutlass.range(
-                        cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True
                     ):
-                        (
-                            tRS_rD_scaled_mn[m, 2 * n],
-                            tRS_rD_scaled_mn[m, 2 * n + 1],
-                        ) = cute.arch.mul_packed_f32x2(
-                            (tRS_rD_mn[m, 2 * n], tRS_rD_mn[m, 2 * n + 1]),
-                            (tDrColVec_mn[m, 0], tDrColVec_mn[m, 0]),
-                        )
         else:
             tRS_rD_scaled.store(tRS_rD.load())
-        if const_expr(self.arch < 100):
-            for i in cutlass.range(cute.size(tRS_rD)):
-                (
-                    tRS_rdXY_f32x2[2 * i],
-                    tRS_rdXY_f32x2[2 * i + 1],
-                    tRS_rOut[i],
-                ) = params.act_bwd_fn(
-                    tRS_rXY_f32x2[2 * i], tRS_rXY_f32x2[2 * i + 1], tRS_rD_scaled[i]
-                )
-        else:
-            for i in cutlass.range(cute.size(tRS_rD) // 2):
-                (
-                    (tRS_rdXY_f32x2[4 * i], tRS_rdXY_f32x2[4 * i + 2]),
-                    (tRS_rdXY_f32x2[4 * i + 1], tRS_rdXY_f32x2[4 * i + 3]),
-                    (tRS_rOut[2 * i], tRS_rOut[2 * i + 1]),
-                ) = params.act_bwd_fn(
-                    (tRS_rXY_f32x2[4 * i], tRS_rXY_f32x2[4 * i + 2]),
-                    (tRS_rXY_f32x2[4 * i + 1], tRS_rXY_f32x2[4 * i + 3]),
-                    (tRS_rD_scaled[2 * i], tRS_rD_scaled[2 * i + 1]),
-                )
         if const_expr(tDrColVecReduce is not None):
             # Accumulate postact * dout before D is scaled by colvec_scale
             colvec_reduce_accumulate(self, tDrColVecReduce, tRS_rOut, rScale=tRS_rD)
         if const_expr(tDrColVec is not None):  # Scale Out by colvec
-            if const_expr(self.arch < 100):
                 tRS_rOut.store(tRS_rOut.load() * tDrColVec.load().to(tRS_rD.element_type))
             else:
                 tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout)
                 tRS_rOut_mn = layout_utils.convert_layout_zero_stride(tRS_rOut, tDrColVec.layout)
                 for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True):
                     for n in cutlass.range(
-                        cute.size(tDrColVec_mn, mode=[1]) // 2, unroll_full=True
                     ):
-                        tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1] = (
-                            cute.arch.mul_packed_f32x2(
-                                (tRS_rOut_mn[m, 2 * n], tRS_rOut_mn[m, 2 * n + 1]),
-                                (tDrColVec_mn[m, 0], tDrColVec_mn[m, 0]),
-                            )
-                        )
         # Type conversion
         tRS_rdXY_f16x2 = cute.make_rmem_tensor(tRS_rdXY_f32x2.layout, implicit_dtype)
         tRS_rdXY_f16x2.store(tRS_rdXY_f32x2.load().to(implicit_dtype))
         tRS_rD.store(cute.recast_tensor(tRS_rdXY_f16x2, Float32).load())
-        return tRS_rOut
     # epi_end is inherited from ComposableEpiMixin → delegates to ColVecReduce.end()
@@ -224,6 +197,10 @@ class GemmDGatedSm90(GemmDGatedMixin, GemmSm90):
     pass
 class GemmDGatedSm100(GemmDGatedMixin, GemmSm100):
     pass
@@ -263,16 +240,21 @@ def _compile_gemm_dact(
 ):
     is_dgated = gemm_cls_name == "dgated"
     sm_to_cls = {
-        "dact": {9: GemmDActSm90, 10: GemmDActSm100, 11: GemmDActSm100, 12: GemmDActSm120},
         "dgated": {
             9: GemmDGatedSm90,
             10: GemmDGatedSm100,
             11: GemmDGatedSm100,
             12: GemmDGatedSm120,
         },
     }
-    if device_capacity[0] == 12 and gemm_cls_name == "dact":
-        raise NotImplementedError("SM120 non-gated dactivation GEMM epilogue is not yet supported")
     GemmCls = sm_to_cls[gemm_cls_name][device_capacity[0]]
     mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors(
         a_dtype,
@@ -289,7 +271,7 @@ def _compile_gemm_dact(
     div_pa = div_for_dtype(postact_dtype)
     pa_leading = 1 if postact_major == "n" else 0
     pa_shape = (m, n) if varlen_m else (m, n, l)
-    mPostAct = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading, divisibility=div_pa)
     if is_dgated:
         act_fn = dgate_fn_map[activation]
@@ -316,7 +298,7 @@ def _compile_gemm_dact(
                 divisibility=1,
             )
         epi_args = GemmCls.EpilogueArguments(
-            mPostAct,
             act_fn,
             mColVecBroadcast=mColVec,
             mColVecReduce=mColVecReduce,
@@ -328,7 +310,7 @@ def _compile_gemm_dact(
         post_init = _set_implicit_dtype
     else:
         act_fn = dact_fn_map[activation]
-        epi_args = GemmCls.EpilogueArguments(mPostAct, act_fn)
         post_init = None
     scheduler_args = make_fake_scheduler_args(
@@ -369,6 +351,7 @@ def gemm_dact(
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
     pingpong: bool = True,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
@@ -432,7 +415,9 @@ def gemm_dact(
     postact_dtype = torch2cute_dtype_map[PostAct.dtype]
     device_capacity = get_device_capacity(A.device)
-    assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported"
     if is_dynamic_persistent and device_capacity[0] == 9:
         assert tile_count_semaphore is not None, (
@@ -451,7 +436,7 @@ def gemm_dact(
         d_major,
         c_major,
         postact_major,
-        (tile_M, tile_N),
         (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
@@ -468,11 +453,6 @@ def gemm_dact(
         use_tma_gather=use_tma_gather,
     )
-    from .cache_utils import COMPILE_ONLY
-    if COMPILE_ONLY:
-        return
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     if is_dgated:
         epi_args = GemmDGatedMixin.EpilogueArguments(
@@ -498,11 +478,9 @@ def gemm_dact(
     varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx)
     if device_capacity[0] in [10, 11]:
-        compiled_fn(
-            A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args, None, None, None
-        )
     else:
-        compiled_fn(A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args, None)
 gemm_dgated = gemm_dact

 import cutlass
 import cutlass.cute as cute
 from cutlass import Int32, Float32, const_expr
+from .gemm_sm80 import GemmSm80
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
 from .gemm_default_epi import GemmDefaultEpiMixin
 from .gemm_act import GemmActMixin
+from .epi_ops import ColVecLoad, ColVecReduce, Scalar, TileStore, colvec_reduce_accumulate
 from .compile_utils import make_fake_tensor as fake_tensor
 from .cute_dsl_utils import (
     mlir_namedtuple,
     torch2cute_dtype_map,
     get_device_capacity,
     make_fake_gemm_tensors,
     compile_gemm_kernel,
 )
+from .cache import jit_cache
 from .rounding import RoundingMode
 from .activation import dact_fn_map, dgate_fn_map
+from . import layout_utils
 class GemmDActMixin(GemmActMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
         assert tRS_rC is not None
         # We don't add C to the accumulator
         GemmDefaultEpiMixin.epi_visit_subtile(self, params, epi_loop_tensors, tRS_rD, tRS_rC=None)
+        tRS_rC_acc = tRS_rC.to(self.acc_dtype)
         # If we don't have .shape here, the compiler generates local stores and loads
         if const_expr(params.act_fn is not None):
+            tRS_rAuxOut = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype)
+            vectorize = const_expr(self.arch == 100)
+            for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize):
+                tRS_rD[i], tRS_rAuxOut[i] = params.act_fn(tRS_rC_acc[i], tRS_rD[i])
         else:
+            tRS_rAuxOut = tRS_rC_acc
+        return (tRS_rAuxOut,)
 class GemmDActSm90(GemmDActMixin, GemmSm90):
     pass
+class GemmDActSm80(GemmDActMixin, GemmSm80):
+    pass
 class GemmDActSm100(GemmDActMixin, GemmSm100):
     pass
 class GemmDGatedMixin(GemmActMixin):
     # Different from GemmActMixin, here act_bwd_fn must take in 3 arguments (x, y, dout)
     # and return 3 arguments (dx, dy, out)
+    _epi_ops = (
+        ColVecLoad("mColVecBroadcast"),
+        Scalar("sr_seed", dtype=Int32),
+        TileStore("mAuxOut"),
+        ColVecReduce("mColVecReduce"),
+    )
     _extra_param_fields = (("act_bwd_fn", cutlass.Constexpr, None),)
     @mlir_namedtuple
     class EpilogueArguments(NamedTuple):
+        mAuxOut: cute.Tensor
         act_bwd_fn: cutlass.Constexpr[Callable] = None
         mColVecBroadcast: Optional[cute.Tensor] = None
         mColVecReduce: Optional[cute.Tensor] = None
         rounding_mode: cutlass.Constexpr[int] = RoundingMode.RN
         assert self.d_dtype.width == 32, "D storage type must be 32 bit"
         assert self.c_dtype.width == 32, "C storage type must be 32 bit"
         self.rounding_mode = args.rounding_mode
+        self.aux_out_dtype = args.mAuxOut.element_type
+        self.aux_out_layout = cutlass.utils.LayoutEnum.from_tensor(args.mAuxOut)
+        self.cta_tile_shape_aux_out_mn = self.cta_tile_shape_mnk[:2]
         d = self._epi_ops_to_params_dict(args)
         d["act_bwd_fn"] = args.act_bwd_fn
         return self.EpilogueParams(**d)
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
+        tDrColVec = epi_loop_tensors.get("mColVecBroadcast")
+        tDrColVecReduce = epi_loop_tensors.get("mColVecReduce")
         assert tRS_rC is not None
         implicit_dtype = self.implicit_dtype
         assert implicit_dtype.width == 16, "GemmDGatedMixin only supports 16bit for now"
         tRS_rOut = cute.make_rmem_tensor_like(tRS_rD, Float32)
         tRS_rD_scaled = cute.make_rmem_tensor_like(tRS_rD)
         if const_expr(tDrColVec is not None):  # Scale D by colvec
+            if const_expr(self.arch != 100):
                 tRS_rD_scaled.store(tRS_rD.load() * tDrColVec.load().to(tRS_rD.element_type))
             else:
                 tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout)
                     tRS_rD_scaled, tDrColVec.layout
                 )
                 for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True):
+                    scale = tDrColVec_mn[m, 0]
                     for n in cutlass.range(
+                        cute.size(tDrColVec_mn, mode=[1]), unroll_full=True, vectorize=True
                     ):
+                        tRS_rD_scaled_mn[m, n] = tRS_rD_mn[m, n] * scale
         else:
             tRS_rD_scaled.store(tRS_rD.load())
+        tRS_rXY_pair = cute.flat_divide(tRS_rXY_f32x2, cute.make_layout(2))
+        tRS_rX = tRS_rXY_pair[0, ...]
+        tRS_rY = tRS_rXY_pair[1, ...]
+        tRS_rdXY_pair = cute.flat_divide(tRS_rdXY_f32x2, cute.make_layout(2))
+        tRS_rdX = tRS_rdXY_pair[0, ...]
+        tRS_rdY = tRS_rdXY_pair[1, ...]
+        vectorize = const_expr(self.arch == 100)
+        for i in cutlass.range(cute.size(tRS_rD), vectorize=vectorize):
+            tRS_rdX[i], tRS_rdY[i], tRS_rOut[i] = params.act_bwd_fn(
+                tRS_rX[i], tRS_rY[i], tRS_rD_scaled[i]
+            )
         if const_expr(tDrColVecReduce is not None):
             # Accumulate postact * dout before D is scaled by colvec_scale
             colvec_reduce_accumulate(self, tDrColVecReduce, tRS_rOut, rScale=tRS_rD)
         if const_expr(tDrColVec is not None):  # Scale Out by colvec
+            if const_expr(self.arch != 100):
                 tRS_rOut.store(tRS_rOut.load() * tDrColVec.load().to(tRS_rD.element_type))
             else:
                 tDrColVec_mn = layout_utils.convert_layout_zero_stride(tDrColVec, tDrColVec.layout)
                 tRS_rOut_mn = layout_utils.convert_layout_zero_stride(tRS_rOut, tDrColVec.layout)
                 for m in cutlass.range(cute.size(tDrColVec_mn, mode=[0]), unroll_full=True):
+                    scale = tDrColVec_mn[m, 0]
                     for n in cutlass.range(
+                        cute.size(tDrColVec_mn, mode=[1]), unroll_full=True, vectorize=True
                     ):
+                        tRS_rOut_mn[m, n] = tRS_rOut_mn[m, n] * scale
         # Type conversion
         tRS_rdXY_f16x2 = cute.make_rmem_tensor(tRS_rdXY_f32x2.layout, implicit_dtype)
         tRS_rdXY_f16x2.store(tRS_rdXY_f32x2.load().to(implicit_dtype))
         tRS_rD.store(cute.recast_tensor(tRS_rdXY_f16x2, Float32).load())
+        return (tRS_rOut,)
     # epi_end is inherited from ComposableEpiMixin → delegates to ColVecReduce.end()
     pass
+class GemmDGatedSm80(GemmDGatedMixin, GemmSm80):
+    pass
 class GemmDGatedSm100(GemmDGatedMixin, GemmSm100):
     pass
 ):
     is_dgated = gemm_cls_name == "dgated"
     sm_to_cls = {
+        "dact": {
+            8: GemmDActSm80,
+            9: GemmDActSm90,
+            10: GemmDActSm100,
+            11: GemmDActSm100,
+            12: GemmDActSm120,
+        },
         "dgated": {
+            8: GemmDGatedSm80,
             9: GemmDGatedSm90,
             10: GemmDGatedSm100,
             11: GemmDGatedSm100,
             12: GemmDGatedSm120,
         },
     }
     GemmCls = sm_to_cls[gemm_cls_name][device_capacity[0]]
     mA, mB, mD, mC, m, n, k, l = make_fake_gemm_tensors(
         a_dtype,
     div_pa = div_for_dtype(postact_dtype)
     pa_leading = 1 if postact_major == "n" else 0
     pa_shape = (m, n) if varlen_m else (m, n, l)
+    mAuxOut = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading, divisibility=div_pa)
     if is_dgated:
         act_fn = dgate_fn_map[activation]
                 divisibility=1,
             )
         epi_args = GemmCls.EpilogueArguments(
+            mAuxOut,
             act_fn,
             mColVecBroadcast=mColVec,
             mColVecReduce=mColVecReduce,
         post_init = _set_implicit_dtype
     else:
         act_fn = dact_fn_map[activation]
+        epi_args = GemmCls.EpilogueArguments(mAuxOut, act_fn)
         post_init = None
     scheduler_args = make_fake_scheduler_args(
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
+    tile_K: int | None = None,
     pingpong: bool = True,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
     postact_dtype = torch2cute_dtype_map[PostAct.dtype]
     device_capacity = get_device_capacity(A.device)
+    assert device_capacity[0] in [8, 9, 10, 11, 12], (
+        "Only SM8x, SM90, SM100, SM110, and SM120 are supported"
+    )
     if is_dynamic_persistent and device_capacity[0] == 9:
         assert tile_count_semaphore is not None, (
         d_major,
         c_major,
         postact_major,
+        (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N),
         (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
         use_tma_gather=use_tma_gather,
     )
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     if is_dgated:
         epi_args = GemmDGatedMixin.EpilogueArguments(
     varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx)
     if device_capacity[0] in [10, 11]:
+        compiled_fn(A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args, None, None)
     else:
+        compiled_fn(A_p, B_p, Out_p, PreAct_p, epi_args, scheduler_args, varlen_args)
 gemm_dgated = gemm_dact

build/torch-cuda/quack/gemm_default_epi.py CHANGED Viewed

@@ -1,5 +1,5 @@
 # Copyright (c) 2025, Wentao Guo, Tri Dao.
-from typing import NamedTuple, Optional
 import cutlass
 import cutlass.cute as cute
@@ -8,6 +8,7 @@ from cutlass import Int32, Float32, const_expr
 from .cute_dsl_utils import mlir_namedtuple
 from .epi_composable import ComposableEpiMixin
 from .epi_ops import Scalar, RowVecLoad, ColVecLoad
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
@@ -41,7 +42,7 @@ class GemmDefaultEpiMixin(ComposableEpiMixin):
         self.rounding_mode = args.rounding_mode
         d = self._epi_ops_to_params_dict(args)
         for key in ("mRowVecBroadcast", "mColVecBroadcast"):
-            if key in self.concat_layout and key in d and d[key] is not None:
                 d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
@@ -52,11 +53,18 @@ class GemmDefaultEpiMixin(ComposableEpiMixin):
         epi_loop_tensors,
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
-        alpha = epi_loop_tensors["alpha"]
-        beta = epi_loop_tensors["beta"]
-        tDrRowVec = epi_loop_tensors["mRowVecBroadcast"]
-        tDrColVec = epi_loop_tensors["mColVecBroadcast"]
         rD = tRS_rD.load()
         # Apply alpha scaling to accumulator if alpha is provided (not None)
         if const_expr(hasattr(params, "alpha") and params.alpha is not None):
@@ -77,27 +85,11 @@ class GemmDefaultEpiMixin(ComposableEpiMixin):
         if const_expr(tDrColVec is not None):
             for i in cutlass.range(cute.size(tDrColVec), unroll_full=True):
                 tRS_rD[i] += tDrColVec[i]
-        return None
-    def epi_setup_postact(
-        self,
-        params,
-        epi_smem_tensors,
-        tiled_copy_r2s,
-        tiled_copy_t2r,
-        tile_coord_mnkl,
-        varlen_manager,
-        tidx,
-    ):
-        """Returns None — default epilogue has no postact output."""
-        return None
-    @cute.jit
-    def epi_convert_postact(
-        self, tRS_rPostAct, sr_seed, tidx, tile_coord_mnkl, num_prev_subtiles, epi_idx
-    ):
-        """Convert postact from acc_dtype to output dtype. Override for custom postprocessing."""
-        return tRS_rPostAct
 class GemmDefaultSm90(GemmDefaultEpiMixin, GemmSm90):

 # Copyright (c) 2025, Wentao Guo, Tri Dao.
+from typing import NamedTuple, Optional, Tuple
 import cutlass
 import cutlass.cute as cute
 from .cute_dsl_utils import mlir_namedtuple
 from .epi_composable import ComposableEpiMixin
 from .epi_ops import Scalar, RowVecLoad, ColVecLoad
+from .gemm_sm80 import GemmSm80
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
         self.rounding_mode = args.rounding_mode
         d = self._epi_ops_to_params_dict(args)
         for key in ("mRowVecBroadcast", "mColVecBroadcast"):
+            if key in self.concat_layout and key in d:
                 d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
         epi_loop_tensors,
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
+        """Return a tuple of register tensors (one per aux output).
+        The returned tuple must be the same length as the tuple returned
+        from :meth:`epi_setup_aux_out`. The default impl returns ``()`` —
+        no aux outputs.
+        """
+        # Use .get(): inactive ops are filtered out of epi_loop_tensors.
+        alpha = epi_loop_tensors.get("alpha")
+        beta = epi_loop_tensors.get("beta")
+        tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast")
+        tDrColVec = epi_loop_tensors.get("mColVecBroadcast")
         rD = tRS_rD.load()
         # Apply alpha scaling to accumulator if alpha is provided (not None)
         if const_expr(hasattr(params, "alpha") and params.alpha is not None):
         if const_expr(tDrColVec is not None):
             for i in cutlass.range(cute.size(tDrColVec), unroll_full=True):
                 tRS_rD[i] += tDrColVec[i]
+        return ()
+class GemmDefaultSm80(GemmDefaultEpiMixin, GemmSm80):
+    pass
 class GemmDefaultSm90(GemmDefaultEpiMixin, GemmSm90):

build/torch-cuda/quack/gemm_interface.py CHANGED Viewed

@@ -3,7 +3,7 @@ from typing import Optional, Tuple, Literal
 from functools import partial
 import torch
-from ._ops_compat import add_quack_op_namespace_prefix
 import torch.nn.functional as F
 from torch import Tensor
@@ -21,12 +21,49 @@ from .rms_final_reduce import rms_final_reduce
 from .rounding import RoundingMode
 # Dictionary mapping activation names to PyTorch functions
 act_to_pytorch_fn_map = {
     None: lambda x: x,
     "relu": F.relu,
     "relu_sq": lambda x: F.relu(x).square(),
     "gelu_tanh_approx": partial(F.gelu, approximate="tanh"),
 }
@@ -34,22 +71,37 @@ act_to_pytorch_fn_map = {
 # Each function takes (gate, up) and returns postact
 gated_to_pytorch_fn_map = {
     "swiglu": lambda gate, up: F.silu(gate) * up,
     "swiglu_oai": lambda gate, up: gate * torch.sigmoid(1.702 * gate) * (up + 1),
     "reglu": lambda gate, up: F.relu(gate) * up,
     "geglu": lambda gate, up: F.gelu(gate, approximate="tanh") * up,
     "glu": lambda gate, up: torch.sigmoid(gate) * up,
 }
-ActActivation = Literal[None, "relu", "relu_sq", "gelu_tanh_approx"]
-GatedActivation = Literal["swiglu", "swiglu_oai", "reglu", "geglu", "glu"]
 Activation = Literal[
     None,
     "relu",
     "relu_sq",
     "gelu_tanh_approx",
     "swiglu",
     "swiglu_oai",
     "reglu",
     "geglu",
     "glu",
@@ -68,9 +120,88 @@ def _concat_interleave_bias(t):
     return t.unflatten(-1, (2, half)).transpose(-2, -1).flatten(-2, -1)
 def default_config(device):
     cap = get_device_capacity(device)[0]
-    if cap in [10, 11]:
         return GemmConfig(
             tile_m=256,
             tile_n=256,
@@ -101,6 +232,31 @@ def default_config(device):
         )
 def nvmmh_config(A, B, device_capacity):
     """Use nvMatmulHeuristics to pick a config for pure GEMM (no varlen/gather/epilogue).
@@ -130,6 +286,21 @@ def prune_invalid_gemm_configs(configs, named_args: dict, **kwargs):
     # use_tma_gather only valid when gather_A is active on SM100/SM110
     if not gather_A or device_capacity not in [10, 11]:
         configs = [conf for conf in configs if not conf.kwargs["config"].use_tma_gather]
     return configs
@@ -157,26 +328,39 @@ def gemm_tuned(
     rounding_mode: int = RoundingMode.RN,
     sr_seed: int | Tensor = 0,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> None:
     if config is None:
-        # Use nvMMH heuristic for pure GEMM (no varlen, no gather, no epilogue)
-        is_pure_gemm = (
-            cu_seqlens_m is None
-            and cu_seqlens_k is None
-            and A_idx is None
-            and C is None
-            and bias is None
-            and not add_to_output
-        )
-        if is_pure_gemm:
-            device_capacity = get_device_capacity(A.device)[0]
-            config = nvmmh_config(A, B, device_capacity)
-        if config is None:
-            config = default_config(A.device)
     varlen_m = cu_seqlens_m is not None
     varlen_k = cu_seqlens_k is not None
     varlen = varlen_m or varlen_k
     gather_A = A_idx is not None
     if gather_A:
         assert varlen, "gather_A requires either varlen_m or varlen_k"
         assert config.cluster_n == 1, "gather_A requires cluster_n=1"
@@ -235,7 +419,8 @@ def gemm_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -252,6 +437,10 @@ def gemm_tuned(
         sr_seed=sr_seed,
         use_tma_gather=config.use_tma_gather,
         concat_layout=swapped_concat,
     )
@@ -274,12 +463,23 @@ def gemm_act_tuned(
     A_idx: Optional[Tensor] = None,  # (total_M,) if gather_A with varlen_m
     dynamic_scheduler: bool = False,
     config: Optional[GemmConfig] = None,
 ) -> None:
     if config is None:
-        config = default_config(A.device)
     varlen_m = cu_seqlens_m is not None
     if varlen_m:
         assert not config.swap_ab, "Variable-length sequences not supported with swap_ab"
     if A.ndim == 2 and not varlen_m:
         A = A.unsqueeze(0)  # (1, M, K)
     B = B.mT  # (N, K) or (L, N, K)
@@ -315,7 +515,8 @@ def gemm_act_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -324,6 +525,8 @@ def gemm_act_tuned(
         cu_seqlens_m=cu_seqlens_m,
         A_idx=A_idx,
         use_tma_gather=config.use_tma_gather,
     )
@@ -383,7 +586,8 @@ def gemm_dact_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -395,8 +599,10 @@ def gemm_dact_tuned(
 def gemm(
     # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k
-    A: Tensor,
-    B: Tensor,  # (K, N) or (L, K, N) or (total_K, N) if varlen_k
     out: Optional[Tensor] = None,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m
     bias: Optional[Tensor] = None,  # (N,) or (L, N)
     alpha: float | Tensor = 1.0,
@@ -412,8 +618,16 @@ def gemm(
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> Tensor:
     """GEMM with optional output tensor and tuning control."""
     if out is None:
-        out_dtype = A.dtype if out_dtype is None else out_dtype
         varlen_m = cu_seqlens_m is not None
         varlen_k = cu_seqlens_k is not None
         if varlen_m:
@@ -428,6 +642,15 @@ def gemm(
                 (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1])
             )
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha = alpha if isinstance(alpha, float) else 1.0
     sr_seed_tensor = sr_seed if isinstance(sr_seed, Tensor) else None
@@ -450,12 +673,14 @@ def gemm(
         sr_seed=sr_seed_int,
         sr_seed_tensor=sr_seed_tensor,
         concat_layout=concat_str,
     )
     return out
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_out"),
     mutates_args=("out",),
     device_types="cuda",
     # We have to split out alpha and alpha_tensor since torch.library requires
@@ -480,11 +705,18 @@ def gemm_out(
     sr_seed: int = 0,
     sr_seed_tensor: Optional[Tensor] = None,
     concat_layout: Optional[str] = None,
 ) -> None:
     """GEMM with pre-allocated output tensor."""
     fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None)
-    alpha = alpha_tensor if alpha_tensor is not None else alpha
-    sr_seed_arg = sr_seed_tensor if sr_seed_tensor is not None else sr_seed
     fn(
         A,
         B,
@@ -499,7 +731,9 @@ def gemm_out(
         dynamic_scheduler=dynamic_scheduler,
         rounding_mode=rounding_mode,
         sr_seed=sr_seed_arg,
-        concat_layout=tuple(concat_layout.split(",")) if concat_layout else None,
     )
@@ -572,10 +806,43 @@ def gemm_ref(
     return out
 def gemm_add(
     # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k
-    A: Tensor,
-    B: Tensor,  # (K, N) or (L, K, N) or (total_K, N) if varlen_k
     C: Tensor,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m or (L, M, N) if varlen_k
     out: Optional[Tensor] = None,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m
     alpha: float | Tensor = 1.0,
@@ -590,8 +857,15 @@ def gemm_add(
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> Tensor:
     """GEMM with addition and optional output tensor."""
     if out is None:
-        out_dtype = A.dtype if out_dtype is None else out_dtype
         varlen_m = cu_seqlens_m is not None
         varlen_k = cu_seqlens_k is not None
         if varlen_m:
@@ -608,17 +882,26 @@ def gemm_add(
             )
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     add_to_output = C is out and isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None
     alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha = alpha if isinstance(alpha, float) else 1.0
     beta_tensor = beta if not isinstance(beta, float) else None
     beta = beta if isinstance(beta, float) else 1.0
-    alpha_arg = alpha_tensor if alpha_tensor is not None else alpha
-    beta_arg = beta_tensor if beta_tensor is not None else beta
     concat_str = ",".join(concat_layout) if concat_layout else None
     if add_to_output:
         gemm_add_inplace(
-            A,
-            B,
             out,
             alpha=alpha_arg,
             beta=beta_arg,
@@ -648,12 +931,14 @@ def gemm_add(
             dynamic_scheduler=dynamic_scheduler,
             tuned=tuned,
             concat_layout=concat_str,
         )
     return out
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_add_out"),
     mutates_args=("out",),
     device_types="cuda",
     # We have to split out alpha and alpha_tensor since torch.library requires
@@ -678,11 +963,13 @@ def gemm_add_out(
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     concat_layout: Optional[str] = None,
 ) -> None:
     """GEMM with addition and pre-allocated output tensor."""
     fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None)
-    alpha = alpha_tensor if alpha_tensor is not None else alpha
-    beta = beta_tensor if beta_tensor is not None else beta
     fn(
         A,
         B,
@@ -690,13 +977,15 @@ def gemm_add_out(
         C,
         alpha=alpha,
         beta=beta,
         cu_seqlens_m=cu_seqlens_m,
         cu_seqlens_k=cu_seqlens_k,
         A_idx=A_idx,
         batch_idx_permute=batch_idx_permute,
         add_to_output=add_to_output,
         dynamic_scheduler=dynamic_scheduler,
-        concat_layout=tuple(concat_layout.split(",")) if concat_layout else None,
     )
@@ -783,8 +1072,9 @@ def gemm_add_ref(
 def gemm_add_inplace(
     # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k
-    A: Tensor,
-    B: Tensor,  # (K, N) or (L, K, N) or (total_K, N) if varlen_k
     out: Tensor,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m or (L, M, N) if varlen_k
     alpha: float | Tensor = 1.0,
     beta: float | Tensor = 1.0,
@@ -808,10 +1098,24 @@ def gemm_add_inplace(
         dynamic_scheduler: Whether to use dynamic scheduler
         tuned: Whether to use autotuned configuration
     """
     alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha = alpha if isinstance(alpha, float) else 1.0
     beta_tensor = beta if not isinstance(beta, float) else None
     beta = beta if isinstance(beta, float) else 1.0
     gemm_add_inplace_op(
         A,
         B,
@@ -829,11 +1133,13 @@ def gemm_add_inplace(
         concat_layout=",".join(concat_layout)
         if isinstance(concat_layout, tuple)
         else concat_layout,
     )
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_add_inplace"),
     mutates_args=("out",),
     device_types="cuda",
     # We have to split out alpha and alpha_tensor since torch.library requires
@@ -856,10 +1162,12 @@ def gemm_add_inplace_op(
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     concat_layout: Optional[str] = None,
 ) -> None:
     fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None)
-    alpha = alpha_tensor if alpha_tensor is not None else alpha
-    beta = beta_tensor if beta_tensor is not None else beta
     add_to_output = isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None
     # Use out as both input bias and output
     fn(
@@ -875,13 +1183,19 @@ def gemm_add_inplace_op(
         batch_idx_permute=batch_idx_permute,
         add_to_output=add_to_output,
         dynamic_scheduler=dynamic_scheduler,
-        concat_layout=tuple(concat_layout.split(",")) if concat_layout else None,
     )
 def gemm_act(
-    A: Tensor,  # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m
-    B: Tensor,  # (K, N) or (L, K, N)
     C: Optional[Tensor] = None,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m
     bias: Optional[Tensor] = None,  # (N,) or (L, N)
     activation: Activation = None,
@@ -897,9 +1211,16 @@ def gemm_act(
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> Tuple[Optional[Tensor], Tensor]:
     """GEMM with activation (or gated activation) and optional output tensors."""
     is_gated = activation in gated_to_pytorch_fn_map
-    out_dtype = A.dtype if out_dtype is None else out_dtype
-    postact_dtype = A.dtype if postact_dtype is None else postact_dtype
     varlen_m = cu_seqlens_m is not None
     # Determine output shape based on gather_A
     if varlen_m:
@@ -914,6 +1235,14 @@ def gemm_act(
         preact_out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     if postact_out is None:
         postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device)
     concat_str = ",".join(concat_layout) if concat_layout else None
     if is_gated:
         gemm_gated_out(
@@ -929,6 +1258,8 @@ def gemm_act(
             dynamic_scheduler,
             tuned,
             concat_layout=concat_str,
         )
     else:
         gemm_act_out(
@@ -943,6 +1274,8 @@ def gemm_act(
             A_idx,
             dynamic_scheduler,
             tuned,
         )
     return preact_out, postact_out
@@ -951,10 +1284,10 @@ gemm_gated = gemm_act
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_act_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
-    schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True) -> ()",
 )
 def gemm_act_out(
     A: Tensor,  # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m
@@ -968,10 +1301,25 @@ def gemm_act_out(
     A_idx: Optional[Tensor] = None,  # (total_M,) if gather_A with varlen_m
     dynamic_scheduler: bool = False,
     tuned: bool = True,
 ) -> None:
     """GEMM with activation and pre-allocated output tensors."""
     fn = gemm_act_tuned if tuned else partial(gemm_act_tuned.fn, config=None)
-    fn(A, B, preact_out, postact_out, C, bias, activation, cu_seqlens_m, A_idx, dynamic_scheduler)
 def gemm_act_ref(
@@ -1048,6 +1396,16 @@ def gemm_dact(
         dx_out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     if postact_out is None:
         postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device)
     if is_dgated:
         colvec_reduce_final = gemm_dgated_out(
             A,
@@ -1063,10 +1421,10 @@ def gemm_dact(
             dynamic_scheduler,
             tuned,
         )
-        if not colvec_reduce:
-            return dx_out, postact_out
-        else:
-            return dx_out, postact_out, colvec_reduce_final
     else:
         gemm_dact_out(
             A,
@@ -1080,14 +1438,15 @@ def gemm_dact(
             dynamic_scheduler,
             tuned,
         )
-        return dx_out, postact_out
 gemm_dgated = gemm_dact
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_dact_out"),
     mutates_args=("dx_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a3!) dx_out, Tensor(a4!) postact_out, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> ()",
@@ -1152,11 +1511,27 @@ def gemm_dact_ref(
 gemm_dgated_ref = gemm_dact_ref
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_symmetric_out"),
     mutates_args=("out",),
     device_types="cuda",
-    schema="(Tensor A, Tensor B, Tensor(a2!) out, Tensor? C=None, bool dynamic_scheduler=False, float alpha=1.0, float beta=1.0) -> ()",
 )
 def gemm_symmetric_out(
     A: Tensor,  # (M, K) or (L, M, K)
@@ -1166,8 +1541,12 @@ def gemm_symmetric_out(
     dynamic_scheduler: bool = False,
     alpha: float = 1.0,
     beta: float = 1.0,
 ) -> None:
     """GEMM with guaranteed symmetric output."""
     if A.ndim == 2:
         A = A.unsqueeze(0)  # (1, M, K)
     B = B.mT  # (M, K) or (L, M, K)
@@ -1184,12 +1563,7 @@ def gemm_symmetric_out(
     )
     sm = get_device_capacity(A.device)[0]
     # We want square tile per cluster
-    tile_m, tile_n, cluster_m, pingpong = {
-        9: (128, 256, 2, False),
-        10: (256, 256, 2, False),
-        11: (256, 256, 2, False),
-        12: (128, 128, 1, True),
-    }[sm]
     gemm_symmetric_dispatch(
         A,
         B,
@@ -1229,11 +1603,29 @@ def gemm_symmetric(
     if out is None:
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     alpha_val = alpha if isinstance(alpha, float) else 1.0
     beta_val = beta if isinstance(beta, float) else 1.0
     gemm_symmetric_out(
-        A, B, out, C, dynamic_scheduler=dynamic_scheduler, alpha=alpha_val, beta=beta_val
     )
     return out
@@ -1258,12 +1650,24 @@ def gemm_gated_tuned(
     dynamic_scheduler: bool = False,
     config: Optional[GemmConfig] = None,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> None:
     if config is None:
-        config = default_config(A.device)
     varlen_m = cu_seqlens_m is not None
     if varlen_m:
         assert not config.swap_ab, "Variable-length sequences not supported with swap_ab"
     if A.ndim == 2 and not varlen_m:
         A = A.unsqueeze(0)  # (1, M, K)
     B = B.mT  # (N, K) or (L, N, K)
@@ -1307,7 +1711,8 @@ def gemm_gated_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -1317,6 +1722,8 @@ def gemm_gated_tuned(
         A_idx=A_idx,
         use_tma_gather=config.use_tma_gather,
         concat_layout=concat_layout,
     )
@@ -1403,7 +1810,8 @@ def gemm_dgated_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -1423,10 +1831,10 @@ def gemm_dgated_tuned(
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_gated_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
-    schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str activation='swiglu', Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True, str? concat_layout=None) -> ()",
 )
 def gemm_gated_out(
     A: Tensor,  # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m
@@ -1441,6 +1849,8 @@ def gemm_gated_out(
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     concat_layout: Optional[str] = None,
 ) -> None:
     """GEMM with gated activation and pre-allocated output tensors."""
     fn = gemm_gated_tuned if tuned else partial(gemm_gated_tuned.fn, config=None)
@@ -1455,12 +1865,14 @@ def gemm_gated_out(
         cu_seqlens_m,
         A_idx,
         dynamic_scheduler,
-        concat_layout=tuple(concat_layout.split(",")) if concat_layout else None,
     )
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_dgated_out"),
     mutates_args=("dx_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a!) dx_out, Tensor(b!) postact_out, Tensor? colvec_scale=None, str activation='swiglu', bool colvec_reduce=False, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> Tensor",
@@ -1499,7 +1911,7 @@ def gemm_dgated_out(
     return result
-@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_dgated_out"))
 def gemm_dgated_out_fake(
     A: Tensor,
     B: Tensor,
@@ -1514,20 +1926,6 @@ def gemm_dgated_out_fake(
     dynamic_scheduler: bool = True,
     tuned: bool = True,
 ) -> Tensor:
-    _precompile_default_config(
-        gemm_dgated_tuned,
-        A,
-        B,
-        PreAct,
-        dx_out,
-        postact_out,
-        colvec_scale=colvec_scale,
-        activation=activation,
-        colvec_reduce=colvec_reduce,
-        cu_seqlens_m=cu_seqlens_m,
-        A_idx=A_idx,
-        dynamic_scheduler=dynamic_scheduler,
-    )
     if not colvec_reduce:
         return torch.empty(0, dtype=torch.float32, device=A.device)
     else:
@@ -1541,25 +1939,6 @@ def gemm_dgated_out_fake(
         return torch.empty(out_shape, dtype=torch.float32, device=A.device)
-def _precompile_default_config(autotuned_fn, *args, **kwargs):
-    """Compile the default config in COMPILE_ONLY mode.
-    Checks COMPILE_ONLY flag and SymInt guard, then calls the unwrapped function with
-    config=None (which selects the default config), triggering compilation (exports .o)
-    without benchmarking or kernel launch.
-    Tests use tuned=False which also selects the default config, so this is sufficient.
-    """
-    from .cache_utils import COMPILE_ONLY
-    A = args[0] if args else kwargs.get("A")
-    if not COMPILE_ONLY or A is None or isinstance(A.shape[0], torch.SymInt):
-        return
-    try:
-        autotuned_fn.fn(*args, config=None, **kwargs)
-    except Exception:
-        pass
 @gemm_add_inplace_op.register_fake
 def gemm_add_inplace_fake(
     A: Tensor,
@@ -1575,73 +1954,71 @@ def gemm_add_inplace_fake(
     batch_idx_permute: Optional[Tensor] = None,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
 ) -> None:
-    alpha_val = alpha_tensor if alpha_tensor is not None else alpha
-    beta_val = beta_tensor if beta_tensor is not None else beta
-    add_to_output = isinstance(beta_val, float) and beta_val == 1.0 and cu_seqlens_m is None
-    _precompile_default_config(
-        gemm_tuned,
-        A,
-        B,
-        out,
-        out if not add_to_output else None,
-        alpha=alpha_val,
-        beta=beta_val,
-        cu_seqlens_m=cu_seqlens_m,
-        cu_seqlens_k=cu_seqlens_k,
-        A_idx=A_idx,
-        batch_idx_permute=batch_idx_permute,
-        add_to_output=add_to_output,
-        dynamic_scheduler=dynamic_scheduler,
-    )
-def _register_precompile_fake(custom_op, autotuned_fn, rewrite=None):
-    """Register a fake that precompiles the default config in COMPILE_ONLY mode.
-    For custom_ops that forward args to their autotuned fn. Binds all args by name,
-    strips 'tuned', applies optional rewrite(kw), then calls _precompile_default_config.
-    PyTorch normalizes all custom_op args to positional, so we use inspect.signature
-    to recover keyword names.
     """
-    import inspect
-    sig = inspect.signature(custom_op._init_fn)
-    @custom_op.register_fake
-    def _fake(*args, **kwargs):
-        bound = sig.bind(*args, **kwargs)
-        bound.apply_defaults()
-        kw = dict(bound.arguments)
-        kw.pop("tuned", None)
-        if rewrite is not None:
-            rewrite(kw)
-        _precompile_default_config(autotuned_fn, **kw)
-def _rewrite_merge_alpha(kwargs):
-    """Merge alpha_tensor into alpha for gemm_tuned; add C=None."""
-    at = kwargs.pop("alpha_tensor", None)
-    if at is not None:
-        kwargs["alpha"] = at
-    kwargs.setdefault("C", None)
-def _rewrite_merge_alpha_beta(kwargs):
-    """Merge alpha_tensor/beta_tensor into alpha/beta for gemm_tuned."""
-    at = kwargs.pop("alpha_tensor", None)
-    if at is not None:
-        kwargs["alpha"] = at
-    bt = kwargs.pop("beta_tensor", None)
-    if bt is not None:
-        kwargs["beta"] = bt
-_register_precompile_fake(gemm_out, gemm_tuned, rewrite=_rewrite_merge_alpha)
-_register_precompile_fake(gemm_add_out, gemm_tuned, rewrite=_rewrite_merge_alpha_beta)
-_register_precompile_fake(gemm_act_out, gemm_act_tuned)
-_register_precompile_fake(gemm_dact_out, gemm_dact_tuned)
-_register_precompile_fake(gemm_gated_out, gemm_gated_tuned)
 @gemm_symmetric_out.register_fake
@@ -1653,35 +2030,12 @@ def gemm_symmetric_out_fake(
     dynamic_scheduler: bool = False,
     alpha: float = 1.0,
     beta: float = 1.0,
 ) -> None:
-    from .cache_utils import COMPILE_ONLY
-    if not COMPILE_ONLY or isinstance(A.shape[0], torch.SymInt):
-        return
-    # gemm_symmetric is not autotuned, compile the single fixed config directly
-    sm = get_device_capacity(A.device)[0]
-    tile_m = 256 if sm == 10 else 128
-    tile_n = 128 if sm == 12 else 256
-    cluster_m = 1 if sm == 12 else 2
-    try:
-        gemm_symmetric_dispatch(
-            A.unsqueeze(0) if A.ndim == 2 else A,
-            (B.mT.unsqueeze(0) if B.ndim == 2 else B.mT),
-            out.unsqueeze(0) if out.ndim == 2 else out,
-            (C.unsqueeze(0) if C.ndim == 2 else C) if C is not None else None,
-            torch.zeros(1, dtype=torch.int32, device=A.device) if dynamic_scheduler else None,
-            tile_M=tile_m,
-            tile_N=tile_n,
-            cluster_M=cluster_m,
-            cluster_N=1,
-            pingpong=False,
-            persistent=True,
-            max_swizzle_size=8,
-            alpha=alpha,
-            beta=beta,
-        )
-    except Exception:
-        pass
 ## ── gemm_rms ────────────────────────────────────────────────────────────────
@@ -1704,6 +2058,7 @@ def _gemm_rms_tuned(
     out: Tensor,  # (M, N) or (L, M, N)
     C: Optional[Tensor] = None,  # (M, N) or (L, M, N)
     norm_weight: Optional[Tensor] = None,  # (N,) or (L, N)
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     config: Optional[GemmConfig] = None,
@@ -1723,6 +2078,8 @@ def _gemm_rms_tuned(
         C = C.unsqueeze(0)
     if norm_weight is not None and norm_weight.ndim == 1:
         norm_weight = norm_weight.unsqueeze(0)  # (L, N)
     # Allocate partial reduction buffer
     tile_n = config.tile_n
     n_tiles = (N + tile_n - 1) // tile_n
@@ -1746,11 +2103,13 @@ def _gemm_rms_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
         rowvec=norm_weight,
     )
     # Final reduction: rstd = rsqrt(sum(partials) / N + eps)
     scale = 1.0 / N
@@ -1763,10 +2122,10 @@ def _gemm_rms_tuned(
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_rms_out"),
-    mutates_args=("out",),
     device_types="cuda",
-    schema="(Tensor A, Tensor B, Tensor(a!) out, Tensor? C=None, Tensor? norm_weight=None, float eps=1e-6, bool dynamic_scheduler=False, bool tuned=True) -> Tensor",
 )
 def _gemm_rms_out(
     A: Tensor,
@@ -1774,6 +2133,7 @@ def _gemm_rms_out(
     out: Tensor,
     C: Optional[Tensor] = None,
     norm_weight: Optional[Tensor] = None,
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
@@ -1781,6 +2141,7 @@ def _gemm_rms_out(
     """GEMM + RMS + optional rowvec scaling.
     D_raw = A @ B (+ C), rstd = rsqrt(mean(D_raw^2) + eps), D_out = D_raw * norm_weight.
     """
     fn = _gemm_rms_tuned if tuned else partial(_gemm_rms_tuned.fn, config=None)
     return fn(
@@ -1789,32 +2150,24 @@ def _gemm_rms_out(
         out,
         C=C,
         norm_weight=norm_weight,
         eps=eps,
         dynamic_scheduler=dynamic_scheduler,
     )
-@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_rms_out"))
 def _gemm_rms_out_fake(
     A: Tensor,
     B: Tensor,
     out: Tensor,
     C: Optional[Tensor] = None,
     norm_weight: Optional[Tensor] = None,
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
 ) -> Tensor:
-    _precompile_default_config(
-        _gemm_rms_tuned,
-        A,
-        B,
-        out,
-        C=C,
-        norm_weight=norm_weight,
-        eps=eps,
-        dynamic_scheduler=dynamic_scheduler,
-    )
     rstd_shape = A.shape[:-1]
     return torch.empty(rstd_shape, dtype=torch.float32, device=A.device)
@@ -1844,6 +2197,7 @@ def gemm_rms(
     norm_weight: Optional[Tensor] = None,  # (N,) or (L, N)
     out: Optional[Tensor] = None,  # (M, N) or (L, M, N)
     out_dtype: Optional[torch.dtype] = None,
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
@@ -1851,6 +2205,7 @@ def gemm_rms(
     """GEMM + RMS statistics + optional rowvec scaling.
     D_raw = A @ B (+ C), rstd = rsqrt(mean(D_raw^2) + eps), D_out = D_raw * norm_weight.
     Returns (D_out, rstd).
     """
     out_dtype = A.dtype if out_dtype is None else out_dtype
@@ -1858,12 +2213,28 @@ def gemm_rms(
     if out is None:
         out_shape = (*A.shape[:-1], N)
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     rstd = _gemm_rms_out(
         A,
         B,
         out,
         C=C,
         norm_weight=norm_weight,
         eps=eps,
         dynamic_scheduler=dynamic_scheduler,
         tuned=tuned,
@@ -1927,7 +2298,8 @@ def gemm_norm_act_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -1989,7 +2361,8 @@ def gemm_norm_gated_tuned(
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
-        config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
@@ -1999,7 +2372,7 @@ def gemm_norm_gated_tuned(
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_norm_act_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? rstd=None, str? activation=None, bool dynamic_scheduler=False, bool tuned=True) -> ()",
@@ -2019,23 +2392,11 @@ def gemm_norm_act_out(
     fn(A, B, preact_out, postact_out, C, rstd, activation, dynamic_scheduler)
-@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_norm_act_out"))
-def _gemm_norm_act_out_fake(
-    A,
-    B,
-    preact_out,
-    postact_out,
-    C=None,
-    rstd=None,
-    activation=None,
-    dynamic_scheduler=False,
-    tuned=True,
-) -> None:
-    pass
 @torch.library.custom_op(
-    add_quack_op_namespace_prefix("gemm_norm_gated_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? rstd=None, str activation='swiglu', bool dynamic_scheduler=False, bool tuned=True) -> ()",
@@ -2055,19 +2416,7 @@ def gemm_norm_gated_out(
     fn(A, B, preact_out, postact_out, C, rstd, activation, dynamic_scheduler)
-@torch.library.register_fake(add_quack_op_namespace_prefix("gemm_norm_gated_out"))
-def _gemm_norm_gated_out_fake(
-    A,
-    B,
-    preact_out,
-    postact_out,
-    C=None,
-    rstd=None,
-    activation="swiglu",
-    dynamic_scheduler=False,
-    tuned=True,
-) -> None:
-    pass
 def gemm_norm_act(
@@ -2101,6 +2450,13 @@ def gemm_norm_act(
         preact_out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     if postact_out is None:
         postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device)
     if is_gated:
         gemm_norm_gated_out(
             A,
@@ -2152,13 +2508,12 @@ def gemm_norm_act_ref(
     if rstd is not None:
         D = D * rstd.unsqueeze(-1)
     preact = D.to(out_dtype) if store_preact else None
-    _act_map = {**act_to_pytorch_fn_map, "silu": F.silu}
     if is_gated:
         gate = D[..., ::2]
         up = D[..., 1::2]
         postact = gated_to_pytorch_fn_map[activation](gate, up).to(postact_dtype)
     else:
-        postact = _act_map[activation](D).to(postact_dtype)
     return preact, postact

 from functools import partial
 import torch
+from ._ops_compat import add_op_namespace_prefix
 import torch.nn.functional as F
 from torch import Tensor
 from .rounding import RoundingMode
+def _empty_k_matmul_into(
+    out: Tensor,
+    *,
+    bias: Optional[Tensor] = None,
+    C: Optional[Tensor] = None,
+    beta: float | Tensor = 1.0,
+) -> None:
+    """K=0 fast path: write `beta * C + bias` (or zero if neither) into `out`.
+    Used by every gemm-flavored wrapper to skip a kernel launch when the
+    contraction dim is empty. The matmul A @ B contributes zero, so the only
+    remaining terms are the C term and the (broadcast) bias.
+    """
+    if C is not None:
+        if isinstance(beta, float) and beta == 1.0:
+            out.copy_(C)
+        else:
+            torch.mul(C, beta, out=out)
+    else:
+        out.zero_()
+    if bias is not None:
+        out += bias
+def _silu_tanh(x: Tensor) -> Tensor:
+    x_half = 0.5 * x
+    return x_half * torch.tanh(x_half) + x_half
+def _swiglu_oai_tanh(gate: Tensor, up: Tensor, alpha: float = 1.702) -> Tensor:
+    gate_half = 0.5 * gate
+    return (gate_half * torch.tanh(alpha * gate_half) + gate_half) * (up + 1)
 # Dictionary mapping activation names to PyTorch functions
 act_to_pytorch_fn_map = {
     None: lambda x: x,
+    "silu": F.silu,
+    "silu-tanh": _silu_tanh,
     "relu": F.relu,
     "relu_sq": lambda x: F.relu(x).square(),
     "gelu_tanh_approx": partial(F.gelu, approximate="tanh"),
+    "tanh": torch.tanh,
 }
 # Each function takes (gate, up) and returns postact
 gated_to_pytorch_fn_map = {
     "swiglu": lambda gate, up: F.silu(gate) * up,
+    "swiglu-tanh": lambda gate, up: _silu_tanh(gate) * up,
     "swiglu_oai": lambda gate, up: gate * torch.sigmoid(1.702 * gate) * (up + 1),
+    "swiglu_oai-tanh": _swiglu_oai_tanh,
     "reglu": lambda gate, up: F.relu(gate) * up,
     "geglu": lambda gate, up: F.gelu(gate, approximate="tanh") * up,
     "glu": lambda gate, up: torch.sigmoid(gate) * up,
 }
+ActActivation = Literal[None, "silu", "silu-tanh", "relu", "relu_sq", "gelu_tanh_approx", "tanh"]
+GatedActivation = Literal[
+    "swiglu",
+    "swiglu-tanh",
+    "swiglu_oai",
+    "swiglu_oai-tanh",
+    "reglu",
+    "geglu",
+    "glu",
+]
 Activation = Literal[
     None,
+    "silu",
+    "silu-tanh",
     "relu",
     "relu_sq",
     "gelu_tanh_approx",
+    "tanh",
     "swiglu",
+    "swiglu-tanh",
     "swiglu_oai",
+    "swiglu_oai-tanh",
     "reglu",
     "geglu",
     "glu",
     return t.unflatten(-1, (2, half)).transpose(-2, -1).flatten(-2, -1)
+# ── Blockscaled (MXFP8 / MXFP4 / NVFP4) helpers ─────────────────────────────
+#
+# A and B may be passed as ``(data, scale_factor)`` tuples. Scale factors use
+# the canonical cuBLAS/CUTLASS 128x4 blocked layout: shape
+# ``(M // 128, K // VEC // 4, 32, 4, 4)`` (optionally with a leading batch L),
+# where the ``(32, 4, 4)`` inner block is ``(m % 32, (m // 32) % 4, k_block % 4)``
+# with strides ``(16, 4, 1)`` — one contiguous 512-byte atom per 128 rows x 4
+# K-blocks, matching torchao's ``to_blocked`` and ``torch._scaled_mm``.
+# VEC (the quantization block along K) is implied by the SF dtype:
+# float8_e8m0fnu -> 32 (MX formats), float8_e4m3fn -> 16 (NVFP4).
+# fp4 operands use ``torch.float4_e2m1fn_x2`` storage: shapes carry packed K
+# (two elements per byte); K here always refers to logical K.
+def _unpack_operand(X) -> Tuple[Tensor, Optional[Tensor]]:
+    """Split an ``A`` / ``B`` argument into (data, scale_factor or None)."""
+    if isinstance(X, (tuple, list)):
+        data, sf = X
+        return data, sf
+    return X, None
+def _sf_normalize(SF: Tensor, name: str) -> Tensor:
+    """Validate a user scale-factor tensor and add the batch dim if missing.
+    Requires ``(rm, rk, 32, 4, 4)`` or ``(L, rm, rk, 32, 4, 4)`` with the inner
+    ``(32, 4, 4)`` block contiguous (strides ``(16, 4, 1)`` — one 512 B atom);
+    outer strides are free, so slices of a larger scale buffer are accepted.
+    Returns a zero-copy ``(L, rm, rk, 32, 4, 4)`` view, which is what the
+    compiled kernel consumes directly (the kernel reads only the base pointer
+    and the outer strides; the inner atom layout is hardware-fixed).
+    """
+    assert SF.ndim in (5, 6) and tuple(SF.shape[-3:]) == (32, 4, 4), (
+        f"{name}: expected (rm, rk, 32, 4, 4) or (L, rm, rk, 32, 4, 4) blocked scale factors, "
+        f"got shape {tuple(SF.shape)}"
+    )
+    assert SF.stride()[-3:] == (16, 4, 1), (
+        f"{name}: inner (32, 4, 4) block must be contiguous with strides (16, 4, 1), "
+        f"got {SF.stride()[-3:]}"
+    )
+    return SF.unsqueeze(0) if SF.ndim == 5 else SF
+def _logical_k(X: Tensor) -> int:
+    """Logical contraction extent of the last dim (fp4 packs two elements per byte)."""
+    return X.shape[-1] * (2 if X.dtype == torch.float4_e2m1fn_x2 else 1)
+def _sf_encode(SF: Tensor) -> Tensor:
+    """View e8m0 scale factors as uint8 for the custom-op boundary.
+    Upstream PyTorch bug: a ``float8_e8m0fnu`` input to any mutable custom op
+    makes Inductor's ``decompose_auto_functionalized`` pass fail with
+    "auto_functionalized_v2 was not removed" (e4m3 is unaffected), breaking
+    torch.compile. The uint8 view is zero-copy and unambiguous (uint8 == e8m0,
+    e4m3 stays itself); :func:`_sf_decode` restores the dtype inside the op.
+    """
+    return SF.view(torch.uint8) if SF.dtype == torch.float8_e8m0fnu else SF
+def _sf_decode(SF: Optional[Tensor]) -> Optional[Tensor]:
+    """Inverse of :func:`_sf_encode` (uint8 -> e8m0), applied inside op bodies."""
+    if SF is not None and SF.dtype == torch.uint8:
+        SF = SF.view(torch.float8_e8m0fnu)
+    return SF
 def default_config(device):
     cap = get_device_capacity(device)[0]
+    if cap == 8:
+        return GemmConfig(
+            tile_m=128,
+            tile_n=128,
+            tile_k=32,
+            num_warps=4,
+            cluster_m=1,
+            cluster_n=1,
+            pingpong=False,
+            is_dynamic_persistent=False,
+            device_capacity=8,
+        )
+    elif cap in [10, 11]:
         return GemmConfig(
             tile_m=256,
             tile_n=256,
         )
+def blockscaled_default_config(m: int, n: int) -> GemmConfig:
+    """Default SM100 config for blockscaled GEMM.
+    Large shapes use a (256, 256) tile: it makes num_acc_stage == 1, which turns
+    on ``overlap_accum_sf`` (a second TMEM accumulator stage) so the per-tile
+    scale-apply + TMEM drain overlaps the next tile's MMA instead of
+    serializing after it.
+    """
+    if m >= 512 and n >= 256:
+        tile_m, tile_n, cluster = 256, 256, (2, 1)
+    elif m >= 512 and n >= 128:
+        tile_m, tile_n, cluster = 256, 128, (2, 1)
+    else:
+        tile_m, tile_n, cluster = 128, 128, (1, 1)
+    return GemmConfig(
+        tile_m=tile_m,
+        tile_n=tile_n,
+        cluster_m=cluster[0],
+        cluster_n=cluster[1],
+        pingpong=False,
+        is_dynamic_persistent=True,
+        device_capacity=10,
+    )
 def nvmmh_config(A, B, device_capacity):
     """Use nvMatmulHeuristics to pick a config for pure GEMM (no varlen/gather/epilogue).
     # use_tma_gather only valid when gather_A is active on SM100/SM110
     if not gather_A or device_capacity not in [10, 11]:
         configs = [conf for conf in configs if not conf.kwargs["config"].use_tma_gather]
+    if kwargs.get("SFA", None) is not None:  # blockscaled (SM100 tcgen05 MMA constraints)
+        def _blockscaled_ok(c: GemmConfig) -> bool:
+            return (
+                c.device_capacity in (10, 11)
+                and not c.swap_ab  # untested with blockscaled; SFA/SFB would swap too
+                and c.tile_k is None  # tile_k is derived from the MMA instruction
+                and c.tile_m in (128, 256)
+                and c.tile_n in (64, 128, 192, 256)
+                # SF multicast is limited to 4 CTAs per cluster dim
+                and c.cluster_m <= 4
+                and c.cluster_n <= 4
+            )
+        configs = [conf for conf in configs if _blockscaled_ok(conf.kwargs["config"])]
     return configs
     rounding_mode: int = RoundingMode.RN,
     sr_seed: int | Tensor = 0,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
+    SFA: Optional[Tensor] = None,  # (L, rm, rk, 32, 4, 4) blocked scale factors
+    SFB: Optional[Tensor] = None,  # (L, rn, rk, 32, 4, 4)
 ) -> None:
+    blockscaled = SFA is not None
+    if blockscaled:
+        SFA, SFB = _sf_decode(SFA), _sf_decode(SFB)
     if config is None:
+        if blockscaled:
+            m = A.shape[-2]
+            config = blockscaled_default_config(m, B.shape[-1])
+        else:
+            # Use nvMMH heuristic for pure GEMM (no varlen, no gather, no epilogue)
+            is_pure_gemm = (
+                cu_seqlens_m is None
+                and cu_seqlens_k is None
+                and A_idx is None
+                and C is None
+                and bias is None
+                and not add_to_output
+            )
+            if is_pure_gemm:
+                device_capacity = get_device_capacity(A.device)[0]
+                config = nvmmh_config(A, B, device_capacity)
+            if config is None:
+                config = default_config(A.device)
     varlen_m = cu_seqlens_m is not None
     varlen_k = cu_seqlens_k is not None
     varlen = varlen_m or varlen_k
     gather_A = A_idx is not None
+    if blockscaled:
+        assert not gather_A, "Blockscaled GEMM does not support gather_A yet"
+        assert not concat_layout, "Blockscaled GEMM does not support concat_layout"
+        assert not config.swap_ab, "Blockscaled GEMM does not support swap_ab yet"
     if gather_A:
         assert varlen, "gather_A requires either varlen_m or varlen_k"
         assert config.cluster_n == 1, "gather_A requires cluster_n=1"
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        config.cluster_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
         sr_seed=sr_seed,
         use_tma_gather=config.use_tma_gather,
         concat_layout=swapped_concat,
+        num_warps=config.num_warps,
+        tile_K=config.tile_k,
+        SFA=SFA,
+        SFB=SFB,
     )
     A_idx: Optional[Tensor] = None,  # (total_M,) if gather_A with varlen_m
     dynamic_scheduler: bool = False,
     config: Optional[GemmConfig] = None,
+    SFA: Optional[Tensor] = None,  # (L, rm, rk, 32, 4, 4) blocked scale factors
+    SFB: Optional[Tensor] = None,  # (L, rn, rk, 32, 4, 4)
 ) -> None:
+    blockscaled = SFA is not None
+    if blockscaled:
+        SFA, SFB = _sf_decode(SFA), _sf_decode(SFB)
     if config is None:
+        if blockscaled:
+            config = blockscaled_default_config(A.shape[-2], B.shape[-1])
+        else:
+            config = default_config(A.device)
     varlen_m = cu_seqlens_m is not None
     if varlen_m:
         assert not config.swap_ab, "Variable-length sequences not supported with swap_ab"
+    if blockscaled:
+        assert not varlen_m and A_idx is None, "Blockscaled GEMM does not support varlen/gather yet"
+        assert not config.swap_ab, "Blockscaled GEMM does not support swap_ab yet"
     if A.ndim == 2 and not varlen_m:
         A = A.unsqueeze(0)  # (1, M, K)
     B = B.mT  # (N, K) or (L, N, K)
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
         cu_seqlens_m=cu_seqlens_m,
         A_idx=A_idx,
         use_tma_gather=config.use_tma_gather,
+        SFA=SFA,
+        SFB=SFB,
     )
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
 def gemm(
     # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k
+    # For blockscaled (MXFP8/MXFP4/NVFP4): a tuple (A, SFA) with A fp8/fp4 and SFA the
+    # blocked scale factors (rm, rk, 32, 4, 4) or (L, rm, rk, 32, 4, 4) — see helpers above.
+    A: Tensor | Tuple[Tensor, Tensor],
+    B: Tensor | Tuple[Tensor, Tensor],  # (K, N) or (L, K, N) or (total_K, N) if varlen_k
     out: Optional[Tensor] = None,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m
     bias: Optional[Tensor] = None,  # (N,) or (L, N)
     alpha: float | Tensor = 1.0,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> Tensor:
     """GEMM with optional output tensor and tuning control."""
+    A, SFA = _unpack_operand(A)
+    B, SFB = _unpack_operand(B)
+    assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors"
+    if SFA is not None:
+        SFA = _sf_encode(_sf_normalize(SFA, "SFA"))
+        SFB = _sf_encode(_sf_normalize(SFB, "SFB"))
     if out is None:
+        if out_dtype is None:
+            # Blockscaled inputs are fp8/fp4; default to bf16 output.
+            out_dtype = torch.bfloat16 if SFA is not None else A.dtype
         varlen_m = cu_seqlens_m is not None
         varlen_k = cu_seqlens_k is not None
         if varlen_m:
                 (A.shape[0], B.shape[-1]) if A.ndim == 2 else (A.shape[0], A.shape[-2], B.shape[-1])
             )
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
+    # Empty-input fast path: skip kernel launch.
+    # M=0 / N=0 — the tile scheduler's ceil_div over a zero dim divides by zero.
+    # K=0 — the kernel rejects stride-0 inputs (stride must be divisible by 8);
+    #       semantically the empty contraction yields a zero matrix.
+    if out.numel() == 0:
+        return out
+    if A.numel() == 0:
+        _empty_k_matmul_into(out, bias=bias)
+        return out
     alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha = alpha if isinstance(alpha, float) else 1.0
     sr_seed_tensor = sr_seed if isinstance(sr_seed, Tensor) else None
         sr_seed=sr_seed_int,
         sr_seed_tensor=sr_seed_tensor,
         concat_layout=concat_str,
+        SFA=SFA,
+        SFB=SFB,
     )
     return out
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_out"),
     mutates_args=("out",),
     device_types="cuda",
     # We have to split out alpha and alpha_tensor since torch.library requires
     sr_seed: int = 0,
     sr_seed_tensor: Optional[Tensor] = None,
     concat_layout: Optional[str] = None,
+    # Blockscaled scale factors, (L, rm/rn, rk, 32, 4, 4); tuples are unpacked
+    # to these flat args before the custom-op boundary since torch.library
+    # schemas have no (Tensor, Tensor) argument type.
+    SFA: Optional[Tensor] = None,
+    SFB: Optional[Tensor] = None,
 ) -> None:
     """GEMM with pre-allocated output tensor."""
     fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None)
+    # Shared helpers: drift between this eager body and the register_fake side
+    # is structurally impossible because both call the same functions.
+    alpha = _merge_tensor(alpha, alpha_tensor)
+    sr_seed_arg = _merge_tensor(sr_seed, sr_seed_tensor)
     fn(
         A,
         B,
         dynamic_scheduler=dynamic_scheduler,
         rounding_mode=rounding_mode,
         sr_seed=sr_seed_arg,
+        concat_layout=_parse_concat_layout(concat_layout),
+        SFA=SFA,
+        SFB=SFB,
     )
     return out
+def gemm_blockscaled_ref(
+    A: Tuple[Tensor, Tensor],  # ((M, K) or (L, M, K) fp8/fp4x2, blocked SFA)
+    B: Tuple[Tensor, Tensor],  # ((K, N) or (L, K, N) fp8/fp4x2 K-contig, blocked SFB)
+    alpha: float | Tensor = 1.0,
+    out_dtype: torch.dtype = torch.bfloat16,
+) -> Tensor:
+    """Dequantize-and-matmul reference for blockscaled GEMM."""
+    from .blockscaled.utils import dequant_operand, unpack_scale_blocked_to_2d
+    A, SFA = _unpack_operand(A)
+    B, SFB = _unpack_operand(B)
+    SFA = _sf_normalize(SFA, "SFA")
+    SFB = _sf_normalize(SFB, "SFB")
+    sf_vec = 32 if SFA.dtype == torch.float8_e8m0fnu else 16
+    batched = A.ndim == 3
+    a3 = A if batched else A.unsqueeze(0)  # (l, m, k_packed)
+    b3 = (B if batched else B.unsqueeze(0)).mT  # (l, n, k_packed)
+    a_val = dequant_operand(a3)  # (l, m, k) fp32
+    b_val = dequant_operand(b3)
+    l, m, k = a_val.shape
+    n = b_val.shape[1]
+    sfa = unpack_scale_blocked_to_2d(SFA, m, k // sf_vec).float()
+    sfb = unpack_scale_blocked_to_2d(SFB, n, k // sf_vec).float()
+    a_dq = a_val * sfa.repeat_interleave(sf_vec, dim=-1)
+    b_dq = b_val * sfb.repeat_interleave(sf_vec, dim=-1)
+    out = torch.einsum("lmk,lnk->lmn", a_dq, b_dq)
+    if not (isinstance(alpha, float) and alpha == 1.0):
+        out = out * alpha
+    out = out.to(out_dtype)
+    return out if batched else out.squeeze(0)
 def gemm_add(
     # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k
+    # For blockscaled: a tuple (A, SFA) — see gemm().
+    A: Tensor | Tuple[Tensor, Tensor],
+    B: Tensor | Tuple[Tensor, Tensor],  # (K, N) or (L, K, N) or (total_K, N) if varlen_k
     C: Tensor,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m or (L, M, N) if varlen_k
     out: Optional[Tensor] = None,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m
     alpha: float | Tensor = 1.0,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> Tensor:
     """GEMM with addition and optional output tensor."""
+    A, SFA = _unpack_operand(A)
+    B, SFB = _unpack_operand(B)
+    assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors"
+    if SFA is not None:
+        SFA = _sf_encode(_sf_normalize(SFA, "SFA"))
+        SFB = _sf_encode(_sf_normalize(SFB, "SFB"))
     if out is None:
+        if out_dtype is None:
+            out_dtype = torch.bfloat16 if SFA is not None else A.dtype
         varlen_m = cu_seqlens_m is not None
         varlen_k = cu_seqlens_k is not None
         if varlen_m:
             )
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     add_to_output = C is out and isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None
+    # Empty-input fast path: skip kernel launch (see gemm() for rationale).
+    # K=0 reduces D = alpha*A@B + beta*C to D = beta*C.
+    if out.numel() == 0:
+        return out
+    if A.numel() == 0:
+        if add_to_output:
+            return out  # out IS C, and out += alpha * 0 is a no-op
+        _empty_k_matmul_into(out, C=C, beta=beta)
+        return out
     alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha = alpha if isinstance(alpha, float) else 1.0
     beta_tensor = beta if not isinstance(beta, float) else None
     beta = beta if isinstance(beta, float) else 1.0
+    alpha_arg = _merge_tensor(alpha, alpha_tensor)
+    beta_arg = _merge_tensor(beta, beta_tensor)
     concat_str = ",".join(concat_layout) if concat_layout else None
     if add_to_output:
         gemm_add_inplace(
+            A if SFA is None else (A, SFA),
+            B if SFB is None else (B, SFB),
             out,
             alpha=alpha_arg,
             beta=beta_arg,
             dynamic_scheduler=dynamic_scheduler,
             tuned=tuned,
             concat_layout=concat_str,
+            SFA=SFA,
+            SFB=SFB,
         )
     return out
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_add_out"),
     mutates_args=("out",),
     device_types="cuda",
     # We have to split out alpha and alpha_tensor since torch.library requires
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     concat_layout: Optional[str] = None,
+    SFA: Optional[Tensor] = None,  # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out)
+    SFB: Optional[Tensor] = None,
 ) -> None:
     """GEMM with addition and pre-allocated output tensor."""
     fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None)
+    alpha = _merge_tensor(alpha, alpha_tensor)
+    beta = _merge_tensor(beta, beta_tensor)
     fn(
         A,
         B,
         C,
         alpha=alpha,
         beta=beta,
+        SFA=SFA,
+        SFB=SFB,
         cu_seqlens_m=cu_seqlens_m,
         cu_seqlens_k=cu_seqlens_k,
         A_idx=A_idx,
         batch_idx_permute=batch_idx_permute,
         add_to_output=add_to_output,
         dynamic_scheduler=dynamic_scheduler,
+        concat_layout=_parse_concat_layout(concat_layout),
     )
 def gemm_add_inplace(
     # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (M, total_K) if varlen_k or (whatever, K) if gather_A with varlen_m or (M, whatever) if gather_A with varlen_k
+    # For blockscaled: a tuple (A, SFA) — see gemm().
+    A: Tensor | Tuple[Tensor, Tensor],
+    B: Tensor | Tuple[Tensor, Tensor],  # (K, N) or (L, K, N) or (total_K, N) if varlen_k
     out: Tensor,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m or (L, M, N) if varlen_k
     alpha: float | Tensor = 1.0,
     beta: float | Tensor = 1.0,
         dynamic_scheduler: Whether to use dynamic scheduler
         tuned: Whether to use autotuned configuration
     """
+    A, SFA = _unpack_operand(A)
+    B, SFB = _unpack_operand(B)
+    assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors"
+    if SFA is not None:
+        SFA = _sf_encode(_sf_normalize(SFA, "SFA"))
+        SFB = _sf_encode(_sf_normalize(SFB, "SFB"))
     alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha = alpha if isinstance(alpha, float) else 1.0
     beta_tensor = beta if not isinstance(beta, float) else None
     beta = beta if isinstance(beta, float) else 1.0
+    # Empty-input fast path: out += alpha * A@B with K=0 reduces to out *= beta.
+    # The matmul contributes zero, so use the helper with C=out.
+    if out.numel() == 0:
+        return
+    if A.numel() == 0:
+        if beta != 1.0 or beta_tensor is not None:
+            out.mul_(_merge_tensor(beta, beta_tensor))
+        return
     gemm_add_inplace_op(
         A,
         B,
         concat_layout=",".join(concat_layout)
         if isinstance(concat_layout, tuple)
         else concat_layout,
+        SFA=SFA,
+        SFB=SFB,
     )
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_add_inplace"),
     mutates_args=("out",),
     device_types="cuda",
     # We have to split out alpha and alpha_tensor since torch.library requires
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     concat_layout: Optional[str] = None,
+    SFA: Optional[Tensor] = None,  # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out)
+    SFB: Optional[Tensor] = None,
 ) -> None:
     fn = gemm_tuned if tuned else partial(gemm_tuned.fn, config=None)
+    alpha = _merge_tensor(alpha, alpha_tensor)
+    beta = _merge_tensor(beta, beta_tensor)
     add_to_output = isinstance(beta, float) and beta == 1.0 and cu_seqlens_m is None
     # Use out as both input bias and output
     fn(
         batch_idx_permute=batch_idx_permute,
         add_to_output=add_to_output,
         dynamic_scheduler=dynamic_scheduler,
+        concat_layout=_parse_concat_layout(concat_layout),
+        SFA=SFA,
+        SFB=SFB,
     )
 def gemm_act(
+    # For blockscaled: a tuple (A, SFA) — see gemm().
+    A: Tensor
+    | Tuple[
+        Tensor, Tensor
+    ],  # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m
+    B: Tensor | Tuple[Tensor, Tensor],  # (K, N) or (L, K, N)
     C: Optional[Tensor] = None,  # (M, N) or (L, M, N) or (total_M, N) if varlen_m
     bias: Optional[Tensor] = None,  # (N,) or (L, N)
     activation: Activation = None,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
 ) -> Tuple[Optional[Tensor], Tensor]:
     """GEMM with activation (or gated activation) and optional output tensors."""
+    A, SFA = _unpack_operand(A)
+    B, SFB = _unpack_operand(B)
+    assert (SFA is None) == (SFB is None), "A and B must both (or neither) carry scale factors"
+    if SFA is not None:
+        SFA = _sf_encode(_sf_normalize(SFA, "SFA"))
+        SFB = _sf_encode(_sf_normalize(SFB, "SFB"))
     is_gated = activation in gated_to_pytorch_fn_map
+    default_dtype = torch.bfloat16 if SFA is not None else A.dtype
+    out_dtype = default_dtype if out_dtype is None else out_dtype
+    postact_dtype = default_dtype if postact_dtype is None else postact_dtype
     varlen_m = cu_seqlens_m is not None
     # Determine output shape based on gather_A
     if varlen_m:
         preact_out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     if postact_out is None:
         postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device)
+    # Empty-input fast path. For M=0 or N=0 the outputs are empty; for K=0
+    # (A@B == 0) the no-bias / no-C surface yields preact=0 and act(0)=0 for
+    # every supported activation, so both outputs are zero.
+    if postact_out.numel() == 0 or A.numel() == 0:
+        if preact_out is not None:
+            _empty_k_matmul_into(preact_out)
+        _empty_k_matmul_into(postact_out)
+        return preact_out, postact_out
     concat_str = ",".join(concat_layout) if concat_layout else None
     if is_gated:
         gemm_gated_out(
             dynamic_scheduler,
             tuned,
             concat_layout=concat_str,
+            SFA=SFA,
+            SFB=SFB,
         )
     else:
         gemm_act_out(
             A_idx,
             dynamic_scheduler,
             tuned,
+            SFA=SFA,
+            SFB=SFB,
         )
     return preact_out, postact_out
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_act_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
+    schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True, Tensor? SFA=None, Tensor? SFB=None) -> ()",
 )
 def gemm_act_out(
     A: Tensor,  # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m
     A_idx: Optional[Tensor] = None,  # (total_M,) if gather_A with varlen_m
     dynamic_scheduler: bool = False,
     tuned: bool = True,
+    SFA: Optional[Tensor] = None,  # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out)
+    SFB: Optional[Tensor] = None,
 ) -> None:
     """GEMM with activation and pre-allocated output tensors."""
     fn = gemm_act_tuned if tuned else partial(gemm_act_tuned.fn, config=None)
+    fn(
+        A,
+        B,
+        preact_out,
+        postact_out,
+        C,
+        bias,
+        activation,
+        cu_seqlens_m,
+        A_idx,
+        dynamic_scheduler,
+        SFA=SFA,
+        SFB=SFB,
+    )
 def gemm_act_ref(
         dx_out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     if postact_out is None:
         postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device)
+    # Empty-input fast path: M=0 / N=0 → outputs are empty; K=0 (A.numel()==0)
+    # makes the upstream GEMM gradient zero, so dx is zero regardless of activation.
+    if dx_out.numel() == 0 or A.numel() == 0:
+        _empty_k_matmul_into(dx_out)
+        _empty_k_matmul_into(postact_out)
+        results = [dx_out, postact_out]
+        if colvec_reduce:
+            colvec_shape = (*out_shape[:-1],)
+            results.append(torch.zeros(colvec_shape, dtype=torch.float32, device=A.device))
+        return tuple(results)
     if is_dgated:
         colvec_reduce_final = gemm_dgated_out(
             A,
             dynamic_scheduler,
             tuned,
         )
+        results = [dx_out, postact_out]
+        if colvec_reduce:
+            results.append(colvec_reduce_final)
+        return tuple(results)
     else:
         gemm_dact_out(
             A,
             dynamic_scheduler,
             tuned,
         )
+        results = [dx_out, postact_out]
+        return tuple(results)
 gemm_dgated = gemm_dact
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_dact_out"),
     mutates_args=("dx_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a3!) dx_out, Tensor(a4!) postact_out, str? activation=None, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> ()",
 gemm_dgated_ref = gemm_dact_ref
+def _symmetric_gemm_config(sm: int) -> tuple[int, int, int, bool]:
+    configs = {
+        8: (128, 128, 1, False),
+        9: (128, 256, 2, False),
+        10: (256, 256, 2, False),
+        11: (256, 256, 2, False),
+        12: (128, 128, 1, True),
+    }
+    if sm not in configs:
+        raise NotImplementedError(
+            "gemm_symmetric is only supported on SM8x, SM90, SM100, SM110, and SM120"
+        )
+    return configs[sm]
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_symmetric_out"),
     mutates_args=("out",),
     device_types="cuda",
+    # alpha/beta split into float + Tensor pair because torch.library requires
+    # each schema arg to have a fixed type. See gemm_add_out for the pattern.
 )
 def gemm_symmetric_out(
     A: Tensor,  # (M, K) or (L, M, K)
     dynamic_scheduler: bool = False,
     alpha: float = 1.0,
     beta: float = 1.0,
+    alpha_tensor: Optional[Tensor] = None,
+    beta_tensor: Optional[Tensor] = None,
 ) -> None:
     """GEMM with guaranteed symmetric output."""
+    alpha = _merge_tensor(alpha, alpha_tensor)
+    beta = _merge_tensor(beta, beta_tensor)
     if A.ndim == 2:
         A = A.unsqueeze(0)  # (1, M, K)
     B = B.mT  # (M, K) or (L, M, K)
     )
     sm = get_device_capacity(A.device)[0]
     # We want square tile per cluster
+    tile_m, tile_n, cluster_m, pingpong = _symmetric_gemm_config(sm)
     gemm_symmetric_dispatch(
         A,
         B,
     if out is None:
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
+    alpha_tensor = alpha if not isinstance(alpha, float) else None
     alpha_val = alpha if isinstance(alpha, float) else 1.0
+    beta_tensor = beta if not isinstance(beta, float) else None
     beta_val = beta if isinstance(beta, float) else 1.0
+    # Empty-input fast path: out = alpha * A@A.T + beta * C reduces to beta * C
+    # when K=0 (or just zeros / empty for M=0).
+    if out.numel() == 0:
+        return out
+    if A.numel() == 0:
+        _empty_k_matmul_into(out, C=C, beta=beta)
+        return out
     gemm_symmetric_out(
+        A,
+        B,
+        out,
+        C,
+        dynamic_scheduler=dynamic_scheduler,
+        alpha=alpha_val,
+        beta=beta_val,
+        alpha_tensor=alpha_tensor,
+        beta_tensor=beta_tensor,
     )
     return out
     dynamic_scheduler: bool = False,
     config: Optional[GemmConfig] = None,
     concat_layout: tuple | None = None,  # tensors whose non-contiguous dim is concat [gate; up]
+    SFA: Optional[Tensor] = None,  # (L, rm, rk, 32, 4, 4) blocked scale factors
+    SFB: Optional[Tensor] = None,  # (L, rn, rk, 32, 4, 4)
 ) -> None:
+    blockscaled = SFA is not None
+    if blockscaled:
+        SFA, SFB = _sf_decode(SFA), _sf_decode(SFB)
     if config is None:
+        if blockscaled:
+            config = blockscaled_default_config(A.shape[-2], B.shape[-1])
+        else:
+            config = default_config(A.device)
     varlen_m = cu_seqlens_m is not None
     if varlen_m:
         assert not config.swap_ab, "Variable-length sequences not supported with swap_ab"
+    if blockscaled:
+        assert not varlen_m and A_idx is None, "Blockscaled GEMM does not support varlen/gather yet"
+        assert not concat_layout, "Blockscaled GEMM does not support concat_layout"
+        assert not config.swap_ab, "Blockscaled GEMM does not support swap_ab yet"
     if A.ndim == 2 and not varlen_m:
         A = A.unsqueeze(0)  # (1, M, K)
     B = B.mT  # (N, K) or (L, N, K)
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
         A_idx=A_idx,
         use_tma_gather=config.use_tma_gather,
         concat_layout=concat_layout,
+        SFA=SFA,
+        SFB=SFB,
     )
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_gated_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
+    schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? bias=None, str activation='swiglu', Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=False, bool tuned=True, str? concat_layout=None, Tensor? SFA=None, Tensor? SFB=None) -> ()",
 )
 def gemm_gated_out(
     A: Tensor,  # (M, K) or (L, M, K) or (total_M, K) if varlen_m or (whatever, K) if gather_A with varlen_m
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     concat_layout: Optional[str] = None,
+    SFA: Optional[Tensor] = None,  # blocked scale factors, (L, rm, rk, 32, 4, 4) (see gemm_out)
+    SFB: Optional[Tensor] = None,
 ) -> None:
     """GEMM with gated activation and pre-allocated output tensors."""
     fn = gemm_gated_tuned if tuned else partial(gemm_gated_tuned.fn, config=None)
         cu_seqlens_m,
         A_idx,
         dynamic_scheduler,
+        concat_layout=_parse_concat_layout(concat_layout),
+        SFA=SFA,
+        SFB=SFB,
     )
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_dgated_out"),
     mutates_args=("dx_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor PreAct, Tensor(a!) dx_out, Tensor(b!) postact_out, Tensor? colvec_scale=None, str activation='swiglu', bool colvec_reduce=False, Tensor? cu_seqlens_m=None, Tensor? A_idx=None, bool dynamic_scheduler=True, bool tuned=True) -> Tensor",
     return result
+@torch.library.register_fake(add_op_namespace_prefix("gemm_dgated_out"))
 def gemm_dgated_out_fake(
     A: Tensor,
     B: Tensor,
     dynamic_scheduler: bool = True,
     tuned: bool = True,
 ) -> Tensor:
     if not colvec_reduce:
         return torch.empty(0, dtype=torch.float32, device=A.device)
     else:
         return torch.empty(out_shape, dtype=torch.float32, device=A.device)
 @gemm_add_inplace_op.register_fake
 def gemm_add_inplace_fake(
     A: Tensor,
     batch_idx_permute: Optional[Tensor] = None,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
+    concat_layout: Optional[str] = None,
+    SFA: Optional[Tensor] = None,
+    SFB: Optional[Tensor] = None,
 ) -> None:
+    # Pure no-op: the op only mutates ``out``; kernel compilation is owned
+    # by jit_cache + the async compile pool at real execution time.
+    return
+# ---------------------------------------------------------------------------
+# Shared schema-split helpers.
+#
+# torch.library.custom_op requires a concrete type per arg, so union-typed
+# autotuned args (e.g. ``alpha: Union[float, Tensor]``, ``sr_seed: Union[int,
+# Tensor]``) are split into two fixed-typed schema kwargs at the custom_op
+# boundary (``alpha: float`` + ``alpha_tensor: Optional[Tensor]``). The eager
+# bodies merge them back into the unified form via :func:`_merge_tensor`
+# before calling the autotuned fn.
+# ---------------------------------------------------------------------------
+def _merge_tensor(value, tensor_value):
+    """Return ``tensor_value`` if non-None, else ``value``.
+    Single source of truth for the ``Union[scalar, Tensor]`` schema-split
+    merge. Used both inside eager bodies (where ``value = alpha,
+    tensor_value = alpha_tensor``) and inside the fake path (which derives
+    the split pairs from the custom_op signature).
+    """
+    return tensor_value if tensor_value is not None else value
+def _parse_concat_layout(value):
+    """Coerce ``concat_layout`` from schema form (``Optional[str]``) to
+    autotuned form (``Optional[tuple[str, ...]]``).
+    custom_op schemas can't express ``tuple[str, ...]``, so callers pass a
+    comma-separated string. The autotuned fn keys on a tuple (via
+    ``tuple(sorted(concat_layout))``); a stray string would be iterated
+    char-by-char and silently produce a wrong, never-used compile signature.
+    Single source of truth used by both eager bodies and the fake path.
     """
+    if value is None or isinstance(value, tuple):
+        return value
+    return tuple(value.split(",")) if value else None
+def _register_noop_fake(custom_op):
+    """Register a pure no-op fake for a mutating custom op.
+    These ops only mutate their ``out`` argument, so Dynamo / AOT autograd
+    need no shape effect from the fake; kernel compilation is owned by
+    jit_cache + the async compile pool at real execution time.
+    """
+    @custom_op.register_fake
+    def _fake(*args, **kwargs):
+        return
+_register_noop_fake(gemm_out)
+_register_noop_fake(gemm_add_out)
+_register_noop_fake(gemm_act_out)
+_register_noop_fake(gemm_dact_out)
+_register_noop_fake(gemm_gated_out)
 @gemm_symmetric_out.register_fake
     dynamic_scheduler: bool = False,
     alpha: float = 1.0,
     beta: float = 1.0,
+    alpha_tensor: Optional[Tensor] = None,
+    beta_tensor: Optional[Tensor] = None,
 ) -> None:
+    # Pure no-op: the op only mutates ``out``; kernel compilation is owned
+    # by jit_cache + the async compile pool at real execution time.
+    return
 ## ── gemm_rms ────────────────────────────────────────────────────────────────
     out: Tensor,  # (M, N) or (L, M, N)
     C: Optional[Tensor] = None,  # (M, N) or (L, M, N)
     norm_weight: Optional[Tensor] = None,  # (N,) or (L, N)
+    premult_out: Optional[Tensor] = None,  # (M, N) or (L, M, N) — pre-norm_weight snapshot
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     config: Optional[GemmConfig] = None,
         C = C.unsqueeze(0)
     if norm_weight is not None and norm_weight.ndim == 1:
         norm_weight = norm_weight.unsqueeze(0)  # (L, N)
+    if premult_out is not None and premult_out.ndim == 2:
+        premult_out = premult_out.unsqueeze(0)
     # Allocate partial reduction buffer
     tile_n = config.tile_n
     n_tiles = (N + tile_n - 1) // tile_n
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
         rowvec=norm_weight,
+        aux_out=premult_out,
     )
     # Final reduction: rstd = rsqrt(sum(partials) / N + eps)
     scale = 1.0 / N
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_rms_out"),
+    mutates_args=("out", "premult_out"),
     device_types="cuda",
+    schema="(Tensor A, Tensor B, Tensor(a!) out, Tensor? C=None, Tensor? norm_weight=None, Tensor(a2!)? premult_out=None, float eps=1e-6, bool dynamic_scheduler=False, bool tuned=True) -> Tensor",
 )
 def _gemm_rms_out(
     A: Tensor,
     out: Tensor,
     C: Optional[Tensor] = None,
     norm_weight: Optional[Tensor] = None,
+    premult_out: Optional[Tensor] = None,
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     """GEMM + RMS + optional rowvec scaling.
     D_raw = A @ B (+ C), rstd = rsqrt(mean(D_raw^2) + eps), D_out = D_raw * norm_weight.
+    If premult_out is provided, D_raw (the pre-norm_weight value) is also written to it.
     """
     fn = _gemm_rms_tuned if tuned else partial(_gemm_rms_tuned.fn, config=None)
     return fn(
         out,
         C=C,
         norm_weight=norm_weight,
+        premult_out=premult_out,
         eps=eps,
         dynamic_scheduler=dynamic_scheduler,
     )
+@torch.library.register_fake(add_op_namespace_prefix("gemm_rms_out"))
 def _gemm_rms_out_fake(
     A: Tensor,
     B: Tensor,
     out: Tensor,
     C: Optional[Tensor] = None,
     norm_weight: Optional[Tensor] = None,
+    premult_out: Optional[Tensor] = None,
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
 ) -> Tensor:
     rstd_shape = A.shape[:-1]
     return torch.empty(rstd_shape, dtype=torch.float32, device=A.device)
     norm_weight: Optional[Tensor] = None,  # (N,) or (L, N)
     out: Optional[Tensor] = None,  # (M, N) or (L, M, N)
     out_dtype: Optional[torch.dtype] = None,
+    premult_out: Optional[Tensor] = None,  # (M, N) or (L, M, N) — pre-norm_weight snapshot
     eps: float = 1e-6,
     dynamic_scheduler: bool = False,
     tuned: bool = True,
     """GEMM + RMS statistics + optional rowvec scaling.
     D_raw = A @ B (+ C), rstd = rsqrt(mean(D_raw^2) + eps), D_out = D_raw * norm_weight.
+    If premult_out is provided, D_raw (the pre-norm_weight value) is also written to it.
     Returns (D_out, rstd).
     """
     out_dtype = A.dtype if out_dtype is None else out_dtype
     if out is None:
         out_shape = (*A.shape[:-1], N)
         out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
+    # Empty-input fast path. Skipping the kernel also avoids a torch.library
+    # adinplaceorview_impl IndexError that fires on empty inputs because
+    # premult_out's positional slot isn't materialized in the boxed args tuple.
+    # K=0 with no C reduces the matmul to zero, so D = 0 and rstd = rsqrt(eps).
+    if out.numel() == 0 or A.numel() == 0:
+        _empty_k_matmul_into(out)
+        if premult_out is not None:
+            _empty_k_matmul_into(premult_out)
+        rstd_shape = A.shape[:-1]
+        if A.numel() == 0 and out.numel() > 0:
+            # K=0: rstd = rsqrt(0 + eps) for every row.
+            rstd = torch.full(rstd_shape, eps**-0.5, dtype=torch.float32, device=A.device)
+        else:
+            rstd = torch.empty(rstd_shape, dtype=torch.float32, device=A.device)
+        return out, rstd
     rstd = _gemm_rms_out(
         A,
         B,
         out,
         C=C,
         norm_weight=norm_weight,
+        premult_out=premult_out,
         eps=eps,
         dynamic_scheduler=dynamic_scheduler,
         tuned=tuned,
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
         config.tile_n,
         config.cluster_m,
         config.cluster_n,
+        tile_K=config.tile_k,
+        pingpong=config.pingpong,
         persistent=True,
         is_dynamic_persistent=dynamic_scheduler,
         max_swizzle_size=config.max_swizzle_size,
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_norm_act_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? rstd=None, str? activation=None, bool dynamic_scheduler=False, bool tuned=True) -> ()",
     fn(A, B, preact_out, postact_out, C, rstd, activation, dynamic_scheduler)
+_register_noop_fake(gemm_norm_act_out)
 @torch.library.custom_op(
+    add_op_namespace_prefix("gemm_norm_gated_out"),
     mutates_args=("preact_out", "postact_out"),
     device_types="cuda",
     schema="(Tensor A, Tensor B, Tensor(a2!)? preact_out, Tensor(a3!) postact_out, Tensor? C=None, Tensor? rstd=None, str activation='swiglu', bool dynamic_scheduler=False, bool tuned=True) -> ()",
     fn(A, B, preact_out, postact_out, C, rstd, activation, dynamic_scheduler)
+_register_noop_fake(gemm_norm_gated_out)
 def gemm_norm_act(
         preact_out = torch.empty(out_shape, dtype=out_dtype, device=A.device)
     if postact_out is None:
         postact_out = torch.empty(postact_shape, dtype=postact_dtype, device=A.device)
+    # Empty-input fast path: skip kernel; zero both outputs (act(0)=0 for all
+    # supported activations under the no-bias/no-C path of this test surface).
+    if postact_out.numel() == 0 or A.numel() == 0:
+        if preact_out is not None:
+            _empty_k_matmul_into(preact_out)
+        _empty_k_matmul_into(postact_out)
+        return preact_out, postact_out
     if is_gated:
         gemm_norm_gated_out(
             A,
     if rstd is not None:
         D = D * rstd.unsqueeze(-1)
     preact = D.to(out_dtype) if store_preact else None
     if is_gated:
         gate = D[..., ::2]
         up = D[..., 1::2]
         postact = gated_to_pytorch_fn_map[activation](gate, up).to(postact_dtype)
     else:
+        postact = act_to_pytorch_fn_map[activation](D).to(postact_dtype)
     return preact, postact

build/torch-cuda/quack/gemm_norm_act.py CHANGED Viewed

@@ -18,13 +18,14 @@ from .cute_dsl_utils import (
     get_device_capacity,
     get_max_active_clusters,
 )
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
-from .gemm_act import GemmActMixin, GemmGatedMixin
 from .epi_ops import vec_multiply
 from .activation import act_fn_map, gate_fn_map
-from .cache_utils import jit_cache
 from .rounding import RoundingMode
 from .gemm_tvm_ffi_utils import (
     get_major,
@@ -54,9 +55,9 @@ class GemmNormActMixin(GemmActMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
-        tDrRowVec = epi_loop_tensors["mRowVecBroadcast"]
-        tDrColVec = epi_loop_tensors["mColVecBroadcast"]
         # Load accumulator and apply alpha/beta/C
         rD = tRS_rD.load()
         if const_expr(hasattr(params, "alpha") and params.alpha is not None):
@@ -73,24 +74,28 @@ class GemmNormActMixin(GemmActMixin):
         vec_multiply(self, tRS_rD, tDrColVec, tDrRowVec)
         # Apply activation
         if const_expr(params.act_fn is not None):
-            tRS_rPostAct = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype)
-            if const_expr(self.arch < 100):
-                for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
-                    tRS_rPostAct[i] = params.act_fn(tRS_rD[i])
             else:
-                for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True):
-                    tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn(
                         (tRS_rD[2 * i], tRS_rD[2 * i + 1])
                     )
         else:
-            tRS_rPostAct = tRS_rD
-        return tRS_rPostAct
 class GemmNormActSm90(GemmNormActMixin, GemmSm90):
     pass
 class GemmNormActSm100(GemmNormActMixin, GemmSm100):
     pass
@@ -109,9 +114,9 @@ class GemmNormGatedMixin(GemmGatedMixin):
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
-    ) -> Optional[cute.Tensor]:
-        tDrRowVec = epi_loop_tensors["mRowVecBroadcast"]
-        tDrColVec = epi_loop_tensors["mColVecBroadcast"]
         # Load accumulator and apply alpha/beta/C
         rD = tRS_rD.load()
         if const_expr(hasattr(params, "alpha") and params.alpha is not None):
@@ -127,29 +132,30 @@ class GemmNormGatedMixin(GemmGatedMixin):
         # Multiply by colvec (rstd) and rowvec (norm_weight)
         vec_multiply(self, tRS_rD, tDrColVec, tDrRowVec)
         # Gated activation on normalized D
-        tRS_rPostAct_layout = cute.recast_layout(2, 1, tRS_rD.layout)
-        tRS_rPostAct = cute.make_rmem_tensor(tRS_rPostAct_layout.shape, self.acc_dtype)
-        if const_expr(self.arch < 100):
-            for i in cutlass.range(cute.size(tRS_rPostAct), unroll_full=True):
-                tRS_rPostAct[i] = params.act_fn(tRS_rD[2 * i], tRS_rD[2 * i + 1])
-        else:
-            for i in cutlass.range(cute.size(tRS_rPostAct) // 2, unroll_full=True):
-                tRS_rPostAct[2 * i], tRS_rPostAct[2 * i + 1] = params.act_fn(
-                    (tRS_rD[4 * i], tRS_rD[4 * i + 2]),
-                    (tRS_rD[4 * i + 1], tRS_rD[4 * i + 3]),
-                )
-        return tRS_rPostAct
 class GemmNormGatedSm90(GemmNormGatedMixin, GemmSm90):
     pass
 class GemmNormGatedSm100(GemmNormGatedMixin, GemmSm100):
     pass
-class GemmNormGatedSm120(GemmNormGatedMixin, GemmSm120):
     pass
@@ -183,12 +189,14 @@ def _compile_gemm_norm_act(
 ):
     sm_to_cls = {
         "norm_act": {
             9: GemmNormActSm90,
             10: GemmNormActSm100,
             11: GemmNormActSm100,
             12: GemmNormActSm120,
         },
         "norm_gated": {
             9: GemmNormGatedSm90,
             10: GemmNormGatedSm100,
             11: GemmNormGatedSm100,
@@ -213,7 +221,7 @@ def _compile_gemm_norm_act(
     pa_n = cute.sym_int() if gemm_cls_name == "norm_gated" else n
     pa_leading_dim = 1 if gemm_cls_name == "norm_gated" else pa_leading
     pa_shape = (m, pa_n) if varlen_m else (m, pa_n, l)
-    mPostAct = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa)
     mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4)
     if colvec_ndim == 2:
@@ -234,7 +242,7 @@ def _compile_gemm_norm_act(
             return make_ptr(dtype, 0, cute.AddressSpace.gmem, assumed_align=4)
     epi_args = GemmCls.EpilogueArguments(
-        mPostAct,
         act_fn,
         mRowVecBroadcast=mRowVec,
         mColVecBroadcast=mColVec,
@@ -277,6 +285,7 @@ def gemm_norm_act_fn(
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
     pingpong: bool = False,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
@@ -326,7 +335,9 @@ def gemm_norm_act_fn(
     colvec_ndim = colvec.ndim if colvec is not None else 0
     device_capacity = get_device_capacity(A.device)
-    assert device_capacity[0] in [9, 10, 11, 12], "Only SM90, SM100, SM110, and SM120 are supported"
     if rounding_mode == RoundingMode.RS:
         assert device_capacity[0] == 10, "Stochastic rounding requires SM100"
@@ -349,7 +360,7 @@ def gemm_norm_act_fn(
         d_major,
         c_major,
         postact_major,
-        (tile_M, tile_N),
         (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
@@ -366,11 +377,6 @@ def gemm_norm_act_fn(
         sr_seed_mode=sr_seed_mode,
     )
-    from .cache_utils import COMPILE_ONLY
-    if COMPILE_ONLY:
-        return
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     def scalar_arg(scalar, mode, dtype=Int32):
@@ -395,6 +401,6 @@ def gemm_norm_act_fn(
     varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx)
     if device_capacity[0] in [10, 11]:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None, None)
     else:
-        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None)

     get_device_capacity,
     get_max_active_clusters,
 )
+from .gemm_sm80 import GemmSm80
 from .gemm_sm90 import GemmSm90
 from .gemm_sm100 import GemmSm100
 from .gemm_sm120 import GemmSm120
+from .gemm_act import GemmActMixin, GemmGatedMixin, GemmGatedSm120Mixin
 from .epi_ops import vec_multiply
 from .activation import act_fn_map, gate_fn_map
+from .cache import jit_cache
 from .rounding import RoundingMode
 from .gemm_tvm_ffi_utils import (
     get_major,
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
+        tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast")
+        tDrColVec = epi_loop_tensors.get("mColVecBroadcast")
         # Load accumulator and apply alpha/beta/C
         rD = tRS_rD.load()
         if const_expr(hasattr(params, "alpha") and params.alpha is not None):
         vec_multiply(self, tRS_rD, tDrColVec, tDrRowVec)
         # Apply activation
         if const_expr(params.act_fn is not None):
+            tRS_rAuxOut = cute.make_rmem_tensor(tRS_rD.layout.shape, self.acc_dtype)
+            if const_expr(self.arch != 100):
+                for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True):
+                    tRS_rAuxOut[i] = params.act_fn(tRS_rD[i])
             else:
+                for i in cutlass.range(cute.size(tRS_rAuxOut) // 2, unroll_full=True):
+                    tRS_rAuxOut[2 * i], tRS_rAuxOut[2 * i + 1] = params.act_fn(
                         (tRS_rD[2 * i], tRS_rD[2 * i + 1])
                     )
         else:
+            tRS_rAuxOut = tRS_rD
+        return (tRS_rAuxOut,)
 class GemmNormActSm90(GemmNormActMixin, GemmSm90):
     pass
+class GemmNormActSm80(GemmNormActMixin, GemmSm80):
+    pass
 class GemmNormActSm100(GemmNormActMixin, GemmSm100):
     pass
         epi_loop_tensors: Tuple[cute.Tensor, ...],
         tRS_rD: cute.Tensor,
         tRS_rC: Optional[cute.Tensor] = None,
+    ) -> Tuple[cute.Tensor, ...]:
+        tDrRowVec = epi_loop_tensors.get("mRowVecBroadcast")
+        tDrColVec = epi_loop_tensors.get("mColVecBroadcast")
         # Load accumulator and apply alpha/beta/C
         rD = tRS_rD.load()
         if const_expr(hasattr(params, "alpha") and params.alpha is not None):
         # Multiply by colvec (rstd) and rowvec (norm_weight)
         vec_multiply(self, tRS_rD, tDrColVec, tDrRowVec)
         # Gated activation on normalized D
+        tRS_rAuxOut_layout = cute.recast_layout(2, 1, tRS_rD.layout)
+        tRS_rAuxOut = cute.make_rmem_tensor(tRS_rAuxOut_layout.shape, self.acc_dtype)
+        tRS_rD_pair = cute.flat_divide(tRS_rD, cute.make_layout(2))
+        tRS_rGate = tRS_rD_pair[0, ...]
+        tRS_rUp = tRS_rD_pair[1, ...]
+        vectorize = const_expr(self.arch == 100)
+        for i in cutlass.range(cute.size(tRS_rAuxOut), unroll_full=True, vectorize=vectorize):
+            tRS_rAuxOut[i] = params.act_fn(tRS_rGate[i], tRS_rUp[i])
+        return (tRS_rAuxOut,)
 class GemmNormGatedSm90(GemmNormGatedMixin, GemmSm90):
     pass
+class GemmNormGatedSm80(GemmNormGatedMixin, GemmSm80):
+    pass
 class GemmNormGatedSm100(GemmNormGatedMixin, GemmSm100):
     pass
+class GemmNormGatedSm120(GemmGatedSm120Mixin, GemmNormGatedMixin, GemmSm120):
     pass
 ):
     sm_to_cls = {
         "norm_act": {
+            8: GemmNormActSm80,
             9: GemmNormActSm90,
             10: GemmNormActSm100,
             11: GemmNormActSm100,
             12: GemmNormActSm120,
         },
         "norm_gated": {
+            8: GemmNormGatedSm80,
             9: GemmNormGatedSm90,
             10: GemmNormGatedSm100,
             11: GemmNormGatedSm100,
     pa_n = cute.sym_int() if gemm_cls_name == "norm_gated" else n
     pa_leading_dim = 1 if gemm_cls_name == "norm_gated" else pa_leading
     pa_shape = (m, pa_n) if varlen_m else (m, pa_n, l)
+    mAuxOut = fake_tensor(postact_dtype, pa_shape, leading_dim=pa_leading_dim, divisibility=div_pa)
     mRowVec = fake_tensor(rowvec_dtype, (l, n), leading_dim=1, divisibility=4)
     if colvec_ndim == 2:
             return make_ptr(dtype, 0, cute.AddressSpace.gmem, assumed_align=4)
     epi_args = GemmCls.EpilogueArguments(
+        mAuxOut,
         act_fn,
         mRowVecBroadcast=mRowVec,
         mColVecBroadcast=mColVec,
     tile_N: int,
     cluster_M: int,
     cluster_N: int,
+    tile_K: int | None = None,
     pingpong: bool = False,
     persistent: bool = True,
     is_dynamic_persistent: bool = False,
     colvec_ndim = colvec.ndim if colvec is not None else 0
     device_capacity = get_device_capacity(A.device)
+    assert device_capacity[0] in [8, 9, 10, 11, 12], (
+        "Only SM8x, SM90, SM100, SM110, and SM120 are supported"
+    )
     if rounding_mode == RoundingMode.RS:
         assert device_capacity[0] == 10, "Stochastic rounding requires SM100"
         d_major,
         c_major,
         postact_major,
+        (tile_M, tile_N, tile_K) if tile_K is not None else (tile_M, tile_N),
         (cluster_M, cluster_N, 1),
         pingpong,
         persistent,
         sr_seed_mode=sr_seed_mode,
     )
     max_active_clusters = get_max_active_clusters(cluster_M * cluster_N) if persistent else 0
     def scalar_arg(scalar, mode, dtype=Int32):
     varlen_args = make_varlen_args(cu_seqlens_m, None, A_idx)
     if device_capacity[0] in [10, 11]:
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args, None, None)
     else:
+        compiled_fn(A_p, B_p, D_p, C_p, epi_args, scheduler_args, varlen_args)

build/torch-cuda/quack/gemm_sm100.py CHANGED Viewed

@@ -1,4 +1,4 @@
-# Copyright (c) 2025-2026, Tri Dao.
 # Based on the cute-dsl example:
 # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
@@ -23,11 +23,22 @@ from cutlass.cute.nvgpu.warp import (
 )
 from cutlass import Int32, Float32, Boolean, const_expr
 from cutlass.utils import LayoutEnum
-from .pipeline import PipelineTmaUmma, PipelineTmaCpAsyncUmma
 from .tile_scheduler import TileSchedulerOptions
 from .varlen_utils import VarlenArguments, VarlenManager
-from .gemm_sm90 import GemmSm90, NamedBarrierGemm
 from . import layout_utils
 from . import copy_utils as copy_utils
 from . import sm100_utils as quack_sm100_utils
@@ -64,29 +75,7 @@ SM100 tcgen05.mma instructions operate as follows:
 - Write accumulator to TMEM
 The accumulator in TMEM must then be loaded to registers before writing back to GMEM.
-Input arguments to this example is same as dense_gemm.py.
-.. code-block:: bash
-    python examples/blackwell/dense_gemm_persistent.py                          \
-      --ab_dtype Float16 --d_dtype Float16 --acc_dtype Float32                  \
-      --mma_tiler_mn 256,128 --cluster_shape_mn 2,1                             \
-      --mnkl 8192,8192,8192,1                                                   \
-      --use_2cta_instrs
-To collect performance with NCU profiler:
-.. code-block:: bash
-    ncu python examples/blackwell/dense_gemm_persistent.py                     \
-      --ab_dtype Float16 --d_dtype Float16 --acc_dtype Float32                 \
-      --mma_tiler_mn 256,128 --cluster_shape_mn 2,1                            \
-      --mnkl 8192,8192,8192,1                                                  \
-      --use_2cta_instrs                                        \
-      --warmup_iterations 1 --iterations 10 --skip_ref_check
-Constraints are same as dense_gemm.py:
 * Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2),
   see detailed valid dtype combinations in below GemmSm100 class documentation
 * A/B tensor must have the same data type
@@ -99,15 +88,15 @@ Constraints are same as dense_gemm.py:
 """
-class GemmSm100(GemmSm90):
     """This class implements batched matrix multiplication (C = A x B) with support for various data types
     and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
     :param acc_dtype: Data type for accumulation during computation
     :type acc_dtype: type[cutlass.Numeric]
-    :param mma_tiler_mn: Shape of the MMA tile. Pass (M, N) to default K to
         4 MMA instructions, or (M, N, K) to set the K tile size explicitly.
-    :type mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]]
     :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
     :type cluster_shape_mn: Tuple[int, int]
@@ -141,7 +130,7 @@ class GemmSm100(GemmSm90):
     Example:
         >>> gemm = GemmSm100(
         ...     acc_dtype=Float32,
-        ...     mma_tiler_mn=(128, 128),
         ...     cluster_shape_mn=(2, 2)
         ... )
         >>> gemm(mA, mB, mD, max_active_clusters, stream)
@@ -149,14 +138,14 @@ class GemmSm100(GemmSm90):
     arch = 100
-    EpilogueArguments = GemmSm90.EpilogueArguments
-    EpilogueParams = GemmSm90.EpilogueParams
     def __init__(
         self,
         acc_dtype: Type[cutlass.Numeric],
         a_dtype: Type[cutlass.Numeric],  # ignored for now
-        mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mnk: Tuple[int, int, int],
         sf_vec_size: Optional[int] = None,
         gather_A: bool = False,
@@ -171,7 +160,7 @@ class GemmSm100(GemmSm90):
         1.  MMA Instruction Settings (tcgen05):
             - acc_dtype: Data types for MMA accumulator.
-            - mma_tiler_mn: The (M, N) shape of the MMA instruction tiler.
             - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant
               with cta_group=2 should be used.
@@ -180,27 +169,29 @@ class GemmSm100(GemmSm90):
         :param acc_dtype: Data type of the accumulator.
         :type acc_dtype: type[cutlass.Numeric]
-        :param mma_tiler_mn: (M, N) or (M, N, K) shape of the MMA tile.
             If only (M, N) is given, K defaults to 4 * instruction K.
-        :type mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]]
         :param cluster_shape_mnk: Tuple (ClusterM, ClusterN) shape of the cluster.
         :type cluster_shape_mnk: Tuple[int, int]
         """
         self.acc_dtype: Type[cutlass.Numeric] = acc_dtype
-        self.use_2cta_instrs = mma_tiler_mn[0] in (256,)
         self.cluster_shape_mnk = cluster_shape_mnk
         assert cluster_shape_mnk[2] == 1, "Cluster shape K must be 1"
         # K dimension: if user provides 3 values, use their K; otherwise default in _setup_attributes
-        if len(mma_tiler_mn) == 3:
-            self.mma_tiler = tuple(mma_tiler_mn)
         else:
-            self.mma_tiler = (*mma_tiler_mn, 0)
-        self.sf_vec_size = sf_vec_size
-        self.blockscaled = sf_vec_size is not None
         self.is_persistent = True
-        self.pingpong = False  # for compatibility with GemmSm90
         self.use_clc_persistence = use_clc_persistence
         self.gather_A = gather_A
         self.concat_layout = concat_layout or ()
         self.use_tma_gather = use_tma_gather
@@ -229,6 +220,23 @@ class GemmSm100(GemmSm90):
             barrier_id=int(NamedBarrierGemm.Epilogue),
             num_threads=self.num_epi_warps * cute.arch.WARP_SIZE,
         )
         # Register reallocation for gather_A (3 warp groups, 504 regs total, 168 per WG default).
         # Heavy epilogues (e.g. colvec_reduce in DGated) override these to avoid register spilling.
         # Without gather_A there are only 2 WGs (512 total, 256 per WG = max), no reallocation needed.
@@ -250,6 +258,15 @@ class GemmSm100(GemmSm90):
         # Multiple of 4 warps to increase/decrease number of registers
         assert self.threads_per_cta % 128 == 0
     def _setup_attributes(self, epilogue_args: EpilogueArguments, varlen_args: VarlenArguments):
         """Set up configurations that are dependent on GEMM inputs
@@ -264,6 +281,8 @@ class GemmSm100(GemmSm90):
         - Computing A/B/C shared memory layout
         - Computing tensor memory allocation columns
         """
         # Compute mma instruction shapes
         mma_inst_bits_k = 256
         # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K)
@@ -284,6 +303,7 @@ class GemmSm100(GemmSm90):
         if const_expr(not self.blockscaled):
             self.tiled_mma = sm100_utils.make_trivial_tiled_mma(
                 self.a_dtype,
                 self.a_major_mode,
                 self.b_major_mode,
                 self.acc_dtype,
@@ -294,6 +314,7 @@ class GemmSm100(GemmSm90):
         else:
             self.tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
                 self.a_dtype,
                 self.a_major_mode,
                 self.b_major_mode,
                 self.sf_dtype,
@@ -303,6 +324,7 @@ class GemmSm100(GemmSm90):
             )
             self.tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
                 self.a_dtype,
                 self.a_major_mode,
                 self.b_major_mode,
                 self.sf_dtype,
@@ -313,6 +335,10 @@ class GemmSm100(GemmSm90):
         # Compute mma/cluster/tile shapes
         if self.mma_tiler[2] > 0:
             mma_inst_tile_k = self.mma_tiler[2] // self.mma_inst_shape_mnk[2]
         else:
             mma_inst_tile_k = 4
@@ -340,6 +366,25 @@ class GemmSm100(GemmSm90):
                 self.mma_tiler_sfb[1],
                 self.mma_tiler_sfb[2],
             )
         else:
             self.cta_tile_shape_mnk_sfb = None
@@ -368,13 +413,25 @@ class GemmSm100(GemmSm90):
             self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
         # Compute epilogue subtile
         self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
             self.cta_tile_shape_mnk,
             self.use_2cta_instrs,
             self.d_layout if self.d_layout is not None else LayoutEnum.ROW_MAJOR,
             self.d_dtype if self.d_dtype is not None else cutlass.BFloat16,
-            layout_c=self.c_layout,
-            elem_ty_c=self.c_dtype,
         )
         # TMA store tile starts must stay aligned when advancing across CTA-N tiles.
         # There's a bug w compute_epilogue_tile_shape (as of cutlass-dsl 4.4.2) where if
@@ -415,8 +472,16 @@ class GemmSm100(GemmSm90):
             prefetch_A_idx,
             cutlass.utils.get_smem_capacity_in_bytes(f"sm_{self.arch}"),  # smem_capacity
             self.occupancy,
         )
-        self.sched_stage = 1
         self.a_prefetch_stage = (
             0
             if not self.gather_A
@@ -512,7 +577,6 @@ class GemmSm100(GemmSm90):
         stream: cuda.CUstream,
         mSFA: Optional[cute.Tensor] = None,
         mSFB: Optional[cute.Tensor] = None,
-        trace_ptr: Optional[cutlass.Int64] = None,
     ):
         """Execute the GEMM operation in steps:
         - Setup static attributes before smem/grid/tma computation
@@ -576,7 +640,7 @@ class GemmSm100(GemmSm90):
             # so non-packed buffers work (e.g. a slice of a larger scale tensor).
             # Only the innermost 512-B tile must be contiguous.
             # For varlen_m, mSFA is sized for per-expert 128-row-padded storage
-            # (dQaccum format), so use its own M dim (= total_padded_rm * 128)
             # instead of mA.shape[0] (= total_m, unpadded).
             if const_expr(cute.rank(mA) == 3):
                 sfa_shape = mA.shape
@@ -599,8 +663,12 @@ class GemmSm100(GemmSm90):
         a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
         b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
         tma_atom_a, tma_tensor_a = None, None
-        a_op = sm100_utils.cluster_shape_to_tma_atom_A(
-            self.cluster_shape_mnk, self.tiled_mma.thr_id
         )
         if const_expr(not self.gather_A):
             tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
@@ -628,9 +696,9 @@ class GemmSm100(GemmSm90):
                 tma_smem_layout.shape,
                 internal_type=(cutlass.TFloat32 if mA.element_type is Float32 else None),
             )
-        b_op = sm100_utils.cluster_shape_to_tma_atom_B(
-            self.cluster_shape_mnk, self.tiled_mma.thr_id
-        )
         tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
             b_op,
             copy_utils.create_ragged_tensor_for_tma(mB, ragged_dim=1) if varlen_k else mB,
@@ -645,9 +713,7 @@ class GemmSm100(GemmSm90):
         tma_atom_sfb, tma_tensor_sfb = None, None
         if const_expr(self.blockscaled):
             # Setup TMA load for SFA
-            sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
-                self.cluster_shape_mnk, self.tiled_mma.thr_id
-            )
             sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
             tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
                 sfa_op,
@@ -659,9 +725,7 @@ class GemmSm100(GemmSm90):
                 internal_type=cutlass.Int16,
             )
             # Setup TMA load for SFB
-            sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
-                self.cluster_shape_mnk, self.tiled_mma.thr_id
-            )
             sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
             tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
                 sfb_op,
@@ -672,9 +736,15 @@ class GemmSm100(GemmSm90):
                 self.cluster_layout_sfb_vmnk.shape,
                 internal_type=cutlass.Int16,
             )
-            if const_expr(
-                self.cta_tile_shape_mnk[1] == 192 and self.sf_dtype is cutlass.Float8E8M0FNU
-            ):
                 x = tma_tensor_sfb.stride[0][1]
                 y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4)
                 tma_tensor_sfb = cute.make_tensor(
@@ -702,28 +772,24 @@ class GemmSm100(GemmSm90):
             self.num_tma_load_bytes += sfa_copy_size + sfb_copy_size
         self.num_tma_load_bytes *= atom_thr_size
-        # Setup TMA store for D
-        tma_atom_d, tma_tensor_d = None, None
-        if const_expr(mD is not None):
-            tma_atom_d, tma_tensor_d = self._make_tma_epi_atoms_and_tensors(
-                copy_utils.create_ragged_tensor_for_tma(mD, ragged_dim=0, ptr_shift=True)
-                if varlen_m
-                else mD,
-                self.epi_smem_layout_staged,
-                self.epi_tile,
-                op_type="store"
-                if not (hasattr(epilogue_args, "add_to_output") and epilogue_args.add_to_output)
-                else "add",
-            )
-        tma_atom_c, tma_tensor_c = None, None
-        if const_expr(mC is not None):
-            tma_atom_c, tma_tensor_c = self._make_tma_epi_atoms_and_tensors(
-                mC, self.epi_c_smem_layout_staged, self.epi_tile, op_type="load"
-            )
         epilogue_params = self.epi_to_underlying_arguments(epilogue_args)
         varlen_params = VarlenManager.to_underlying_arguments(varlen_args)
         TileSchedulerCls = self.get_scheduler_class(varlen_m=varlen_m)
         tile_sched_args = self.get_scheduler_arguments(
             mA, mB, mD, scheduler_args, varlen_args, epilogue_args
@@ -750,19 +816,21 @@ class GemmSm100(GemmSm90):
                 self.cta_tile_shape_mnk[0] if varlen_m else self.cta_tile_shape_mnk[2]
             )
-        # Define shared storage for kernel
-        @cute.struct
         class SharedStorage:
-            ab_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.ab_stage * 2]
-            epi_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.epi_c_stage * 2]
-            acc_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
-            sched_pipeline_array_ptr: cute.struct.MemRange[cutlass.Int64, self.sched_stage * 2]
-            a_prefetch_pipeline_array_ptr: cute.struct.MemRange[
-                cutlass.Int64, self.a_prefetch_stage * 2
             ]
-            sched_data: cute.struct.MemRange[Int32, self.sched_stage * 12]
-            tmem_dealloc_mbar_ptr: cutlass.Int64
-            tmem_holding_buf: Int32
             sAIdx: cute.struct.Align[cute.struct.MemRange[Int32, a_idx_smem_size], 16]
             # (EPI_TILE_M, EPI_TILE_N, STAGE)
             sD: cute.struct.Align[
@@ -831,7 +899,6 @@ class GemmSm100(GemmSm90):
             self.epi_tile,
             tile_sched_params,
             TileSchedulerCls,
-            trace_ptr,
         ).launch(
             grid=grid,
             block=[self.threads_per_cta, 1, 1],
@@ -874,16 +941,11 @@ class GemmSm100(GemmSm90):
         epi_tile: cute.Tile,
         tile_sched_params,
         TileSchedulerCls: cutlass.Constexpr[Callable],
-        trace_ptr: Optional[cutlass.Int64] = None,
     ):
         """
         GPU device kernel performing the Persistent batched GEMM computation.
         """
-        from .trace import TraceContext
-        tctx = TraceContext.create(trace_ptr)
         varlen_m = const_expr(varlen_params.cu_seqlens_m is not None)
         varlen_k = const_expr(varlen_params.cu_seqlens_k is not None)
         assert not (varlen_m and varlen_k)
@@ -891,6 +953,7 @@ class GemmSm100(GemmSm90):
             assert varlen_m or varlen_k
         has_D = const_expr(mD_mnl is not None)
         has_C = const_expr(mC_mnl is not None)
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
@@ -920,39 +983,26 @@ class GemmSm100(GemmSm90):
         # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
         smem = cutlass.utils.SmemAllocator()
-        storage = smem.allocate(self.shared_storage)
         # Initialize pipelines and states
         ab_pipeline = self.make_ab_pipeline(
             tiled_mma=tiled_mma,
             cluster_layout_vmnk=cluster_layout_vmnk,
-            ab_pipeline_mbar_ptr=storage.ab_pipeline_array_ptr.data_ptr(),
             is_leader_cta=is_leader_cta,
         )
         epi_pipeline = None
-        if const_expr(has_C):
-            epi_pipeline = self.make_epi_pipeline(
-                c_smem_layout=cute.slice_(epi_c_smem_layout, (None, None, 0)),
-                epi_pipeline_mbar_ptr=storage.epi_pipeline_array_ptr.data_ptr(),
-            )
-        acc_pipeline = self.make_acc_pipeline(
-            cluster_layout_vmnk=cluster_layout_vmnk,
-            acc_pipeline_mbar_ptr=storage.acc_pipeline_array_ptr.data_ptr(),
-        )
         sched_pipeline = None
         sched_data = None
         if const_expr(self.is_persistent):
-            sched_pipeline = self.make_sched_pipeline(
-                self.cluster_shape_mnk,
-                sched_pipeline_mbar_ptr=storage.sched_pipeline_array_ptr.data_ptr(),
-                has_C=has_C,
-            )
-            sched_data = storage.sched_data.get_tensor((12, self.sched_stage))
         a_prefetch_pipeline = None
         if const_expr(self.gather_A):
-            a_prefetch_pipeline = self.make_a_prefetch_pipeline(
-                storage.a_prefetch_pipeline_array_ptr.data_ptr(),
-            )
         tmem_alloc_barrier = pipeline.NamedBarrier(
             barrier_id=int(NamedBarrierGemm.TmemPtr),
@@ -960,11 +1010,9 @@ class GemmSm100(GemmSm90):
         )
         # Tensor memory dealloc barrier init
         tmem = cutlass.utils.TmemAllocator(
-            storage.tmem_holding_buf,
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.epilog_warp_id[0],
             is_two_cta=use_2cta_instrs,
-            two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
         )
         # Cluster arrive after barrier init
@@ -1020,13 +1068,18 @@ class GemmSm100(GemmSm90):
         )
         TileSchedulerCls = partial(
-            TileSchedulerCls.create, tile_sched_params, sched_data, sched_pipeline
         )
         epi_load_barrier = None
-        if const_expr(has_C):
             epi_load_barrier = pipeline.NamedBarrier(
-                barrier_id=int(NamedBarrierGemm.EpilogueLoad), num_threads=2 * cute.arch.WARP_SIZE
             )
         # Cluster wait before tensor memory alloc
@@ -1042,30 +1095,6 @@ class GemmSm100(GemmSm90):
                 cute.arch.griddepcontrol_wait()
             if const_expr(self.gather_A):
                 cute.arch.setmaxregister_decrease(self.num_regs_other)
-            # Compute multicast mask for A/B buffer full
-            block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
-            block_in_cluster_coord_sfb_vmnk = None
-            if const_expr(self.blockscaled):
-                block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
-                    cta_rank_in_cluster
-                )
-            a_mcast_mask, b_mcast_mask = None, None
-            sfa_mcast_mask, sfb_mcast_mask = None, None
-            if const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
-                a_mcast_mask = cpasync.create_tma_multicast_mask(
-                    cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
-                )
-                b_mcast_mask = cpasync.create_tma_multicast_mask(
-                    cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
-                )
-                if const_expr(self.blockscaled):
-                    sfa_mcast_mask = cpasync.create_tma_multicast_mask(
-                        cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
-                    )
-                    sfb_mcast_mask = cpasync.create_tma_multicast_mask(
-                        cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
-                    )
             # Persistent tile scheduling loop
             tile_scheduler = TileSchedulerCls()
             work_tile = tile_scheduler.initial_work_tile_info()
@@ -1076,7 +1105,16 @@ class GemmSm100(GemmSm90):
                 pipeline.PipelineUserType.Consumer, self.a_prefetch_stage
             )
             do_epi_load_barrier_arrive = Boolean(True)
             while work_tile.is_valid_tile:
                 tile_coord_mnkl = work_tile.tile_idx
                 batch_idx = tile_coord_mnkl[3]
                 # Local_tile partition global tensors
@@ -1102,7 +1140,7 @@ class GemmSm100(GemmSm90):
                 )
                 if const_expr(self.blockscaled):
                     # (bM, bK)
-                    # SFA uses padded per-expert offset (dQaccum format), not
                     # the A-data offset — allows varlen_m seqlens that aren't
                     # multiples of 128.
                     gSFA_mkl = cute.local_tile(
@@ -1111,38 +1149,36 @@ class GemmSm100(GemmSm90):
                         (mma_tile_coord_mnl[0], None),
                     )
                     # (bN, bK)
-                    # SFB uses padded per-expert K offset in varlen_k (dQaccum format).
                     gSFB_nkl = cute.local_tile(
                         varlen_manager.offset_batch_SFB(mSFB_nkl, batch_idx),
                         cute.select(self.mma_tiler_sfb, [1, 2]),
-                        (
-                            (
-                                mma_tile_coord_mnl[1] // 2
-                                if self.cta_tile_shape_mnk[1] == 64
-                                else mma_tile_coord_mnl[1]
-                            ),
-                            None,
-                        ),
                     )
                 # Partition global tensor for TiledMMA_A/B/D
                 # Then partition global/shared tensor for TMA load A/B
                 len_k = varlen_manager.len_k(batch_idx)
-                # TMA load A partition_S/D
-                a_cta_layout = cute.make_layout(
-                    cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
-                )
                 copy_A, prefetch_A = None, None
                 if const_expr(not self.gather_A):
                     # (MMA, MMA_M, MMA_K, RestK)
                     tCgA = thr_mma.partition_A(gA_mk)
-                    copy_A, _, _ = copy_utils.tma_get_copy_fn(
-                        tma_atom_a,
-                        cta_coord=block_in_cluster_coord_vmnk[2],
-                        cta_layout=a_cta_layout,
-                        src_tensor=tCgA,
-                        dst_tensor=sA,
-                        mcast_mask=a_mcast_mask,
                     )
                 else:
                     # For varlen_m paths (TMA or cp.async): consume indices from
@@ -1162,9 +1198,7 @@ class GemmSm100(GemmSm90):
                         warp_idx,
                     )
                     if const_expr(varlen_m):
-                        cute.arch.sync_warp()
-                        with cute.arch.elect_one():
-                            a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
                         a_prefetch_consumer_state.advance()
                     if const_expr(prefetch_A is not None):
                         prefetch_A = partial(prefetch_A, a_prefetch_pipeline)
@@ -1176,52 +1210,33 @@ class GemmSm100(GemmSm90):
                     # (MMA, MMA_N, MMA_K)
                     tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
                 # TMA load B partition_S/D
-                copy_B, _, _ = copy_utils.tma_get_copy_fn(
-                    tma_atom_b,
-                    cta_coord=block_in_cluster_coord_vmnk[1],
-                    cta_layout=cute.make_layout(
-                        cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
-                    ),
-                    src_tensor=tCgB,
-                    dst_tensor=sB,
-                    mcast_mask=b_mcast_mask,
                 )
                 copy_SFA, copy_SFB = None, None
                 if const_expr(self.blockscaled):
                     #  TMA load SFA partition_S/D
-                    copy_SFA, _, _ = copy_utils.tma_get_copy_fn(
                         tma_atom_sfa,
-                        cta_coord=block_in_cluster_coord_vmnk[2],
-                        cta_layout=a_cta_layout,
                         src_tensor=tCgSFA,
                         dst_tensor=sSFA,
-                        filter_zeros=True,
-                        mcast_mask=sfa_mcast_mask,
                     )
                     # TMA load SFB partition_S/D
-                    sfb_cta_layout = cute.make_layout(
-                        cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
-                    )
-                    copy_SFB, _, _ = copy_utils.tma_get_copy_fn(
                         tma_atom_sfb,
-                        cta_coord=block_in_cluster_coord_sfb_vmnk[1],
-                        cta_layout=sfb_cta_layout,
                         src_tensor=tCgSFB,
                         dst_tensor=sSFB,
-                        filter_zeros=True,
-                        mcast_mask=sfb_mcast_mask,
                     )
                 k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2])
-                tctx.b("tma_load")
                 if const_expr(not self.gather_A):
-                    ab_producer_state = self.load_AB(
                         ab_pipeline,
                         ab_producer_state,
-                        copy_A,
-                        copy_B,
                         k_tile_cnt,
-                        copy_SFA,
-                        copy_SFB,
                     )
                 elif const_expr(self.use_tma_gather):
                     ab_producer_state, a_prefetch_consumer_state = self.load_AB_tma_gather(
@@ -1243,7 +1258,7 @@ class GemmSm100(GemmSm90):
                         copy_B,
                         k_tile_cnt,
                     )
-                tctx.e("tma_load")
                 if const_expr(epi_load_barrier is not None):
                     # In the first work tile, the epi load warp will wait for the signal
                     # from the mainloop load warp to start loading C, to avoid interfering
@@ -1252,8 +1267,10 @@ class GemmSm100(GemmSm90):
                         epi_load_barrier.arrive()
                         do_epi_load_barrier_arrive = Boolean(False)
                 # Advance to next tile
                 tile_scheduler.advance_to_next_work()
                 work_tile = tile_scheduler.get_current_work()
             # Wait A/B buffer empty
             if warp_idx == self.ab_load_warp_id:
                 ab_pipeline.producer_tail(ab_producer_state)
@@ -1274,11 +1291,19 @@ class GemmSm100(GemmSm90):
                 work_tile = tile_scheduler.initial_work_tile_info()
                 while work_tile.is_valid_tile:
                     # Advance to next tile
                     tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp)
                     work_tile = tile_scheduler.get_current_work()
                     # End of persistent scheduler loop
                 if is_scheduler_warp:
                     tile_scheduler.producer_tail()
         # Specialized A-index prefetch warp (gather_A only)
         if const_expr(self.gather_A):
@@ -1359,9 +1384,9 @@ class GemmSm100(GemmSm90):
             if const_expr(self.gather_A):
                 cute.arch.setmaxregister_decrease(self.num_regs_other)
             # PDL: wait for prior kernel before any C TMA loads (matches cutlass C++ epi_load)
-            if const_expr(self.use_pdl and mC_mnl is not None):
                 cute.arch.griddepcontrol_wait()
-            if const_expr(mC_mnl is not None):
                 epi_producer_state = pipeline.make_pipeline_state(
                     pipeline.PipelineUserType.Producer, self.epi_c_stage
                 )
@@ -1373,22 +1398,41 @@ class GemmSm100(GemmSm90):
                     # Get tile coord from tile scheduler
                     tile_coord_mnkl = work_tile.tile_idx
                     batch_idx = tile_coord_mnkl[3]
-                    copy_C_fn, _, bGS_gC = self.epilog_gmem_copy_and_partition(
-                        tma_atom_c,
-                        varlen_manager.offset_batch_epi(mC_mnl, batch_idx),
-                        self.cta_tile_shape_mnk[:2],
-                        epi_tile,
-                        sC,
                         tile_coord_mnkl,
                     )
-                    copy_C = copy_utils.tma_producer_copy_fn(copy_C_fn, epi_pipeline)
                     if do_epi_load_barrier_wait:
                         epi_load_barrier.arrive_and_wait()
                         do_epi_load_barrier_wait = Boolean(False)
-                    epi_tile_num = const_expr(cute.size(bGS_gC, mode=[1]))
                     for epi_idx in cutlass.range(epi_tile_num, unroll=1):
                         epi_pipeline.producer_acquire(epi_producer_state)
-                        copy_C(src_idx=epi_idx, producer_state=epi_producer_state)
                         # Epi pipeline's producer commit is a NOP
                         epi_pipeline.producer_commit(epi_producer_state)
                         epi_producer_state.advance()
@@ -1446,21 +1490,8 @@ class GemmSm100(GemmSm90):
                     cute.slice_(sfb_smem_layout, (None, None, None, 0)),
                 )
                 tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
-                # Partition for S2T copy of SFA/SFB
-                (
-                    tiled_copy_s2t_sfa,
-                    tCsSFA_compact_s2t,
-                    tCtSFA_compact_s2t,
-                ) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
-                (
-                    tiled_copy_s2t_sfb,
-                    tCsSFB_compact_s2t,
-                    tCtSFB_compact_s2t,
-                ) = self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
             else:
                 tCtSFA, tCtSFB = None, None
-                tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = None, None, None
-                tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = None, None, None
             # Persistent tile scheduling loop
             tile_scheduler = TileSchedulerCls()
@@ -1486,7 +1517,10 @@ class GemmSm100(GemmSm90):
                 )
                 tCtAcc = tCtAcc_base[None, None, None, acc_stage_idx]
                 tCtSFB_mma = tCtSFB
-                if const_expr(self.blockscaled and self.mma_inst_shape_mnk[1] in (64, 192)):
                     tCtSFB_mma = cute.make_tensor(
                         cute.recast_ptr(
                             sfb_tmem_base_ptr + Int32((tile_coord_mnkl[1] % 2) * 2),
@@ -1494,7 +1528,25 @@ class GemmSm100(GemmSm90):
                         ),
                         tCtSFB.layout,
                     )
-                tctx.b("mma")
                 ab_consumer_state, acc_producer_state, tiled_mma = self.mma(
                     ab_pipeline,
                     acc_pipeline,
@@ -1509,17 +1561,14 @@ class GemmSm100(GemmSm90):
                     cta_rank_in_cluster,
                     tCtSFA,
                     tCtSFB_mma,
-                    tiled_copy_s2t_sfa,
-                    tiled_copy_s2t_sfb,
-                    tCsSFA_compact_s2t,
-                    tCsSFB_compact_s2t,
-                    tCtSFA_compact_s2t,
-                    tCtSFB_compact_s2t,
                 )
                 if const_expr(self.overlap_accum_sf):
                     # After iter 0, 2, ..., shift tmem ptr by -256.
                     # After iter 1, 3, ..., shift tmem ptr by 256.
-                    tCtSFA, tCtSFB, tCtSFA_compact_s2t, tCtSFB_compact_s2t = [
                         cute.make_tensor(
                             cute.recast_ptr(
                                 # Doing tmem ptr arithmetic requires 32-bit type, wrong otherwise
@@ -1532,9 +1581,9 @@ class GemmSm100(GemmSm90):
                             ),
                             mT.layout,
                         )
-                        for mT in [tCtSFA, tCtSFB, tCtSFA_compact_s2t, tCtSFB_compact_s2t]
                     ]
-                tctx.e("mma")
                 # Advance to next tile
                 tile_scheduler.advance_to_next_work()
                 work_tile = tile_scheduler.get_current_work()
@@ -1592,6 +1641,13 @@ class GemmSm100(GemmSm90):
                 pipeline.PipelineUserType.Consumer, self.epi_c_stage
             )
             while work_tile.is_valid_tile:
                 # Get tile coord from tile scheduler
                 tile_coord_mnkl = work_tile.tile_idx
                 batch_idx = tile_coord_mnkl[3]
@@ -1635,10 +1691,10 @@ class GemmSm100(GemmSm90):
                     acc_release_idx=self.iter_acc_early_release
                     if const_expr(self.overlap_accum_sf)
                     else epi_tile_num - 1,
-                    clear_acc=varlen_k and k_len == 0,
                 )
-                tctx.b("epilogue")
                 epi_read_state, _ = self.epilogue(
                     epilogue_params,
                     epi_smem_tensors,
@@ -1667,11 +1723,11 @@ class GemmSm100(GemmSm90):
                 )
                 # acc_pipeline.consumer_release was already called in self.epi_load_acc_subtile
                 acc_consumer_state.advance()
-                tctx.e("epilogue")
                 # Advance to next tile
                 tile_scheduler.advance_to_next_work()
-                work_tile = tile_scheduler.get_current_work()
             # Wait for D store complete
             if is_tma_warp:
@@ -1682,8 +1738,6 @@ class GemmSm100(GemmSm90):
             tmem_alloc_barrier.arrive_and_wait()
             tmem.free(acc_tmem_ptr)
-        tctx.flush()
     @cute.jit
     def _make_gather_A_copy(
         self,
@@ -1876,19 +1930,15 @@ class GemmSm100(GemmSm90):
         cta_rank_in_cluster: Int32,
         tCtSFA: Optional[cute.Tensor] = None,
         tCtSFB: Optional[cute.Tensor] = None,
-        tiled_copy_s2t_sfa: Optional[cute.TiledCopy] = None,
-        tiled_copy_s2t_sfb: Optional[cute.TiledCopy] = None,
-        tCsSFA_compact_s2t: Optional[cute.Tensor] = None,
-        tCsSFB_compact_s2t: Optional[cute.Tensor] = None,
-        tCtSFA_compact_s2t: Optional[cute.Tensor] = None,
-        tCtSFB_compact_s2t: Optional[cute.Tensor] = None,
     ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState, cute.TiledMma]:
-        blockscaled = const_expr(tiled_copy_s2t_sfa is not None)
         if const_expr(blockscaled):
             assert all(x is not None for x in (tCtSFA, tCtSFB))
-            assert all(x is not None for x in (tiled_copy_s2t_sfa, tiled_copy_s2t_sfb))
-            assert all(x is not None for x in (tCsSFA_compact_s2t, tCsSFB_compact_s2t))
-            assert all(x is not None for x in (tCtSFA_compact_s2t, tCtSFB_compact_s2t))
         # If gather_A and use_2cta_instrs, the cp.async for the non-leader CTA will
         # arrive at an mbarrier on the non-leader CTA side, then the mma warp of the non-leader
         # CTA will wait for that then arrive at the mbarrier on the leader CTA.
@@ -1911,29 +1961,67 @@ class GemmSm100(GemmSm90):
                 if not is_leader_cta:
                     ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
                     with cute.arch.elect_one():
-                        # The odd CTA signals the even CTA
-                        ab_pipeline.sync_object_full.arrive_mbarrier(
-                            ab_consumer_state.index, dst_rank=cta_rank_in_cluster & 0xFE
                         )
             if is_leader_cta:
                 # Conditionally wait for AB buffer full
                 ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
                 #  Copy SFA/SFB from smem to tmem
                 if const_expr(blockscaled):
-                    s2t_stage_coord = (None, None, None, None, ab_consumer_state.index)
-                    tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
-                    tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
-                    cute.copy(tiled_copy_s2t_sfa, tCsSFA_compact_s2t_staged, tCtSFA_compact_s2t)
-                    cute.copy(tiled_copy_s2t_sfb, tCsSFB_compact_s2t_staged, tCtSFB_compact_s2t)
                 for k_blk_idx in cutlass.range(num_k_blocks, unroll_full=True):
                     k_blk_coord = (None, None, k_blk_idx, ab_consumer_state.index)
                     if const_expr(blockscaled):
                         # Set SFA/SFB tensor to tiled_mma
                         sf_kblock_coord = (None, None, k_blk_idx)
-                        tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
-                        tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock_coord].iterator)
-                    cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc)
-                    tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                 # Async arrive AB buffer empty
                 ab_pipeline.consumer_release(ab_consumer_state)
             ab_consumer_state.advance()
@@ -1957,7 +2045,7 @@ class GemmSm100(GemmSm90):
         tTR_tAcc: cute.Tensor,
         tTR_rAcc: cute.Tensor,
         tRS_rD: cute.Tensor,
-        epi_idx: int,
         acc_pipeline: pipeline.PipelineAsync,
         acc_consumer_state: pipeline.PipelineState,
         acc_release_idx: int,
@@ -1965,50 +2053,15 @@ class GemmSm100(GemmSm90):
     ):
         if not clear_acc:
             # Load accumulator from tensor memory buffer to register
-            cute.copy(tiled_copy_t2r, tTR_tAcc[None, None, None, epi_idx], tTR_rAcc)
             tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
             tRS_rD.store(tRS_rAcc.load())
         else:
             tRS_rD.fill(0.0)
-        if epi_idx == acc_release_idx:
             cute.arch.fence_view_async_tmem_load()
-            with cute.arch.elect_one():
-                acc_pipeline.consumer_release(acc_consumer_state)
-    def mainloop_s2t_copy_and_partition(
-        self,
-        sSF: cute.Tensor,
-        tSF: cute.Tensor,
-    ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
-        """
-        Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination).
-        :param sSF: The scale factor tensor in smem
-        :type sSF: cute.Tensor
-        :param tSF: The scale factor tensor in tmem
-        :type tSF: cute.Tensor
-        :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
-            - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
-            - tCsSF_compact_s2t: The partitioned scale factor tensor in smem
-            - tSF_compact_s2t: The partitioned scale factor tensor in tmem
-        :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
-        """
-        # (MMA, MMA_MN, MMA_K, STAGE)
-        tCsSF_compact = cute.filter_zeros(sSF)
-        # (MMA, MMA_MN, MMA_K)
-        tCtSF_compact = cute.filter_zeros(tSF)
-        # Make S2T CopyAtom and tiledCopy
-        copy_atom_s2t = cute.make_copy_atom(tcgen05.Cp4x32x128bOp(self.cta_group), self.sf_dtype)
-        tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
-        thr_copy_s2t = tiled_copy_s2t.get_slice(0)
-        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
-        tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
-        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
-        tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_)
-        # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
-        tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
-        return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
     def epilog_tmem_copy_and_partition(
         self,
@@ -2142,7 +2195,6 @@ class GemmSm100(GemmSm90):
         self,
         tiled_mma: cute.TiledMma,
         cluster_layout_vmnk: cute.Layout,
-        ab_pipeline_mbar_ptr: cute.Pointer,
         is_leader_cta: Boolean,
     ) -> pipeline.PipelineAsync:
         # If gather_A and use_2cta_instrs, the cp.async for the non-leader CTA will
@@ -2155,9 +2207,11 @@ class GemmSm100(GemmSm90):
         if const_expr(not self.gather_A or self.use_tma_gather):
             producer_cnt = 1
         else:
-            producer_cnt = self.num_ab_load_warps * 32 + (
-                1 if const_expr(not self.use_2cta_instrs) else 2
-            )
         ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt)
         # Each warp will contribute to the arrive count with the number of mcast size
         mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -2165,19 +2219,8 @@ class GemmSm100(GemmSm90):
         ab_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, consumer_arrive_cnt
         )
-        if const_expr(not self.gather_A):
-            pipeline_ab = pipeline.PipelineTmaUmma.create(
-                barrier_storage=ab_pipeline_mbar_ptr,
-                num_stages=self.ab_stage,
-                producer_group=ab_pipeline_producer_group,
-                consumer_group=ab_pipeline_consumer_group,
-                tx_count=self.num_tma_load_bytes,
-                cta_layout_vmnk=cluster_layout_vmnk,
-                defer_sync=True,
-            )
-        elif const_expr(self.use_tma_gather):
             pipeline_ab = PipelineTmaUmma.create(
-                barrier_storage=ab_pipeline_mbar_ptr,
                 num_stages=self.ab_stage,
                 producer_group=ab_pipeline_producer_group,
                 consumer_group=ab_pipeline_consumer_group,
@@ -2187,40 +2230,35 @@ class GemmSm100(GemmSm90):
             )
         else:
             pipeline_ab = PipelineTmaCpAsyncUmma.create(
-                barrier_storage=ab_pipeline_mbar_ptr,
                 num_stages=self.ab_stage,
                 producer_group=ab_pipeline_producer_group,
                 consumer_group=ab_pipeline_consumer_group,
                 tx_count=self.num_tma_load_bytes,
                 cta_layout_vmnk=cluster_layout_vmnk,
-                producer_drop_count=None
-                if not self.use_2cta_instrs
-                else (2 if not is_leader_cta else 0),
                 defer_sync=True,
             )
         return pipeline_ab
-    def make_acc_pipeline(
-        self, cluster_layout_vmnk: cute.Layout, acc_pipeline_mbar_ptr: cute.Pointer
-    ) -> pipeline.PipelineAsync:
         acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
         num_acc_consumer_threads = self.num_epi_warps * (2 if self.use_2cta_instrs else 1)
         acc_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, num_acc_consumer_threads
         )
-        return pipeline.PipelineUmmaAsync.create(
-            barrier_storage=acc_pipeline_mbar_ptr,
             num_stages=self.num_acc_stage,
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
             defer_sync=True,
         )
     def make_sched_pipeline(
         self,
         cluster_layout_mnk: cute.Layout,
-        sched_pipeline_mbar_ptr: cute.Pointer,
         has_C: bool = False,
     ) -> pipeline.PipelineAsync:
         # Threads/warps participating in this pipeline
@@ -2237,32 +2275,40 @@ class GemmSm100(GemmSm90):
         sched_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, consumer_arrive_cnt
         )
-        return pipeline.PipelineAsync.create(
-            barrier_storage=sched_pipeline_mbar_ptr,
             num_stages=self.sched_stage,
             producer_group=sched_pipeline_producer_group,
             consumer_group=sched_pipeline_consumer_group,
             # If there's cluster, the consumers must arrive at the mbar of CTA 0 in the cluster.
             consumer_mask=None if const_expr(cluster_size == 1) else 0,
             defer_sync=True,
         )
     @cute.jit
-    def make_a_prefetch_pipeline(
-        self, a_prefetch_pipeline_mbar_ptr: cute.Pointer
-    ) -> pipeline.PipelineAsync:
         producer_cnt = 32
         a_prefetch_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt)
         consumer_arrive_cnt = self.num_ab_load_warps
         a_prefetch_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, consumer_arrive_cnt
         )
-        return pipeline.PipelineCpAsync.create(
-            barrier_storage=a_prefetch_pipeline_mbar_ptr,
             num_stages=self.a_prefetch_stage,
             producer_group=a_prefetch_producer_group,
             consumer_group=a_prefetch_consumer_group,
             defer_sync=True,
         )
     @classmethod
@@ -2284,6 +2330,7 @@ class GemmSm100(GemmSm90):
         prefetch_A_idx: Literal[None, "varlen_m", "varlen_k"],
         smem_capacity: int,
         occupancy: int,
     ) -> Tuple[int, int, int]:
         """Computes the number of stages for A/B/C operands based on heuristics.
@@ -2319,7 +2366,15 @@ class GemmSm100(GemmSm90):
         # Default D stages
         epi_stage = 4 if cute.size(epi_tile[1]) <= 16 else 2
-        epi_c_stage = 0 if c_dtype is None else (4 if cute.size(epi_tile[1]) <= 16 else 2)
         # Calculate smem layout and size for one stage of A, B, and C
         a_smem_layout_staged_one = sm100_utils.make_smem_layout_a(
@@ -2373,13 +2428,13 @@ class GemmSm100(GemmSm90):
         d_bytes_per_stage = (
             cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) if d_dtype is not None else 0
         )
-        epi_bytes_per_stage = d_bytes_per_stage + cls.epi_smem_bytes_per_stage(
-            epilogue_args, cta_tile_shape_mnk, epi_tile
-        )
-        epi_bytes = epi_bytes_per_stage * epi_stage
         if const_expr(c_dtype is not None):
             c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
             epi_bytes += c_bytes_per_stage * epi_c_stage
         # Calculate A/B/SFA/SFB stages:
         # Start with total smem per CTA (capacity / occupancy)
@@ -2391,7 +2446,8 @@ class GemmSm100(GemmSm90):
         # Refine epilogue stages:
         # Calculate remaining smem after allocating for A/B stages and reserved bytes
         # Add remaining unused smem to epilogue
-        epi_stage += (remaining_bytes - ab_bytes_per_stage * ab_stage) // (epi_bytes_per_stage)
         return num_acc_stage, ab_stage, epi_stage, epi_c_stage
     @staticmethod
@@ -2553,15 +2609,15 @@ class GemmSm100(GemmSm90):
     @staticmethod
     def is_valid_mma_tiler_and_cluster_shape(
-        mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mn: Tuple[int, int],
         blockscaled: bool,
     ) -> bool:
         """
         Check if the mma tiler and cluster shape are valid
-        :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
-        :type mma_tiler_mn: Tuple[int, int]
         :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
         :type cluster_shape_mn: Tuple[int, int]
@@ -2571,20 +2627,20 @@ class GemmSm100(GemmSm90):
         is_valid = True
         # Skip invalid mma tile shape
         if not blockscaled:
-            if mma_tiler_mn[0] not in [64, 128, 256]:
                 is_valid = False
         else:
-            if mma_tiler_mn[0] not in [128, 256]:
                 is_valid = False
-        mma_inst_n = mma_tiler_mn[1] if mma_tiler_mn[1] <= 256 else mma_tiler_mn[1] // 2
         if not blockscaled:
             if mma_inst_n not in range(32, 257, 32):
                 is_valid = False
         else:
             # Blockscaled currently supports tile_n in {64, 128, 192, 256}.
-            if mma_tiler_mn[1] not in [64, 128, 192, 256]:
                 is_valid = False
-        if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0:
             is_valid = False
         # Skip invalid cluster shape
         is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0
@@ -2662,7 +2718,7 @@ class GemmSm100(GemmSm90):
         sf_dtype: Type[cutlass.Numeric],
         sf_vec_size: int,
         d_dtype: Type[cutlass.Numeric],
-        mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mn: Tuple[int, int],
         m: int,
         n: int,
@@ -2680,13 +2736,9 @@ class GemmSm100(GemmSm90):
         if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"):
             can_implement = False
         if not GemmSm100.is_valid_mma_tiler_and_cluster_shape(
-            mma_tiler_mn, cluster_shape_mn, blockscaled=True
         ):
             can_implement = False
-        # Multi-tile N iteration with an asymmetric SFB atom size needs the same
-        # kind of special-case layout rewriting as tile_n==192.
-        if mma_tiler_mn[1] == 224 and n > 224:
-            can_implement = False
         if not GemmSm100.is_valid_tensor_alignment(
             m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major
         ):
@@ -2698,7 +2750,7 @@ class GemmSm100(GemmSm90):
         ab_dtype: Type[cutlass.Numeric],
         acc_dtype: Type[cutlass.Numeric],
         d_dtype: Type[cutlass.Numeric],
-        mma_tiler_mn: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mn: Tuple[int, int],
         m: int,
         n: int,
@@ -2717,8 +2769,8 @@ class GemmSm100(GemmSm90):
         :type acc_dtype: Type[cutlass.Numeric]
         :param d_dtype: The data type of the output tensor
         :type d_dtype: Type[cutlass.Numeric]
-        :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler
-        :type mma_tiler_mn: Tuple[int, int]
         :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
         :type cluster_shape_mn: Tuple[int, int]
         :param m: The number of rows in the A tensor
@@ -2745,7 +2797,7 @@ class GemmSm100(GemmSm90):
             can_implement = False
         # Skip invalid mma tile shape and cluster shape
         if not GemmSm100.is_valid_mma_tiler_and_cluster_shape(
-            mma_tiler_mn, cluster_shape_mn, blockscaled=False
         ):
             can_implement = False
         # Skip illegal problem shape for load/store alignment

+# Copyright (c) 2025-2026, QuACK team.
 # Based on the cute-dsl example:
 # https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
 )
 from cutlass import Int32, Float32, Boolean, const_expr
 from cutlass.utils import LayoutEnum
+from cutlass.cute.experimental import iket
+from .pipeline import (
+    PipelineAsync as QuackPipelineAsync,
+    PipelineCpAsync,
+    PipelineTmaUmma,
+    PipelineTmaCpAsyncUmma,
+    PipelineUmmaAsync,
+    mbarrier_arrive_release_cluster,
+    mbarrier_acquire_cluster,
+)
+from .dsl.smem_struct import Reserved, partitioned_struct
 from .tile_scheduler import TileSchedulerOptions
 from .varlen_utils import VarlenArguments, VarlenManager
+from .gemm_base import GemmTmaBase, NamedBarrierGemm
 from . import layout_utils
 from . import copy_utils as copy_utils
 from . import sm100_utils as quack_sm100_utils
 - Write accumulator to TMEM
 The accumulator in TMEM must then be loaded to registers before writing back to GMEM.
+Constraints:
 * Supported input data types: fp16, bf16, tf32, int8, uint8, fp8 (e4m3fn, e5m2),
   see detailed valid dtype combinations in below GemmSm100 class documentation
 * A/B tensor must have the same data type
 """
+class GemmSm100(GemmTmaBase):
     """This class implements batched matrix multiplication (C = A x B) with support for various data types
     and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
     :param acc_dtype: Data type for accumulation during computation
     :type acc_dtype: type[cutlass.Numeric]
+    :param mma_tiler_mnk: Shape of the MMA tile. Pass (M, N) to default K to
         4 MMA instructions, or (M, N, K) to set the K tile size explicitly.
+    :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]]
     :param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
     :type cluster_shape_mn: Tuple[int, int]
     Example:
         >>> gemm = GemmSm100(
         ...     acc_dtype=Float32,
+        ...     mma_tiler_mnk=(128, 128),
         ...     cluster_shape_mn=(2, 2)
         ... )
         >>> gemm(mA, mB, mD, max_active_clusters, stream)
     arch = 100
+    EpilogueArguments = GemmTmaBase.EpilogueArguments
+    EpilogueParams = GemmTmaBase.EpilogueParams
     def __init__(
         self,
         acc_dtype: Type[cutlass.Numeric],
         a_dtype: Type[cutlass.Numeric],  # ignored for now
+        mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mnk: Tuple[int, int, int],
         sf_vec_size: Optional[int] = None,
         gather_A: bool = False,
         1.  MMA Instruction Settings (tcgen05):
             - acc_dtype: Data types for MMA accumulator.
+            - mma_tiler_mnk: The (M, N) or (M, N, K) shape of the MMA instruction tiler.
             - use_2cta_instrs: Boolean indicating if the tcgen05 MMA variant
               with cta_group=2 should be used.
         :param acc_dtype: Data type of the accumulator.
         :type acc_dtype: type[cutlass.Numeric]
+        :param mma_tiler_mnk: (M, N) or (M, N, K) shape of the MMA tile.
             If only (M, N) is given, K defaults to 4 * instruction K.
+        :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]]
         :param cluster_shape_mnk: Tuple (ClusterM, ClusterN) shape of the cluster.
         :type cluster_shape_mnk: Tuple[int, int]
         """
         self.acc_dtype: Type[cutlass.Numeric] = acc_dtype
+        self.sf_vec_size = sf_vec_size
+        self.blockscaled = sf_vec_size is not None
+        assert len(mma_tiler_mnk) in [2, 3], "MMA tiler must be (M, N) or (M, N, K)"
+        valid_2cta_m = (128, 256) if not self.blockscaled else (256,)
+        self.use_2cta_instrs = cluster_shape_mnk[0] % 2 == 0 and mma_tiler_mnk[0] in valid_2cta_m
         self.cluster_shape_mnk = cluster_shape_mnk
         assert cluster_shape_mnk[2] == 1, "Cluster shape K must be 1"
         # K dimension: if user provides 3 values, use their K; otherwise default in _setup_attributes
+        if len(mma_tiler_mnk) == 3:
+            self.mma_tiler = tuple(mma_tiler_mnk)
         else:
+            self.mma_tiler = (*mma_tiler_mnk, 0)
         self.is_persistent = True
         self.use_clc_persistence = use_clc_persistence
+        self.epi_m_major = True
         self.gather_A = gather_A
         self.concat_layout = concat_layout or ()
         self.use_tma_gather = use_tma_gather
             barrier_id=int(NamedBarrierGemm.Epilogue),
             num_threads=self.num_epi_warps * cute.arch.WARP_SIZE,
         )
+        # CLC throttle: paces query issue to tile consumption so the multi-stage
+        # lookahead can't over-cancel the pending pool. Producer = CTA0 load warp
+        # (arrive per tile started), consumer = CTA0 scheduler warp (sync per
+        # query); 2 warps => 64 threads. Lives here (not the scheduler) because a
+        # NamedBarrier id is a whole-CTA resource coordinated by NamedBarrierGemm,
+        # and the participating-thread count is arch-specific (warp layout). A
+        # single barrier suffices: the dependency chain commit(k+1) <- fetch(k+1)
+        # <- query(k+1) <- sync(k) forces strict producer/consumer alternation, so
+        # <= 1 credit is ever outstanding; bar.sync also gives a hardware wakeup vs
+        # an mbarrier pipeline's PHASECHK + NANOSLEEP polling.
+        self.clc_throttle_barrier = (
+            pipeline.NamedBarrier(
+                barrier_id=int(NamedBarrierGemm.ClcThrottle), num_threads=2 * cute.arch.WARP_SIZE
+            )
+            if self.use_clc_persistence
+            else None
+        )
         # Register reallocation for gather_A (3 warp groups, 504 regs total, 168 per WG default).
         # Heavy epilogues (e.g. colvec_reduce in DGated) override these to avoid register spilling.
         # Without gather_A there are only 2 WGs (512 total, 256 per WG = max), no reallocation needed.
         # Multiple of 4 warps to increase/decrease number of registers
         assert self.threads_per_cta % 128 == 0
+    def epi_smem_warp_shape_mnk(self):
+        # Mirrors cutlass.utils.blackwell_helpers.compute_epilogue_tile_shape:
+        # the epilogue tmem layout uses two M warps and two N warps when the
+        # per-CTA M tile is 64 and the kernel uses 2-CTA instructions.
+        warp_m, warp_n = (
+            (2, 2) if self.cta_tile_shape_mnk[0] == 64 and self.use_2cta_instrs else (4, 1)
+        )
+        return (warp_m, warp_n, 1)
     def _setup_attributes(self, epilogue_args: EpilogueArguments, varlen_args: VarlenArguments):
         """Set up configurations that are dependent on GEMM inputs
         - Computing A/B/C shared memory layout
         - Computing tensor memory allocation columns
         """
+        self.epi_m_major = self.resolve_epi_m_major(epilogue_args)
         # Compute mma instruction shapes
         mma_inst_bits_k = 256
         # (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K)
         if const_expr(not self.blockscaled):
             self.tiled_mma = sm100_utils.make_trivial_tiled_mma(
                 self.a_dtype,
+                self.b_dtype,
                 self.a_major_mode,
                 self.b_major_mode,
                 self.acc_dtype,
         else:
             self.tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
                 self.a_dtype,
+                self.b_dtype,
                 self.a_major_mode,
                 self.b_major_mode,
                 self.sf_dtype,
             )
             self.tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
                 self.a_dtype,
+                self.b_dtype,
                 self.a_major_mode,
                 self.b_major_mode,
                 self.sf_dtype,
         # Compute mma/cluster/tile shapes
         if self.mma_tiler[2] > 0:
+            assert self.mma_tiler[2] % self.mma_inst_shape_mnk[2] == 0, (
+                f"MMA tiler K ({self.mma_tiler[2]}) must be divisible by "
+                f"MMA instruction K ({self.mma_inst_shape_mnk[2]})"
+            )
             mma_inst_tile_k = self.mma_tiler[2] // self.mma_inst_shape_mnk[2]
         else:
             mma_inst_tile_k = 4
                 self.mma_tiler_sfb[1],
                 self.mma_tiler_sfb[2],
             )
+            # The SF atom fixed by the tcgen05 MMA (BlockScaledBasicChunk) is 128 wide
+            # in N, but cta_tile_n need not be a multiple of 128. Two derived
+            # quantities localize all the resulting special handling:
+            # - sfb_tiles_per_atom: adjacent N-tiles that share one 128-wide atom
+            #   (tile_n=64) load the same SFB atom; gmem N-tile coords are divided
+            #   by this.
+            # - sfb_n_atom_misaligned: tile_n an odd multiple of 64 (64, 192) puts
+            #   odd N-tiles 64 into an atom; the MMA's SFB tmem base shifts by 2
+            #   columns for odd N-tile coords, and tile_n=192 additionally needs
+            #   the overlapped-window TMA remap at the SFB TMA setup in __call__.
+            # tile_n=224 is rejected: its tiles start at 32-column offsets within
+            # the atom ((224*j) % 128 cycles 0/96/64/32), which neither mechanism
+            # covers.
+            assert self.cta_tile_shape_mnk[1] in (64, 128, 192, 256), (
+                f"blockscaled tile_n must be in (64, 128, 192, 256), "
+                f"got {self.cta_tile_shape_mnk[1]}"
+            )
+            self.sfb_tiles_per_atom = max(128 // self.cta_tile_shape_mnk[1], 1)
+            self.sfb_n_atom_misaligned = (self.cta_tile_shape_mnk[1] // 64) % 2 == 1
         else:
             self.cta_tile_shape_mnk_sfb = None
             self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
         # Compute epilogue subtile
+        tile_load_layout = None
+        tile_load_dtype = None
+        # If TileLoad exists without C, use the first non-None tile-load tensor as
+        # the C-like input for SM100's epilogue tile shape. Multiple TileLoads
+        # share the same epi_tile shape.
+        for op in getattr(self, "_epi_ops", ()):
+            if op.is_tile_load():
+                tile_load_tensor = getattr(epilogue_args, op.name, None)
+                if tile_load_tensor is not None:
+                    tile_load_layout = LayoutEnum.from_tensor(tile_load_tensor)
+                    tile_load_dtype = tile_load_tensor.element_type
+                    break
         self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
             self.cta_tile_shape_mnk,
             self.use_2cta_instrs,
             self.d_layout if self.d_layout is not None else LayoutEnum.ROW_MAJOR,
             self.d_dtype if self.d_dtype is not None else cutlass.BFloat16,
+            layout_c=self.c_layout if self.c_layout is not None else tile_load_layout,
+            elem_ty_c=self.c_dtype if self.c_dtype is not None else tile_load_dtype,
         )
         # TMA store tile starts must stay aligned when advancing across CTA-N tiles.
         # There's a bug w compute_epilogue_tile_shape (as of cutlass-dsl 4.4.2) where if
             prefetch_A_idx,
             cutlass.utils.get_smem_capacity_in_bytes(f"sm_{self.arch}"),  # smem_capacity
             self.occupancy,
+            self.epi_smem_warp_shape_mnk(),
         )
+        # With CLC the try_cancel response lands directly in the consumer slot, so
+        # the next query can only be issued once all consumers (cluster-wide)
+        # release that slot. >=2 stages keep a query in flight while the previous
+        # tile's info is still being read (cutlass's SchedulerPipelineStageCount
+        # >= 2); the 3rd stage buys response slack for epilogue-bound tiles (e.g.
+        # symmetric's double store, ~3% at M=8192 K=512) and costs only 12 smem
+        # ints + one mbarrier pair.
+        self.sched_stage = 3 if self.use_clc_persistence else 1
         self.a_prefetch_stage = (
             0
             if not self.gather_A
         stream: cuda.CUstream,
         mSFA: Optional[cute.Tensor] = None,
         mSFB: Optional[cute.Tensor] = None,
     ):
         """Execute the GEMM operation in steps:
         - Setup static attributes before smem/grid/tma computation
             # so non-packed buffers work (e.g. a slice of a larger scale tensor).
             # Only the innermost 512-B tile must be contiguous.
             # For varlen_m, mSFA is sized for per-expert 128-row-padded storage
+            # (tile-aligned per-batch padding), so use its own M dim (= total_padded_rm * 128)
             # instead of mA.shape[0] (= total_m, unpadded).
             if const_expr(cute.rank(mA) == 3):
                 sfa_shape = mA.shape
         a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
         b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
         tma_atom_a, tma_tensor_a = None, None
+        a_op = (
+            cpasync.CopyBulkTensorTileG2SOp(self.cta_group)
+            if const_expr(not self.gather_A)
+            else sm100_utils.cluster_shape_to_tma_atom_A(
+                self.cluster_shape_mnk, self.tiled_mma.thr_id
+            )
         )
         if const_expr(not self.gather_A):
             tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
                 tma_smem_layout.shape,
                 internal_type=(cutlass.TFloat32 if mA.element_type is Float32 else None),
             )
+        # block_copy takes compiler-driven multicast metadata at the copy site,
+        # so the TMA atom itself must stay the non-multicast variant here.
+        b_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group)
         tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
             b_op,
             copy_utils.create_ragged_tensor_for_tma(mB, ragged_dim=1) if varlen_k else mB,
         tma_atom_sfb, tma_tensor_sfb = None, None
         if const_expr(self.blockscaled):
             # Setup TMA load for SFA
+            sfa_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group)
             sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0))
             tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
                 sfa_op,
                 internal_type=cutlass.Int16,
             )
             # Setup TMA load for SFB
+            sfb_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group)
             sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
             tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
                 sfb_op,
                 self.cluster_layout_sfb_vmnk.shape,
                 internal_type=cutlass.Int16,
             )
+            # tile_n=192 spans 1.5 SF atoms, so consecutive N-tiles straddle atom
+            # boundaries and a TMA box can't be placed at a half-atom offset.
+            # Instead each tile loads a 2-atom (256-wide) window: remap the gmem
+            # atom sequence to [a0 a1 | a1 a2 | a3 a4 | a4 a5 | ...] (groups of 4
+            # presented atoms at offsets (0, x, x, 2x), advancing by 3 atoms) so
+            # that tile j's window lands on atoms (3j//2, 3j//2 + 1). Odd tiles
+            # start 64 into their first atom; the mma warp corrects for that via
+            # the sfb_n_atom_misaligned tmem offset.
+            if const_expr(self.cta_tile_shape_mnk[1] == 192):
                 x = tma_tensor_sfb.stride[0][1]
                 y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4)
                 tma_tensor_sfb = cute.make_tensor(
             self.num_tma_load_bytes += sfa_copy_size + sfb_copy_size
         self.num_tma_load_bytes *= atom_thr_size
+        # Setup TMA store for D and TMA load for C.
+        tma_atom_d, tma_tensor_d, tma_atom_c, tma_tensor_c = (
+            self.make_tma_epilogue_atoms_and_tensors(mD, mC, epilogue_args, varlen_m)
+        )
         epilogue_params = self.epi_to_underlying_arguments(epilogue_args)
         varlen_params = VarlenManager.to_underlying_arguments(varlen_args)
+        self.epi_load_bytes_per_stage = self.epi_smem_bytes(
+            epilogue_args,
+            self.cta_tile_shape_mnk,
+            self.epi_tile,
+            self.epi_smem_warp_shape_mnk(),
+        ).c_stage
+        if const_expr(mC is not None):
+            c_smem_layout = cute.slice_(self.epi_c_smem_layout_staged, (None, None, 0))
+            self.epi_load_bytes_per_stage += cute.size_in_bytes(self.c_dtype, c_smem_layout)
         TileSchedulerCls = self.get_scheduler_class(varlen_m=varlen_m)
         tile_sched_args = self.get_scheduler_arguments(
             mA, mB, mD, scheduler_args, varlen_args, epilogue_args
                 self.cta_tile_shape_mnk[0] if varlen_m else self.cta_tile_shape_mnk[2]
             )
+        # Define shared storage for kernel. sched_data lives in the RESERVED
+        # smem partition (with the pipeline mbarriers / TMEM holding buf): a
+        # small buffer before the 1024-byte aligned epilogue tensors would add
+        # a 1 KiB pad; CLC responses use i128 copies, so it stays 16-byte
+        # aligned.
+        # 4 Int32 per stage, shared by the two (mode-exclusive) users: STATIC/DYNAMIC
+        # store the STAS-broadcast (pid_m, pid_n, batch_idx, is_valid); CLC stores the
+        # 16-byte try_cancel response (16B-aligned since each stage slot is 16 bytes).
+        sched_smem_size = 4 * self.sched_stage if self.is_persistent else 0
+        @partitioned_struct
         class SharedStorage:
+            sched_data: Reserved[
+                cute.struct.Align[cute.struct.MemRange[Int32, sched_smem_size], 16]
             ]
             sAIdx: cute.struct.Align[cute.struct.MemRange[Int32, a_idx_smem_size], 16]
             # (EPI_TILE_M, EPI_TILE_N, STAGE)
             sD: cute.struct.Align[
             self.epi_tile,
             tile_sched_params,
             TileSchedulerCls,
         ).launch(
             grid=grid,
             block=[self.threads_per_cta, 1, 1],
         epi_tile: cute.Tile,
         tile_sched_params,
         TileSchedulerCls: cutlass.Constexpr[Callable],
     ):
         """
         GPU device kernel performing the Persistent batched GEMM computation.
         """
         varlen_m = const_expr(varlen_params.cu_seqlens_m is not None)
         varlen_k = const_expr(varlen_params.cu_seqlens_k is not None)
         assert not (varlen_m and varlen_k)
             assert varlen_m or varlen_k
         has_D = const_expr(mD_mnl is not None)
         has_C = const_expr(mC_mnl is not None)
+        has_epi_load = const_expr(self.epi_c_stage > 0)
         warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
         # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
         smem = cutlass.utils.SmemAllocator()
+        storage = self.shared_storage.allocate(smem)
         # Initialize pipelines and states
         ab_pipeline = self.make_ab_pipeline(
             tiled_mma=tiled_mma,
             cluster_layout_vmnk=cluster_layout_vmnk,
             is_leader_cta=is_leader_cta,
         )
         epi_pipeline = None
+        if const_expr(has_epi_load):
+            epi_pipeline = self.make_epi_pipeline(tx_count=self.epi_load_bytes_per_stage)
+        acc_pipeline = self.make_acc_pipeline(cluster_layout_vmnk=cluster_layout_vmnk)
         sched_pipeline = None
         sched_data = None
         if const_expr(self.is_persistent):
+            sched_pipeline = self.make_sched_pipeline(self.cluster_shape_mnk, has_C=has_epi_load)
+            sched_data = storage.sched_data.get_tensor(cute.make_layout((4, self.sched_stage)))
         a_prefetch_pipeline = None
         if const_expr(self.gather_A):
+            a_prefetch_pipeline = self.make_a_prefetch_pipeline()
         tmem_alloc_barrier = pipeline.NamedBarrier(
             barrier_id=int(NamedBarrierGemm.TmemPtr),
         )
         # Tensor memory dealloc barrier init
         tmem = cutlass.utils.TmemAllocator(
             barrier_for_retrieve=tmem_alloc_barrier,
             allocator_warp_id=self.epilog_warp_id[0],
             is_two_cta=use_2cta_instrs,
         )
         # Cluster arrive after barrier init
         )
         TileSchedulerCls = partial(
+            TileSchedulerCls.create,
+            tile_sched_params,
+            sched_data,
+            sched_pipeline,
+            throttle_barrier=self.clc_throttle_barrier,
         )
         epi_load_barrier = None
+        if const_expr(has_epi_load):
             epi_load_barrier = pipeline.NamedBarrier(
+                barrier_id=int(NamedBarrierGemm.EpilogueLoad),
+                num_threads=(self.num_ab_load_warps + 1) * cute.arch.WARP_SIZE,
             )
         # Cluster wait before tensor memory alloc
                 cute.arch.griddepcontrol_wait()
             if const_expr(self.gather_A):
                 cute.arch.setmaxregister_decrease(self.num_regs_other)
             # Persistent tile scheduling loop
             tile_scheduler = TileSchedulerCls()
             work_tile = tile_scheduler.initial_work_tile_info()
                 pipeline.PipelineUserType.Consumer, self.a_prefetch_stage
             )
             do_epi_load_barrier_arrive = Boolean(True)
+            # CLC throttle producer: only the first load warp of CTA 0 in the
+            # cluster signals; commit once per work tile started, or the scheduler
+            # warp starves of credits.
+            is_throttle_producer = Boolean(warp_idx == self.ab_load_warp_id)
+            if const_expr(cute.size(cluster_layout_vmnk) > 1):
+                is_throttle_producer = is_throttle_producer & Boolean(
+                    cute.arch.block_idx_in_cluster() == 0
+                )
             while work_tile.is_valid_tile:
+                tile_scheduler.throttle_producer_commit(is_throttle_producer)
                 tile_coord_mnkl = work_tile.tile_idx
                 batch_idx = tile_coord_mnkl[3]
                 # Local_tile partition global tensors
                 )
                 if const_expr(self.blockscaled):
                     # (bM, bK)
+                    # SFA uses the tile-aligned per-batch offset (padded SF layout), not
                     # the A-data offset — allows varlen_m seqlens that aren't
                     # multiples of 128.
                     gSFA_mkl = cute.local_tile(
                         (mma_tile_coord_mnl[0], None),
                     )
                     # (bN, bK)
+                    # SFB uses the tile-aligned per-batch K offset in varlen_k (padded SF layout).
+                    # N-tiles sharing one 128-wide SF atom (tile_n=64) load the same
+                    # atom, so the gmem N-tile coord is divided by sfb_tiles_per_atom.
                     gSFB_nkl = cute.local_tile(
                         varlen_manager.offset_batch_SFB(mSFB_nkl, batch_idx),
                         cute.select(self.mma_tiler_sfb, [1, 2]),
+                        (mma_tile_coord_mnl[1] // self.sfb_tiles_per_atom, None),
                     )
                 # Partition global tensor for TiledMMA_A/B/D
                 # Then partition global/shared tensor for TMA load A/B
                 len_k = varlen_manager.len_k(batch_idx)
+                # block_copy's lowering wants the coordinate held fixed by the
+                # multicast mask: A/SFA are same-M across N peers, while B/SFB
+                # are same-N across M peers. Degenerate cluster dimensions are
+                # left for the compiler lowering to simplify.
+                a_tma_multicast = {
+                    "cluster_shape": self.cluster_shape_mnk[:2],
+                    "multicast_dim": "M",
+                }
+                b_tma_multicast = {
+                    "cluster_shape": self.cluster_shape_mnk[:2],
+                    "multicast_dim": "N",
+                }
                 copy_A, prefetch_A = None, None
                 if const_expr(not self.gather_A):
                     # (MMA, MMA_M, MMA_K, RestK)
                     tCgA = thr_mma.partition_A(gA_mk)
+                    copy_A = copy_utils.tma_get_block_copy_fn(
+                        tma_atom_a, src_tensor=tCgA, dst_tensor=sA, tma_multicast=a_tma_multicast
                     )
                 else:
                     # For varlen_m paths (TMA or cp.async): consume indices from
                         warp_idx,
                     )
                     if const_expr(varlen_m):
+                        a_prefetch_pipeline.consumer_release(a_prefetch_consumer_state)
                         a_prefetch_consumer_state.advance()
                     if const_expr(prefetch_A is not None):
                         prefetch_A = partial(prefetch_A, a_prefetch_pipeline)
                     # (MMA, MMA_N, MMA_K)
                     tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
                 # TMA load B partition_S/D
+                copy_B = copy_utils.tma_get_block_copy_fn(
+                    tma_atom_b, src_tensor=tCgB, dst_tensor=sB, tma_multicast=b_tma_multicast
                 )
                 copy_SFA, copy_SFB = None, None
                 if const_expr(self.blockscaled):
                     #  TMA load SFA partition_S/D
+                    copy_SFA = copy_utils.tma_get_block_copy_fn(
                         tma_atom_sfa,
                         src_tensor=tCgSFA,
                         dst_tensor=sSFA,
+                        tma_multicast=a_tma_multicast,
                     )
                     # TMA load SFB partition_S/D
+                    copy_SFB = copy_utils.tma_get_block_copy_fn(
                         tma_atom_sfb,
                         src_tensor=tCgSFB,
                         dst_tensor=sSFB,
+                        tma_multicast=b_tma_multicast,
                     )
                 k_tile_cnt = cute.ceil_div(len_k, self.cta_tile_shape_mnk[2])
+                iket.range_push("tma_load")
                 if const_expr(not self.gather_A):
+                    ab_producer_state = self.load_tma(
                         ab_pipeline,
                         ab_producer_state,
+                        [copy_A, copy_B, copy_SFA, copy_SFB],
                         k_tile_cnt,
                     )
                 elif const_expr(self.use_tma_gather):
                     ab_producer_state, a_prefetch_consumer_state = self.load_AB_tma_gather(
                         copy_B,
                         k_tile_cnt,
                     )
+                iket.range_pop()
                 if const_expr(epi_load_barrier is not None):
                     # In the first work tile, the epi load warp will wait for the signal
                     # from the mainloop load warp to start loading C, to avoid interfering
                         epi_load_barrier.arrive()
                         do_epi_load_barrier_arrive = Boolean(False)
                 # Advance to next tile
+                iket.range_push("sched_fetch")
                 tile_scheduler.advance_to_next_work()
                 work_tile = tile_scheduler.get_current_work()
+                iket.range_pop()
             # Wait A/B buffer empty
             if warp_idx == self.ab_load_warp_id:
                 ab_pipeline.producer_tail(ab_producer_state)
                 work_tile = tile_scheduler.initial_work_tile_info()
                 while work_tile.is_valid_tile:
                     # Advance to next tile
+                    iket.range_push("clc_produce")
                     tile_scheduler.advance_to_next_work(is_scheduler_warp=is_scheduler_warp)
+                    iket.range_pop()
+                    iket.range_push("clc_consume")
                     work_tile = tile_scheduler.get_current_work()
+                    iket.range_pop()
                     # End of persistent scheduler loop
                 if is_scheduler_warp:
                     tile_scheduler.producer_tail()
+                    # Drain the pending-cluster tail (varlen padding) with unobserved
+                    # cancels so it never launches; see cancel_pending_tail for the
+                    # grant-monotonicity assumption this relies on.
+                    tile_scheduler.cancel_pending_tail()
         # Specialized A-index prefetch warp (gather_A only)
         if const_expr(self.gather_A):
             if const_expr(self.gather_A):
                 cute.arch.setmaxregister_decrease(self.num_regs_other)
             # PDL: wait for prior kernel before any C TMA loads (matches cutlass C++ epi_load)
+            if const_expr(self.use_pdl and has_epi_load):
                 cute.arch.griddepcontrol_wait()
+            if const_expr(has_epi_load):
                 epi_producer_state = pipeline.make_pipeline_state(
                     pipeline.PipelineUserType.Producer, self.epi_c_stage
                 )
                     # Get tile coord from tile scheduler
                     tile_coord_mnkl = work_tile.tile_idx
                     batch_idx = tile_coord_mnkl[3]
+                    copy_C = None
+                    if const_expr(has_C):
+                        copy_C_fn, _, _ = self.epilog_gmem_copy_and_partition(
+                            tma_atom_c,
+                            varlen_manager.offset_batch_epi(mC_mnl, batch_idx),
+                            self.cta_tile_shape_mnk[:2],
+                            epi_tile,
+                            sC,
+                            tile_coord_mnkl,
+                        )
+                        copy_C = copy_utils.tma_producer_copy_fn(copy_C_fn, epi_pipeline)
+                    tile_load_copy_fns = self.epi_tile_load_g2s_copy_fns(
+                        epilogue_params,
+                        epi_smem_tensors,
                         tile_coord_mnkl,
+                        varlen_manager,
+                        epi_pipeline,
+                    )
+                    copy_epi_load = copy_utils.chain_tma_producer_copy_fns(
+                        (copy_C, *tile_load_copy_fns)
                     )
                     if do_epi_load_barrier_wait:
                         epi_load_barrier.arrive_and_wait()
                         do_epi_load_barrier_wait = Boolean(False)
+                    epi_tile_num = const_expr(
+                        cute.size(
+                            cute.zipped_divide(
+                                cute.make_layout(self.cta_tile_shape_mnk[:2]), epi_tile
+                            ),
+                            mode=[1],
+                        )
+                    )
                     for epi_idx in cutlass.range(epi_tile_num, unroll=1):
                         epi_pipeline.producer_acquire(epi_producer_state)
+                        copy_epi_load(src_idx=epi_idx, producer_state=epi_producer_state)
                         # Epi pipeline's producer commit is a NOP
                         epi_pipeline.producer_commit(epi_producer_state)
                         epi_producer_state.advance()
                     cute.slice_(sfb_smem_layout, (None, None, None, 0)),
                 )
                 tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
             else:
                 tCtSFA, tCtSFB = None, None
             # Persistent tile scheduling loop
             tile_scheduler = TileSchedulerCls()
                 )
                 tCtAcc = tCtAcc_base[None, None, None, acc_stage_idx]
                 tCtSFB_mma = tCtSFB
+                if const_expr(self.blockscaled and self.sfb_n_atom_misaligned):
+                    # Odd N-tiles start 64 into a 128-wide SF atom: shift the SFB
+                    # tmem base by 2 columns (in the atom layout (32,4):(16,4),
+                    # N+64 is element offset 8 = 2 tmem columns).
                     tCtSFB_mma = cute.make_tensor(
                         cute.recast_ptr(
                             sfb_tmem_base_ptr + Int32((tile_coord_mnkl[1] % 2) * 2),
                         ),
                         tCtSFB.layout,
                     )
+                copy_s2t_sfa, copy_s2t_sfb = None, None
+                sf_valid_insts = None
+                if const_expr(self.blockscaled):
+                    copy_s2t_sfa = copy_utils.s2t_get_copy_fn(sSFA, tCtSFA, self.cta_group)
+                    copy_s2t_sfb = copy_utils.s2t_get_copy_fn(sSFB, tCtSFB, self.cta_group)
+                    # Exploits the fact that for mxfp8 the MMA instruction K size
+                    # equals the SF vec size (== 32), so one instruction consumes
+                    # exactly one SF block and the mma loop can skip the
+                    # instructions for SF pad blocks on a ragged-K last tile (see
+                    # the comment in self.mma). fp4 has inst_k 64 spanning
+                    # multiple SF blocks, but we don't do varlen_k for
+                    # mxfp4/nvfp4. Valid instructions in that tile; % maps
+                    # "aligned or full tile" to 0 = nothing to skip.
+                    if const_expr(self.mma_inst_shape_mnk[2] == self.sf_vec_size):
+                        num_insts = self.mma_tiler[2] // self.mma_inst_shape_mnk[2]
+                        sf_valid_insts = (
+                            cute.ceil_div(k_len % self.mma_tiler[2], self.sf_vec_size) % num_insts
+                        )
+                iket.range_push("mma")
                 ab_consumer_state, acc_producer_state, tiled_mma = self.mma(
                     ab_pipeline,
                     acc_pipeline,
                     cta_rank_in_cluster,
                     tCtSFA,
                     tCtSFB_mma,
+                    copy_s2t_sfa,
+                    copy_s2t_sfb,
+                    sf_valid_insts,
                 )
                 if const_expr(self.overlap_accum_sf):
                     # After iter 0, 2, ..., shift tmem ptr by -256.
                     # After iter 1, 3, ..., shift tmem ptr by 256.
+                    tCtSFA, tCtSFB = [
                         cute.make_tensor(
                             cute.recast_ptr(
                                 # Doing tmem ptr arithmetic requires 32-bit type, wrong otherwise
                             ),
                             mT.layout,
                         )
+                        for mT in [tCtSFA, tCtSFB]
                     ]
+                iket.range_pop()
                 # Advance to next tile
                 tile_scheduler.advance_to_next_work()
                 work_tile = tile_scheduler.get_current_work()
                 pipeline.PipelineUserType.Consumer, self.epi_c_stage
             )
             while work_tile.is_valid_tile:
+                # Prefetch the next work tile before the epilogue: the response is
+                # already in smem (3-stage sched pipeline), and consuming it here
+                # hides the ~300ns decode (swizzle + async fence) behind this tile's
+                # epilogue — the pacing chain for small-K / double-store epilogues.
+                # advance_to_next_work stays after the body: num_tiles_executed must
+                # count completed tiles during the body (sD stage cycling).
+                next_work_tile = tile_scheduler.get_current_work()
                 # Get tile coord from tile scheduler
                 tile_coord_mnkl = work_tile.tile_idx
                 batch_idx = tile_coord_mnkl[3]
                     acc_release_idx=self.iter_acc_early_release
                     if const_expr(self.overlap_accum_sf)
                     else epi_tile_num - 1,
+                    clear_acc=(varlen_k and k_len == 0),
                 )
+                iket.range_push("epilogue")
                 epi_read_state, _ = self.epilogue(
                     epilogue_params,
                     epi_smem_tensors,
                 )
                 # acc_pipeline.consumer_release was already called in self.epi_load_acc_subtile
                 acc_consumer_state.advance()
+                iket.range_pop()
                 # Advance to next tile
                 tile_scheduler.advance_to_next_work()
+                work_tile = next_work_tile
             # Wait for D store complete
             if is_tma_warp:
             tmem_alloc_barrier.arrive_and_wait()
             tmem.free(acc_tmem_ptr)
     @cute.jit
     def _make_gather_A_copy(
         self,
         cta_rank_in_cluster: Int32,
         tCtSFA: Optional[cute.Tensor] = None,
         tCtSFB: Optional[cute.Tensor] = None,
+        copy_s2t_sfa: Optional[Callable] = None,
+        copy_s2t_sfb: Optional[Callable] = None,
+        sf_valid_insts_last_tile: Optional[Int32] = None,
     ) -> Tuple[cutlass.pipeline.PipelineState, cutlass.pipeline.PipelineState, cute.TiledMma]:
+        blockscaled = const_expr(copy_s2t_sfa is not None)
         if const_expr(blockscaled):
             assert all(x is not None for x in (tCtSFA, tCtSFB))
+            assert copy_s2t_sfb is not None
+        skip_sf_pad_insts = const_expr(sf_valid_insts_last_tile is not None)
         # If gather_A and use_2cta_instrs, the cp.async for the non-leader CTA will
         # arrive at an mbarrier on the non-leader CTA side, then the mma warp of the non-leader
         # CTA will wait for that then arrive at the mbarrier on the leader CTA.
                 if not is_leader_cta:
                     ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
                     with cute.arch.elect_one():
+                        # The odd CTA signals the even CTA. The arrive must release this
+                        # CTA's cp.async smem writes at cluster scope so that the leader's
+                        # 2-CTA MMA, which reads our smem over DSMEM, is guaranteed to
+                        # observe them; a plain mbarrier.arrive is only release.cta
+                        # (https://github.com/Dao-AILab/quack/issues/63).
+                        mbarrier_arrive_release_cluster(
+                            ab_pipeline.sync_object_full.get_barrier(ab_consumer_state.index),
+                            cta_rank_in_cluster & 0xFE,
                         )
             if is_leader_cta:
                 # Conditionally wait for AB buffer full
                 ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
+                if const_expr(need_nonleader_cta):
+                    # consumer_wait acquires at cta scope only; pair the non-leader's
+                    # cluster-scope release with a cluster-scope acquire of the (already
+                    # completed) phase before the MMA reads the peer CTA's smem.
+                    mbarrier_acquire_cluster(
+                        ab_pipeline.sync_object_full.get_barrier(ab_consumer_state.index),
+                        ab_consumer_state.phase,
+                    )
                 #  Copy SFA/SFB from smem to tmem
                 if const_expr(blockscaled):
+                    copy_s2t_sfa(ab_consumer_state.index)
+                    copy_s2t_sfb(ab_consumer_state.index)
+                # Ragged K: the last k-tile's SF atom holds pad bytes beyond the
+                # valid scale blocks. We exploit the fact that for mxfp8 the MMA
+                # instruction K size equals the SF vec size (both 32), i.e. each
+                # instruction consumes exactly one SF block: skipping the
+                # instructions for pad blocks — whose A/B values are
+                # TMA-zero-filled and contribute nothing — means the pad scales
+                # are never consumed and the gmem pad may be arbitrary (e8m0 0xFF
+                # = NaN would otherwise poison the accumulator via 0-value x
+                # NaN-scale products). Instruction issue is a leader-only
+                # decision, so this covers 2-CTA MMA too. fp4 has inst_k 64 (2
+                # SF blocks for mxfp4, 4 for nvfp4), but we don't do varlen_k
+                # for those formats.
+                # (The set/gemm sequence is duplicated below because the DSL
+                # rejects closures capturing staged values inside a dynamic if.)
+                if const_expr(skip_sf_pad_insts):
+                    num_mma_insts = Int32(num_k_blocks)
+                    if sf_valid_insts_last_tile > 0 and k_tile == k_tile_cnt - 1:
+                        num_mma_insts = sf_valid_insts_last_tile
                 for k_blk_idx in cutlass.range(num_k_blocks, unroll_full=True):
                     k_blk_coord = (None, None, k_blk_idx, ab_consumer_state.index)
                     if const_expr(blockscaled):
                         # Set SFA/SFB tensor to tiled_mma
                         sf_kblock_coord = (None, None, k_blk_idx)
+                        if const_expr(skip_sf_pad_insts):
+                            if k_blk_idx < num_mma_insts:
+                                tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
+                                tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock_coord].iterator)
+                                cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc)
+                                tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                        else:
+                            tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kblock_coord].iterator)
+                            tiled_mma.set(tcgen05.Field.SFB, tCtSFB[sf_kblock_coord].iterator)
+                            cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc)
+                            tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
+                    else:
+                        cute.gemm(tiled_mma, acc, tCrA[k_blk_coord], tCrB[k_blk_coord], acc)
+                        tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
                 # Async arrive AB buffer empty
                 ab_pipeline.consumer_release(ab_consumer_state)
             ab_consumer_state.advance()
         tTR_tAcc: cute.Tensor,
         tTR_rAcc: cute.Tensor,
         tRS_rD: cute.Tensor,
+        epi_coord: [int, int],
         acc_pipeline: pipeline.PipelineAsync,
         acc_consumer_state: pipeline.PipelineState,
         acc_release_idx: int,
     ):
         if not clear_acc:
             # Load accumulator from tensor memory buffer to register
+            cute.copy(tiled_copy_t2r, tTR_tAcc[None, None, None, epi_coord], tTR_rAcc)
             tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
             tRS_rD.store(tRS_rAcc.load())
         else:
             tRS_rD.fill(0.0)
+        assert epi_coord[0] == 0  # For Sm100, we assume epi_M = 1
+        if epi_coord[1] == acc_release_idx:
             cute.arch.fence_view_async_tmem_load()
+            acc_pipeline.consumer_release(acc_consumer_state)
     def epilog_tmem_copy_and_partition(
         self,
         self,
         tiled_mma: cute.TiledMma,
         cluster_layout_vmnk: cute.Layout,
         is_leader_cta: Boolean,
     ) -> pipeline.PipelineAsync:
         # If gather_A and use_2cta_instrs, the cp.async for the non-leader CTA will
         if const_expr(not self.gather_A or self.use_tma_gather):
             producer_cnt = 1
         else:
+            producer_cnt = self.num_ab_load_warps * cute.arch.WARP_SIZE
+            if const_expr(not self.use_2cta_instrs):
+                producer_cnt += 1
+            else:
+                producer_cnt += Int32(2) if is_leader_cta else Int32(0)
         ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt)
         # Each warp will contribute to the arrive count with the number of mcast size
         mcast_size = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
         ab_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, consumer_arrive_cnt
         )
+        if const_expr(not self.gather_A or self.use_tma_gather):
             pipeline_ab = PipelineTmaUmma.create(
                 num_stages=self.ab_stage,
                 producer_group=ab_pipeline_producer_group,
                 consumer_group=ab_pipeline_consumer_group,
             )
         else:
             pipeline_ab = PipelineTmaCpAsyncUmma.create(
                 num_stages=self.ab_stage,
                 producer_group=ab_pipeline_producer_group,
                 consumer_group=ab_pipeline_consumer_group,
                 tx_count=self.num_tma_load_bytes,
                 cta_layout_vmnk=cluster_layout_vmnk,
                 defer_sync=True,
             )
         return pipeline_ab
+    def make_acc_pipeline(self, cluster_layout_vmnk: cute.Layout) -> pipeline.PipelineAsync:
         acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
         num_acc_consumer_threads = self.num_epi_warps * (2 if self.use_2cta_instrs else 1)
         acc_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, num_acc_consumer_threads
         )
+        return PipelineUmmaAsync.create(
             num_stages=self.num_acc_stage,
             producer_group=acc_pipeline_producer_group,
             consumer_group=acc_pipeline_consumer_group,
             cta_layout_vmnk=cluster_layout_vmnk,
             defer_sync=True,
+            elect_one_release=True,
+            # TMEM load consumers are already ordered by fence_view_async_tmem_load()
+            syncwarp_before_release=False,
         )
     def make_sched_pipeline(
         self,
         cluster_layout_mnk: cute.Layout,
         has_C: bool = False,
     ) -> pipeline.PipelineAsync:
         # Threads/warps participating in this pipeline
         sched_pipeline_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, consumer_arrive_cnt
         )
+        # Plain PipelineAsync on purpose (vs the DSL example's PipelineClcFetchAsync):
+        # expect_tx is per-phase mbarrier state, so each mode's producer arms the full
+        # barrier as a transaction barrier itself — CLC's multicast try_cancel or
+        # STATIC/DYNAMIC's STAS st.async, both arrive_and_expect_tx(16) per CTA — and
+        # only the consumer protocol (wait full, elect-one arrive at CTA 0's empty
+        # barrier) is shared across modes. A CLC-specific pipeline would hardwire the
+        # producer and still need this one for STATIC/DYNAMIC.
+        return QuackPipelineAsync.create(
             num_stages=self.sched_stage,
             producer_group=sched_pipeline_producer_group,
             consumer_group=sched_pipeline_consumer_group,
             # If there's cluster, the consumers must arrive at the mbar of CTA 0 in the cluster.
             consumer_mask=None if const_expr(cluster_size == 1) else 0,
             defer_sync=True,
+            # One arrive per consumer warp (consumer_arrive_cnt counts warps): syncwarp
+            # so every lane's slot read is complete, then one elected lane signals.
+            elect_one_release=True,
         )
     @cute.jit
+    def make_a_prefetch_pipeline(self) -> pipeline.PipelineAsync:
         producer_cnt = 32
         a_prefetch_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, producer_cnt)
         consumer_arrive_cnt = self.num_ab_load_warps
         a_prefetch_consumer_group = pipeline.CooperativeGroup(
             pipeline.Agent.Thread, consumer_arrive_cnt
         )
+        return PipelineCpAsync.create(
             num_stages=self.a_prefetch_stage,
             producer_group=a_prefetch_producer_group,
             consumer_group=a_prefetch_consumer_group,
             defer_sync=True,
+            elect_one_release=True,
+            syncwarp_before_release=True,
         )
     @classmethod
         prefetch_A_idx: Literal[None, "varlen_m", "varlen_k"],
         smem_capacity: int,
         occupancy: int,
+        warp_shape_mnk: Tuple[int, int, int] | None = None,
     ) -> Tuple[int, int, int]:
         """Computes the number of stages for A/B/C operands based on heuristics.
         # Default D stages
         epi_stage = 4 if cute.size(epi_tile[1]) <= 16 else 2
+        epi_smem_bytes = cls.epi_smem_bytes(
+            epilogue_args, cta_tile_shape_mnk, epi_tile, warp_shape_mnk
+        )
+        has_tile_load = epi_smem_bytes.c_stage > 0
+        epi_c_stage = (
+            0
+            if c_dtype is None and not has_tile_load
+            else (4 if cute.size(epi_tile[1]) <= 16 else 2)
+        )
         # Calculate smem layout and size for one stage of A, B, and C
         a_smem_layout_staged_one = sm100_utils.make_smem_layout_a(
         d_bytes_per_stage = (
             cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) if d_dtype is not None else 0
         )
+        epi_bytes_per_stage = d_bytes_per_stage + epi_smem_bytes.d_stage
+        epi_bytes = epi_smem_bytes.unstaged + epi_bytes_per_stage * epi_stage
         if const_expr(c_dtype is not None):
             c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
             epi_bytes += c_bytes_per_stage * epi_c_stage
+        if const_expr(has_tile_load):
+            epi_bytes += epi_smem_bytes.c_stage * epi_c_stage
         # Calculate A/B/SFA/SFB stages:
         # Start with total smem per CTA (capacity / occupancy)
         # Refine epilogue stages:
         # Calculate remaining smem after allocating for A/B stages and reserved bytes
         # Add remaining unused smem to epilogue
+        if epi_bytes_per_stage > 0:
+            epi_stage += (remaining_bytes - ab_bytes_per_stage * ab_stage) // epi_bytes_per_stage
         return num_acc_stage, ab_stage, epi_stage, epi_c_stage
     @staticmethod
     @staticmethod
     def is_valid_mma_tiler_and_cluster_shape(
+        mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mn: Tuple[int, int],
         blockscaled: bool,
     ) -> bool:
         """
         Check if the mma tiler and cluster shape are valid
+        :param mma_tiler_mnk: The (M, N) or (M, N, K) shape of the MMA instruction tiler
+        :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]]
         :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
         :type cluster_shape_mn: Tuple[int, int]
         is_valid = True
         # Skip invalid mma tile shape
         if not blockscaled:
+            if mma_tiler_mnk[0] not in [64, 128, 256]:
                 is_valid = False
         else:
+            if mma_tiler_mnk[0] not in [128, 256]:
                 is_valid = False
+        mma_inst_n = mma_tiler_mnk[1] if mma_tiler_mnk[1] <= 256 else mma_tiler_mnk[1] // 2
         if not blockscaled:
             if mma_inst_n not in range(32, 257, 32):
                 is_valid = False
         else:
             # Blockscaled currently supports tile_n in {64, 128, 192, 256}.
+            if mma_tiler_mnk[1] not in [64, 128, 192, 256]:
                 is_valid = False
+        if cluster_shape_mn[0] % (2 if mma_tiler_mnk[0] == 256 else 1) != 0:
             is_valid = False
         # Skip invalid cluster shape
         is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0
         sf_dtype: Type[cutlass.Numeric],
         sf_vec_size: int,
         d_dtype: Type[cutlass.Numeric],
+        mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mn: Tuple[int, int],
         m: int,
         n: int,
         if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"):
             can_implement = False
         if not GemmSm100.is_valid_mma_tiler_and_cluster_shape(
+            mma_tiler_mnk, cluster_shape_mn, blockscaled=True
         ):
             can_implement = False
         if not GemmSm100.is_valid_tensor_alignment(
             m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major
         ):
         ab_dtype: Type[cutlass.Numeric],
         acc_dtype: Type[cutlass.Numeric],
         d_dtype: Type[cutlass.Numeric],
+        mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]],
         cluster_shape_mn: Tuple[int, int],
         m: int,
         n: int,
         :type acc_dtype: Type[cutlass.Numeric]
         :param d_dtype: The data type of the output tensor
         :type d_dtype: Type[cutlass.Numeric]
+        :param mma_tiler_mnk: The (M, N) or (M, N, K) shape of the MMA instruction tiler
+        :type mma_tiler_mnk: Union[Tuple[int, int], Tuple[int, int, int]]
         :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
         :type cluster_shape_mn: Tuple[int, int]
         :param m: The number of rows in the A tensor
             can_implement = False
         # Skip invalid mma tile shape and cluster shape
         if not GemmSm100.is_valid_mma_tiler_and_cluster_shape(
+            mma_tiler_mnk, cluster_shape_mn, blockscaled=False
         ):
             can_implement = False
         # Skip illegal problem shape for load/store alignment