feat(mamba3): fused smem-resident SSD forward (F0+F1+F2) + fix N=64 F2 prev_states dtype segfault

NEW: cppmega_mlx/nn/_tilelang/mamba3_ssd_fused_fwd.py — ONE smem-resident CUDA
kernel (gb10/sm_121) that fuses F0 precompute + intra-chunk scan + F1 inter-chunk
carry per (batch,chunk,head). Persistent per-(batch,head) threadgroup loops the
chunk axis serially carrying state[headdim,dstate] fp32 RESIDENT in shared memory
— cb/dA_cumsum/summary_states/prev_states never round-trip to global (the MEASURED
16.37 ms F0 lever). Env-gated CPPMEGA_MAMBA3_SSD_FUSED_FWD (default OFF); the
un-fused F0/F1/F2 + Metal path stay byte-identical.

smem budget @ prod N=64,P=64,L=64 = 82,688 B (80.75 KiB) < GB10 99 KB (101,376 B)
cap, 18,688 B headroom (fp16 operand tiles + fp32 accumulation; all-fp32 staging
would be 112.75 KiB and OVERFLOW). The builder RAISES (RULE #1) if the budget
exceeds the cap — never truncates dstate columns, never falls back to the slow
un-fused chain. Algebra verified fp64-exact (5.6e-17) vs the eager F0+F1+F2
reference for BOTH Output and final_state.

FIX (Phase 0, the prod N=64 'segfault'): chunk_scan_fwd_cuda_prim declared
prev_states fp16, but F1 writes it fp32 (accum_dtype, design §3.3) and the bwd B2
prim already declares it fp32 — the 2B-typed read of a 4B buffer walked the dstate
stride OOB -> launch crash at N=64. Declare prev_states accum_dtype (the producer's
dtype); the T.copy downcasts to the fp16 GEMM operand. Metal prim untouched.

No GPU on Mac: validated py_compile + ast + cross-symbol + fp64 algebra parity.
GB10 phase builds/measures.

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_chunked_scan_core.py

@@ -409,7 +409,18 @@ def main(
         dt: T.Tensor((batch, nheads, nchunks, chunk_size), dtype),  # type: ignore
         dA_cumsum: T.Tensor((batch, nheads, nchunks, chunk_size), dtype),  # type: ignore
         C: T.Tensor((batch, seqlen, ngroups, dstate), dtype),  # type: ignore
-        prev_states: T.Tensor((batch, nchunks, nheads, headdim, dstate), dtype),  # type: ignore
+        # DTYPE CONTRACT FIX (the prod N=64 "segfault"): F1
+        # (``inter_chunk_recur_fwd_metal_prim``) WRITES ``prev_states`` as
+        # ``accum_dtype`` (fp32, design §3.3/§6.6) and the bench/caller allocates it
+        # fp32; the BACKWARD B2 prim already declares it fp32 (backward_core line
+        # 191). This FORWARD prim previously declared it ``dtype`` (fp16) — a
+        # 4-byte/elem buffer read by a 2-byte-typed kernel walks the dstate stride
+        # out of bounds -> OOB global read -> the launch "segfault" at N=64. Reading
+        # it as fp32 (the producer's dtype) is the ONE-path fix; ``prev_state_shared``
+        # below stays fp16 because the off-diagonal ``T.gemm`` contracts in fp16
+        # (the value is downcast on the smem copy, matching the Metal twin's fp16
+        # GEMM operand). RULE #1: surface+fix the contract break, do not paper over.
+        prev_states: T.Tensor((batch, nchunks, nheads, headdim, dstate), accum_dtype),  # type: ignore
         D: T.Tensor((nheads), dtype),  # type: ignore
         Output: T.Tensor((batch, seqlen, nheads, headdim), dtype),  # type: ignore
     ):
@@ -432,6 +443,12 @@ def main(
             dA_cs_m_shared = T.alloc_shared((block_M), dtype)
             scale_m_local = T.alloc_fragment((block_M), accum_dtype)
             C_shared = T.alloc_shared((block_M, block_Dstate), dtype)
+            # prev_states is fp32 in global (accum_dtype, F1's output). The
+            # ``T.copy`` below downcasts the 4B/elem global read into this fp16 smem
+            # GEMM operand (matching the Metal twin's fp16 off-diagonal contraction).
+            # The FIX is the fp32 prim-parameter dtype above: with the buffer typed
+            # fp32 the copy walks the correct dstate stride (4B) instead of the
+            # previous fp16 stride (2B), which read out of bounds at N=64.
             prev_state_shared = T.alloc_shared((block_N, block_Dstate), dtype)
             D_local = T.alloc_fragment((1), accum_dtype)
             x_residual_shared = T.alloc_shared((block_M, block_N), dtype)