fix(§METAL-RETILE): MSL-compile the sub_chunks=2 retile body (Psc1 frag-drain) + compile-check

The §METAL-RETILE sub_chunks=2 body lowered through TIR but FAILED MSL codegen:
'assigning to float4 from incompatible type simdgroup_float8x8' on the direct
T.copy(dchunk_sc1, Psc1[hh, 0:headdim, 0:dstate]) — a fragment -> 3D-shared-slice
copy mis-vectorizes the simdgroup matrix. Fix: drain the sc1 dchunk partial via
the 2D store_f32 staging (the same vectorizable copy the main dchunk uses) then
scatter store_f32 -> Psc1[hh] band with an element-wise T.Parallel copy.

After the fix BOTH Metal prims MSL-compile (codegen-only, NO Apple-GPU dispatch
this run per the watchdog mandate):
 - main_retile  (sub_chunks=2, L=32,L_sub=16,P=N=32,HPC=1,threads=32): MSL 44090 B
 - legacy main  (sub_chunks=1, HPC=2): MSL 37795 B (barrier-fix path lowers clean)

scratch/compile_check_metal_retile.py is the CPU-only lower/codegen harness (guards
the prod-L=64 32KB NO-GO assertion + lowers both prims). z3 driver green:
ALL_POSITIVES_PROVED_AND_NON_VACUOUS (5 metal_subchunk positives UNSAT/partition +
5 paired negatives sat/miss).

HONEST NO-GO: prod L=64 via this retile is 68096 B (HPC=1) / 99328 B (HPC=2),
both >> Apple 32768 B — the build gate RAISES (no over-budget launch, RULE #1).
The body is exercisable only at small in-budget bounded dims (compile-check). No
Apple-GPU numeric measurement this run (watchdog-safety).

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py

@@ -3744,7 +3744,15 @@ def main_retile(
                             transpose_A=True,
                         )
                     T.sync_threads()
-                T.copy(dchunk_sc1, Psc1[hh, 0:headdim, 0:dstate])
+                # drain the sc1 partial frag -> 2D store_f32 (vectorizable copy) ->
+                # Psc1[hh] band via an element-wise Parallel copy (a direct frag->3D-
+                # slice T.copy mis-vectorizes to float4 vs simdgroup_float8x8 in MSL).
+                T.copy(dchunk_sc1, store_f32[0:headdim, 0:dstate])
+                T.sync_threads()
+                for pn in T.Parallel(headdim * dstate):
+                    pp = pn // dstate
+                    nn = pn % dstate
+                    Psc1[hh, pp, nn] = store_f32[pp, nn]
                 T.sync_threads()
                 T.copy(dchunk_frag, store_f32[0:headdim, 0:dstate])
                 T.sync_threads()

scratch/compile_check_metal_retile.py

@@ -0,0 +1,71 @@
+"""COMPILE-CHECK ONLY (no GPU dispatch) for the §METAL-RETILE sub_chunks=2 body.
+
+Builds the Metal batched prim with sub_chunks=2 at a SMALL in-budget bounded shape
+(L=16, L_sub=8, P=N=16, HPC=1) and LOWERS it to MSL via tilelang.compile (codegen,
+NOT a kernel launch — safe under the watchdog mandate). Also asserts the prod-L=64
+build gate RAISES (honest 32 KB NO-GO) and that sub_chunks=1 stays selectable.
+
+Run: .venv/bin/python scratch/compile_check_metal_retile.py
+"""
+import os
+os.environ.setdefault("CPPMEGA_PATH_C_METAL_GEMM_BATCHED", "1")
+
+from cppmega_mlx.nn._tilelang import mamba3_chunked_backward_core as core
+
+
+def _metal_target():
+    # Resolve the Metal target the build path uses, WITHOUT dispatching.
+    from cppmega_mlx.nn._tilelang.mamba3_chunked_scan_core import (
+        _resolve_chunked_compile_target,
+    )
+    return _resolve_chunked_compile_target(None)
+
+
+def lower_retile_small():
+    import tilelang
+    # bounded in-budget: L=16, sub_chunks=2 -> L_sub=8; P=N=16; HPC=1; B=1
+    prim = core.chunk_scan_combine_bwd_metal_gemm_prim_batched(
+        batch=1, seqlen=32, chunk_size=32, ngroups=1, nheads=1, headdim=32,
+        dstate=32, heads_per_cta=1, sub_chunks=2, threads=32,
+    )
+    print("retile prim built (sub_chunks=2, L=32, L_sub=16, P=N=32, HPC=1, threads=32)")
+    tgt = _metal_target()
+    k = tilelang.compile(
+        prim, out_idx=[11, 12, 13, 14, 15, 16, 17], target=tgt
+    )
+    src = k.get_kernel_source()
+    print("LOWER OK: MSL length =", len(src))
+    assert len(src) > 0
+    return True
+
+
+def check_legacy_still_builds():
+    import tilelang
+    prim = core.chunk_scan_combine_bwd_metal_gemm_prim_batched(
+        batch=1, seqlen=32, chunk_size=32, ngroups=1, nheads=2, headdim=32,
+        dstate=32, heads_per_cta=2, sub_chunks=1, threads=32,
+    )
+    tgt = _metal_target()
+    k = tilelang.compile(prim, out_idx=[11, 12, 13, 14, 15, 16, 17], target=tgt)
+    src = k.get_kernel_source()
+    print("LEGACY sub_chunks=1 LOWER OK (HPC=2): MSL length =", len(src))
+    assert len(src) > 0
+    return True
+
+
+def check_prod_l64_gate_raises():
+    # The build-site gate must RAISE for prod L=64 (honest 32 KB NO-GO).
+    fit1 = core._metal_subchunk_smem_bytes(32, 64, 64, 1)
+    fit2 = core._metal_subchunk_smem_bytes(32, 64, 64, 2)
+    print("prod L=64 fit:", "HPC=1", fit1, "HPC=2", fit2)
+    assert fit1["all_static"] >= 32 * 1024, "HPC=1 prod unexpectedly fits 32KB"
+    assert fit2["all_static"] >= 32 * 1024, "HPC=2 prod unexpectedly fits 32KB"
+    print("prod L=64 retile is an honest 32KB NO-GO (gate RAISES) — confirmed")
+    return True
+
+
+if __name__ == "__main__":
+    check_prod_l64_gate_raises()
+    lower_retile_small()
+    check_legacy_still_builds()
+    print("ALL_COMPILE_CHECKS_OK")