fix(gemm): restore capture-safe K6 small-M dispatch

b12x/moe/_shared/kernels/w4a16/kernel.py

@@ -4199,12 +4199,10 @@ def _load_b_registers_trellis256_pair_bits(
             for jj in cutlass.range_constexpr(4):
                 local_n16 = Int32(4) * w_n + Int32(jj)
                 tile_base = Int32(0)
-                if cutlass.const_expr(int(low_bits) == int(high_bits)):
-                    tile_base = kt_base_u32 + local_n16 * Int32(8 * low_bits)
-                    wa[jj], wb[jj] = self._load_trellis256_pair_tile_windows(
-                        b_region, tile_base, lane, low_bits
-                    )
-                elif logical_k16 < Int32(8):
+                if (
+                    cutlass.const_expr(int(low_bits) == int(high_bits))
+                    or logical_k16 < Int32(8)
+                ):
                     tile_base = kt_base_u32 + local_n16 * Int32(8 * low_bits)
                     wa[jj], wb[jj] = self._load_trellis256_pair_tile_windows(
                         b_region, tile_base, lane, low_bits
@@ -11241,6 +11239,28 @@ def _trellis_dense_buffer(
     return buffer
 
 
+def _use_k6_mcg_small(
+    *,
+    device: torch.device,
+    m: int,
+    trellis_bits: int,
+    trellis_codebook: str,
+    trellis_pair_kind,
+    compute_dtype: torch.dtype,
+    external_hadamard_128,
+) -> bool:
+    """Select the capture-safe K6/MCG kernel only on its compiled target."""
+    return (
+        tuple(torch.cuda.get_device_capability(device)) == (12, 0)
+        and m <= 128
+        and trellis_bits == 6
+        and trellis_codebook == "mcg"
+        and trellis_pair_kind is None
+        and compute_dtype == torch.float16
+        and external_hadamard_128 is None
+    )
+
+
 def _run_trellis256_dense_current_device(
     x: torch.Tensor,
     prepared_dense,
@@ -11324,6 +11344,70 @@ def _run_trellis256_dense_current_device(
         None if hadamard_128 is None else _resolve_exl3_hadamard_128(hadamard_128)
     )
 
+    # Keep the established K6/MCG decode path independent from the generic
+    # Trellis scheduler. It owns both H128 rotations, needs no GEMM scratch,
+    # and is safe to capture with only caller-owned output/rotation storage.
+    # Compact pair payloads and the newer SQG codebooks use the generic path.
+    use_k6_mcg_small = _use_k6_mcg_small(
+        device=x.device,
+        m=m,
+        trellis_bits=trellis_bits,
+        trellis_codebook=trellis_codebook,
+        trellis_pair_kind=trellis_pair_kind,
+        compute_dtype=compute_dtype,
+        external_hadamard_128=external_hadamard_128,
+    )
+    if use_k6_mcg_small:
+        if x.dtype == torch.float16:
+            x_f16 = x
+        else:
+            input_f16 = _trellis_dense_buffer(
+                "input_f16",
+                input_f16,
+                shape=(m, size_k),
+                dtype=torch.float16,
+                device=x.device,
+            )
+            input_f16.copy_(x)
+            x_f16 = input_f16
+        rotated_f16 = _trellis_dense_buffer(
+            "rotated_f16",
+            rotated_f16,
+            shape=(m, size_k),
+            dtype=torch.float16,
+            device=x.device,
+        )
+        if output.dtype == torch.float16:
+            small_output = output
+        else:
+            output_f16 = _trellis_dense_buffer(
+                "output_f16",
+                output_f16,
+                shape=(m, size_n),
+                dtype=torch.float16,
+                device=x.device,
+            )
+            small_output = output_f16
+        from b12x.gemm.trellis_linear._small_m import run_k6_mcg
+
+        trellis_i16 = prepared_dense.trellis.view(torch.int16).view(
+            size_k // 16,
+            size_n // 16,
+            96,
+        )
+        run_k6_mcg(
+            x_f16,
+            trellis_i16,
+            small_output,
+            prepared_dense.suh,
+            rotated_f16,
+            prepared_dense.svh,
+            prepared_dense.workspace,
+        )
+        if output.dtype != torch.float16:
+            output.copy_(small_output)
+        return output
+
     gemm_output = _trellis_dense_buffer(
         "gemm_output",
         gemm_output,

tests/gemm/test_trellis_linear.py

@@ -36,6 +36,7 @@
 from b12x.moe._shared.kernels.w4a16.kernel import (
     _run_trellis_dense_hadamard128,
     _trellis256_dense_launch_geometry,
+    _use_k6_mcg_small,
 )
 from b12x.moe._shared.kernels.w4a16.prepare import (
     prepare_qsrt_pair_moe_weights,
@@ -413,6 +414,39 @@ def test_k6_small_m_rejects_unsupported_arch_before_jit(monkeypatch) -> None:
         _small_m.run_k6_mcg(*(torch.empty(0) for _ in range(7)))
 
 
+@pytest.mark.parametrize(
+    ("capability", "expected"),
+    [
+        ((12, 0), True),
+        ((12, 1), False),
+        ((9, 0), False),
+    ],
+)
+def test_k6_small_m_dispatch_requires_compiled_target(
+    monkeypatch,
+    capability: tuple[int, int],
+    expected: bool,
+) -> None:
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_capability",
+        lambda _device: capability,
+    )
+
+    assert (
+        _use_k6_mcg_small(
+            device=torch.device("cuda"),
+            m=128,
+            trellis_bits=6,
+            trellis_codebook="mcg",
+            trellis_pair_kind=None,
+            compute_dtype=torch.float16,
+            external_hadamard_128=None,
+        )
+        is expected
+    )
+
+
 @pytest.mark.parametrize(
     ("size_m", "size_k", "size_n", "expected"),
     [