perf(path_c): gridded CUDA SSD chunked-scan replaces MR single-block serial mamba recurrence (6.56s -> ~7.2ms forward, sm_121)

The Path-C MR mamba sub-region's forward was a single-block T.Kernel(1, threads=1024)
O(S)=4096-step SERIAL SSD recurrence (external_state_recurrence) = 6.56 s/call (95%
of the train step @8l). This ports the already-Metal-validated chunked SSD grid
kernels (F0/F1/F2) to target=cuda/sm_121 and drives the MR mamba region with them.

WHAT CHANGED (all in cppmega_mlx; no tilelang/codegen changes needed):
- _msl_transform: add _as_cuda_target (CUDA sibling of _as_metal_target).
- mamba3_chunked_*_core: thread `target=` through all 7 build_*_metal builders
  (default Metal => existing Metal callers byte-identical; "cuda" => CUDA backend).
  Single _resolve_chunked_compile_target() maps intent -> tvm.target.Target.
- mamba3_chunked_scan_core: NEW chunk_scan_fwd_cuda_prim — CUDA twin of the Metal F2.
  IDENTICAL SSD math; only the backend-shaped codegen differs:
    * acc_o is a REGISTER FRAGMENT (CUDA T.gemm asserts fragment C accumulator;
      Metal staged C in shared for the 8x8 simdgroup route).
    * x_shared plain "shared" (drop Metal "shared.dyn"); no make_swizzled_layout.
    * disable_tma=True on every global<->shared copy + pass_configs
      {tl.disable_tma_lower, tl.disable_warp_specialized}: on sm_121 the GEMM/copy
      TMA tensormap descriptor is 32/16-byte aligned where TMA needs 64-byte
      ("Invalid TMA descriptor arguments"); non-TMA cp.async path is correct.
  compile_chunk_scan_fwd_metal selects the CUDA prim + pass_configs when target=cuda.
- path_c_fusion_schedules: _mamba3_chunked_grid_delegation_prim now passes
  target=cuda on a CUDA host and block_Dstate=dstate to F2 (the default 128 would
  over-slice a production dstate=64 state axis).

MEASURED on gb10 tvm.cuda(0) / sm_121 (scratch/probe_chunked_scan_cuda_gb10.py),
production local_gb10_quarter mamba shape S=4096 chunk=64 heads=112 head_dim=64
state_dim=64 groups=8, grid (112,4,64)=28672 threadgroups vs the serial forward's 1:
  F2 scan+combine : 0.984 ms/call  (parity vs serial fwd 4.75e-4 fp16, PASS<5e-4)
  F0 precompute   : 5.025 ms/call  (cb exact, dA_cumsum 9.8e-4)
  F1 inter-chunk  : 1.209 ms/call  (prev_states 1.31e-4)
  TOTAL forward   : ~7.2 ms/call   vs the serial scan's 6559.7 ms/call (~900x).
The scan recurrence itself (F2, what literally replaced the serial carry) is
0.984 ms vs 6.56 s (~6670x). RULE #1 numeric equivalence: same SSD math, fp16 parity
verified vs the explicit per-timestep serial reference at all probed shapes.

NO REGRESSION on Metal: 21 chunked Metal tests pass (F0/F1/F2 + chained + region-flip);
disable_tma annotations are inert on Metal (TMA does not exist there), CUDA prim only
selected for target=cuda.

REMAINING (not in this commit): the BACKWARD chunked cores (B0/B1/B2) are now
target-threaded too but their CUDA run + the full flag-ON train_step (fwd+bwd) e2e
drive is a follow-up — B0/B1/B2 use NO T.gemm (plain TileLang, like F0/F1) so they
are expected to port with the same target threading; to be validated separately.

Author: apstenku123

cppmega_mlx/nn/_tilelang/_msl_transform.py

@@ -1469,6 +1469,42 @@ def _as_metal_target_uncached(target: Any) -> Any:
             return spec
 
 
+def _as_cuda_target(target: Any = "cuda") -> Any:
+    """Coerce a CUDA target spec into a ``tvm.target.Target`` for TileLang.
+
+    The mamba3 chunked-SSD grid kernels (F0/F1/F2 + B0/B1/B2) were authored
+    against the Metal codegen path via :func:`_as_metal_target`. On a CUDA host
+    (gb10 / sm_121) the SAME ``@T.prim_func`` bodies lower to CUDA unchanged, but
+    the builders must select a CUDA ``tvm.target.Target`` instead of a Metal one.
+    This helper is the CUDA sibling of :func:`_as_metal_target`: it returns a
+    ``tvm.target.Target`` for the bare ``"cuda"`` kind (TileLang's
+    ``determine_target`` allowlist accepts the Target object) or coerces a
+    passed-through spec/Target.
+
+    The arch (sm_121) is inferred by TVM from the live device at compile time;
+    we do NOT pin ``-arch`` here so the same call works on any CUDA device the
+    runtime actually owns. Returns the input unchanged when TVM is unavailable so
+    a non-TVM host sees a clear downstream error (RULE #1: no silent degrade).
+    """
+    try:
+        from tilelang import tvm  # type: ignore
+    except Exception:
+        try:
+            import tvm  # type: ignore
+        except Exception:
+            return target
+
+    if not isinstance(target, str):
+        try:
+            return tvm.target.Target(target)
+        except Exception:
+            return target
+    try:
+        return tvm.target.Target(target)
+    except Exception:
+        return target
+
+
 # Wave 5 grok finding #1: optional self-validation of the full-signature
 # corpus at import time. Off by default to keep import fast; enable in CI or
 # during local development by setting ``CPPMEGA_VALIDATE_PARSE_TESTS=1``.
@@ -1496,6 +1532,7 @@ def _as_metal_target_uncached(target: Any) -> Any:
     "MSLDispatchUnsupported",
     "MetalKernel",
     "TileLangMSLLowering",
+    "_as_cuda_target",
     "_as_metal_target",
     "_assert_path_c_metal_fp8_intrinsics_registered",
     "can_run_metal",

cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py

@@ -399,25 +399,30 @@ def build_chunk_scan_combine_bwd_metal(
     nheads: int,
     headdim: int,
     dstate: int,
+    *,
+    target: Any = None,
     **kwargs: Any,
 ) -> Any:
-    """Compile the B2 scan+combine backward kernel to a Metal ``JITKernel``.
+    """Compile the B2 scan+combine backward kernel to a ``JITKernel``.
 
-    Outputs (dC, dx, dz, dchunk_states, dinp, dA_cumsum_y, dD) are the trailing 7
-    params; pass PRE-ZEROED contiguous buffers positionally. RULE #1: compile
-    failures propagate (no serial fallback).
+    ``target`` selects Metal (default) or CUDA (``"cuda"`` / sm_121). Outputs
+    (dC, dx, dz, dchunk_states, dinp, dA_cumsum_y, dD) are the trailing 7 params;
+    pass PRE-ZEROED contiguous buffers positionally. RULE #1: compile failures
+    propagate (no serial fallback).
     """
     import tilelang
 
-    from cppmega_mlx.nn._tilelang import _msl_transform
+    from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import (
+        _resolve_chunked_compile_target,
+    )
 
     prim = chunk_scan_combine_bwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
     return tilelang.compile(
         prim,
         out_idx=[11, 12, 13, 14, 15, 16, 17],
-        target=_msl_transform._as_metal_target("metal -thread_warp_size=32"),
+        target=_resolve_chunked_compile_target(target),
     )
 
 
@@ -571,25 +576,30 @@ def build_inter_chunk_recur_bwd_metal(
     nheads: int,
     headdim: int,
     dstate: int,
+    *,
+    target: Any = None,
     **kwargs: Any,
 ) -> Any:
-    """Compile the B1 reverse inter-chunk recurrence kernel to a Metal JITKernel.
+    """Compile the B1 reverse inter-chunk recurrence kernel to a ``JITKernel``.
 
-    Outputs (dstates, dh0, dA_cumsum_tail) are the trailing 3 params; pass
-    PRE-ZEROED contiguous fp32 buffers positionally. RULE #1: compile failures
-    propagate (no serial fallback).
+    ``target`` selects Metal (default) or CUDA (``"cuda"`` / sm_121). Outputs
+    (dstates, dh0, dA_cumsum_tail) are the trailing 3 params; pass PRE-ZEROED
+    contiguous fp32 buffers positionally. RULE #1: compile failures propagate
+    (no serial fallback).
     """
     import tilelang
 
-    from cppmega_mlx.nn._tilelang import _msl_transform
+    from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import (
+        _resolve_chunked_compile_target,
+    )
 
     prim = inter_chunk_recur_bwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
     return tilelang.compile(
         prim,
         out_idx=[4, 5, 6],
-        target=_msl_transform._as_metal_target("metal -thread_warp_size=32"),
+        target=_resolve_chunked_compile_target(target),
     )
 
 
@@ -822,23 +832,28 @@ def build_chunk_precompute_bwd_metal(
     nheads: int,
     headdim: int,
     dstate: int,
+    *,
+    target: Any = None,
     **kwargs: Any,
 ) -> Any:
-    """Compile the B0 precompute backward kernel to a Metal ``JITKernel``.
+    """Compile the B0 precompute backward kernel to a ``JITKernel``.
 
-    Outputs (dx, dB, dlog_decay, ddt) are the trailing 4 params; pass PRE-ZEROED
+    ``target`` selects Metal (default) or CUDA (``"cuda"`` / sm_121). Outputs
+    (dx, dB, dlog_decay, ddt) are the trailing 4 params; pass PRE-ZEROED
     contiguous fp32 buffers positionally (dx is accumulated into — B2 wrote the
     D-skip path first). RULE #1: compile failures propagate (no serial fallback).
     """
     import tilelang
 
-    from cppmega_mlx.nn._tilelang import _msl_transform
+    from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import (
+        _resolve_chunked_compile_target,
+    )
 
     prim = chunk_precompute_bwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
     return tilelang.compile(
         prim,
         out_idx=[9, 10, 11, 12],
-        target=_msl_transform._as_metal_target("metal -thread_warp_size=32"),
+        target=_resolve_chunked_compile_target(target),
     )

cppmega_mlx/nn/_tilelang/mamba3_chunked_precompute_core.py

@@ -221,25 +221,31 @@ def build_chunk_precompute_metal(
     nheads: int,
     headdim: int,
     dstate: int,
+    *,
+    target: Any = None,
     **kwargs: Any,
 ) -> Any:
-    """Compile the F0 precompute kernel to a Metal ``JITKernel``.
+    """Compile the F0 precompute kernel to a ``JITKernel`` (Metal or CUDA).
 
-    Outputs (cb, dA_cumsum, summary_states) are the 6th/7th/8th params; pass
-    PRE-ZEROED contiguous buffers positionally (no ``out_idx`` allocation), and
-    ``torch.mps.synchronize()`` after. RULE #1: compile failures propagate.
+    ``target`` selects the codegen backend (``None`` => Metal default, ``"cuda"``
+    => CUDA / sm_121). Outputs (cb, dA_cumsum, summary_states) are the 6th/7th/8th
+    params; pass PRE-ZEROED contiguous buffers positionally (no ``out_idx``
+    allocation), and synchronize the device after. RULE #1: compile failures
+    propagate.
     """
     import tilelang
 
-    from cppmega_mlx.nn._tilelang import _msl_transform
+    from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import (
+        _resolve_chunked_compile_target,
+    )
 
     prim = chunk_precompute_fwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
     return tilelang.compile(
         prim,
         out_idx=[5, 6, 7],
-        target=_msl_transform._as_metal_target("metal -thread_warp_size=32"),
+        target=_resolve_chunked_compile_target(target),
     )
 
 
@@ -363,23 +369,28 @@ def build_inter_chunk_recur_metal(
     nheads: int,
     headdim: int,
     dstate: int,
+    *,
+    target: Any = None,
     **kwargs: Any,
 ) -> Any:
-    """Compile the F1 inter-chunk recurrence kernel to a Metal ``JITKernel``.
+    """Compile the F1 inter-chunk recurrence kernel to a ``JITKernel``.
 
-    Outputs (prev_states, final_state) are the 4th/5th params; pass PRE-ZEROED
-    contiguous fp32 buffers positionally, ``torch.mps.synchronize()`` after.
-    RULE #1: compile failures propagate (no serial fallback).
+    ``target`` selects Metal (default) or CUDA (``"cuda"`` / sm_121). Outputs
+    (prev_states, final_state) are the 4th/5th params; pass PRE-ZEROED contiguous
+    fp32 buffers positionally, synchronize the device after. RULE #1: compile
+    failures propagate (no serial fallback).
     """
     import tilelang
 
-    from cppmega_mlx.nn._tilelang import _msl_transform
+    from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import (
+        _resolve_chunked_compile_target,
+    )
 
     prim = inter_chunk_recur_fwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
     return tilelang.compile(
         prim,
         out_idx=[3, 4],
-        target=_msl_transform._as_metal_target("metal -thread_warp_size=32"),
+        target=_resolve_chunked_compile_target(target),
     )