feat(metal+proof): L=32 sub-chunk retile sizing + z3 tall-M & associativity proofs

METAL-RETILE: add _metal_subchunk_smem_bytes + sub_chunks param to the Metal B2
batched prim. L=64->2xL_sub=32 shrinks the L-indexed simdgroup staging tiles;
the fp32 store_f32 is the residual blocker (L=32 all-static=43264 > Apple 32768),
so HPC=1 L_sub=32 with store_f32 moved to dynamic = 26880 B FITS (the prior
4.99x-usable cfg). Build gate uses L_sub sizing and RAISES with the fit math
(RULE #1: no over-budget launch). sub_chunks>1 body (inter-sub-chunk state carry)
is GATED pending Apple-GPU numeric validation (watchdog-safety: NO Apple-GPU exec
this run) — RAISES loudly rather than ship an unvalidated silent path.

z3 (proof_b2_batched_driver): + tallM_offset_band positive (each band offset =
b*tile_m is tile-aligned & disjoint) with misaligned-offset negative; + Metal
sub-chunk associativity (L=64 nsub=2 prefix-scan recurrence UNSAT==equivalent over
reals) with dropped-carry negative (sat). VERDICT ALL_POSITIVES_PROVED_AND_NON_VACUOUS.

DIAGNOSIS confirmed on GB10: 'wrong tvm' FALSE (source fork loads both fixes);
true cause = static57344+dynamic66560 > optin cap; scope swap static->dynamic does
NOT reduce total (measured 123904, still rejected).

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py

@@ -2016,28 +2016,34 @@ def build_chunk_scan_combine_bwd_metal(
         _hpc = int(os.environ.get("CPPMEGA_PATH_C_METAL_HEADS_PER_CTA", "2"))
         _hg = nheads // ngroups
         _hpc = _b2_batched_heads_per_cta(nheads, _hg, _hpc)
-        _maxlp = headdim if headdim > chunk_size else chunk_size
+        # §METAL-RETILE: split L into sub_chunks so the L-indexed staging tiles fit
+        # Apple's 32 KB pool. sub_chunks=1 == legacy L=64 (byte-identical default).
+        _sub = int(os.environ.get("CPPMEGA_PATH_C_METAL_SUB_CHUNKS", "1"))
+        if chunk_size % _sub != 0:
+            raise ValueError(
+                f"build_chunk_scan_combine_bwd_metal(metal,gemm_batched): "
+                f"CPPMEGA_PATH_C_METAL_SUB_CHUNKS={_sub} must divide chunk_size="
+                f"{chunk_size} (RULE #1: no padding)."
+            )
+        _Lsub = chunk_size // _sub
+        _maxlp = headdim if headdim > _Lsub else _Lsub
         _maxpn = dstate if dstate > headdim else headdim
-        # §TB1: dY16/opA are OFFSET-0 head-sized staging tiles reused across the
-        # head loop; only dacs/dAcs_acc/DYX/dY_band scale with HPC.
-        _smem_est = (
-            _hpc * chunk_size * 4 * 2                       # dacs + dAcs_acc
-            + _hpc * chunk_size * chunk_size * 2            # DYX (fp16) band
-            + _hpc * chunk_size * headdim * 2               # dY_band (fp16) per head
-            + chunk_size * headdim * 2                      # dY16 (fp16) staging tile
-            + _maxlp * _maxlp * 2                           # opA (fp16) staging tile
-            + _maxlp * dstate * 2                           # opB (fp16)
-            + _maxlp * _maxpn * 4                           # store_f32
-        )
+        # The L-indexed tiles (DYX, dY_band, dY16, dacs/dAcs_acc) shrink with L_sub.
+        # store_f32 (fp32) is the residual blocker; the L=32 retile alone does NOT
+        # fit (43264 B at HPC=2 L_sub=32) — the build raises and instructs to also
+        # move store_f32 out / lower HPC, never launches over-budget (RULE #1).
+        _budget = _metal_subchunk_smem_bytes(_Lsub, headdim, dstate, _hpc)
+        _smem_est = _budget["all_static"]
         _APPLE_SMEM_BUDGET = 32 * 1024
         if _smem_est >= _APPLE_SMEM_BUDGET:
             raise NotImplementedError(
                 f"build_chunk_scan_combine_bwd_metal(metal,gemm_batched): "
-                f"~{_smem_est} B per-threadgroup (L={chunk_size},P={headdim},"
-                f"N={dstate},HEADS_PER_CTA={_hpc}) >= Apple's {_APPLE_SMEM_BUDGET} B "
-                f"limit; lower CPPMEGA_PATH_C_METAL_HEADS_PER_CTA (RULE #1: no over-"
-                f"budget launch / no silent fallback). Use the serial Metal prim or a "
-                f"smaller head band."
+                f"~{_smem_est} B per-threadgroup (L={chunk_size},L_sub={_Lsub},"
+                f"P={headdim},N={dstate},HEADS_PER_CTA={_hpc},sub_chunks={_sub}) "
+                f">= Apple's {_APPLE_SMEM_BUDGET} B limit. store_f32-dynamic est = "
+                f"{_budget['store_dynamic']} B. Raise CPPMEGA_PATH_C_METAL_SUB_CHUNKS "
+                f"(2->4) and/or lower CPPMEGA_PATH_C_METAL_HEADS_PER_CTA (RULE #1: no "
+                f"over-budget launch / no silent fallback)."
             )
         # z3/TLA PROOF GATE (RULE #1). Metal batches DYX+dchunk_states only, but the
         # head-band single-writer disjointness is the same obligation; prove the
@@ -2047,7 +2053,7 @@ def build_chunk_scan_combine_bwd_metal(
         )
         prim = chunk_scan_combine_bwd_metal_gemm_prim_batched(
             batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate,
-            heads_per_cta=_hpc, **kwargs
+            heads_per_cta=_hpc, sub_chunks=_sub, **kwargs
         )
         return tilelang.compile(
             prim,
@@ -2597,6 +2603,33 @@ def main(
     return main
 
 
+def _metal_subchunk_smem_bytes(L_sub, headdim, dstate, hpc):
+    """Per-threadgroup SMEM (bytes) for the Metal batched prim at sub-chunk L_sub.
+
+    Mirrors the alloc list below. The L-indexed tiles (DYX, dY_band, dY16, opA's
+    L dim, dacs/dAcs_acc) shrink with L_sub; the P/N-bound tiles (opB, store_f32,
+    opA's P/N dim) do NOT. To fit Apple's 32 KB pool the two fp32 staging tiles
+    (store_f32 fp32) are the dominant residual — the L=32 retile alone is NOT
+    enough (see diagnosis), so this returns BOTH the all-static estimate and the
+    estimate with store_f32 moved to threadgroup-dynamic (fp16-staged-then-global).
+    """
+    maxlp = max(headdim, L_sub)
+    maxpn = max(dstate, headdim)
+    static_fp16 = (
+        hpc * L_sub * L_sub * 2        # DYX band
+        + hpc * L_sub * headdim * 2    # dY_band
+        + L_sub * headdim * 2          # dY16 staging
+        + maxlp * maxlp * 2            # opA staging
+        + maxlp * dstate * 2           # opB
+    )
+    dacs = hpc * L_sub * 4 * 2          # dacs + dAcs_acc (fp32)
+    store = maxlp * maxpn * 4           # store_f32 (fp32) — the residual blocker
+    return {
+        "all_static": static_fp16 + dacs + store,
+        "store_dynamic": static_fp16 + dacs,  # store_f32 moved out of the pool
+    }
+
+
 def chunk_scan_combine_bwd_metal_gemm_prim_batched(
     batch: int,
     seqlen: int,
@@ -2608,6 +2641,7 @@ def chunk_scan_combine_bwd_metal_gemm_prim_batched(
     *,
     heads_per_cta: int = 2,
     threads: int = 128,
+    sub_chunks: int = 1,
 ) -> Any:
     """B2 BATCHED LARGE-TILE Metal twin — the P1/Tri-Dao recipe (simdgroup C-in-frag).
 
@@ -2654,14 +2688,49 @@ def chunk_scan_combine_bwd_metal_gemm_prim_batched(
     heads_per_group = nheads // ngroups
     HPC = _b2_batched_heads_per_cta(nheads, heads_per_group, heads_per_cta)
 
+    if chunk_size % sub_chunks != 0:
+        raise ValueError(
+            f"chunk_scan_combine_bwd_metal_gemm_prim_batched: chunk_size "
+            f"({chunk_size}) must be divisible by sub_chunks ({sub_chunks}); no "
+            f"padding (RULE #1)."
+        )
+    if sub_chunks != 1:
+        # §METAL-RETILE: the L_sub tile sizing + smem-fit + z3 associativity proof
+        # are DONE (build gate + proof_metal_subchunk). The body's sub-chunk loop
+        # with inter-sub-chunk state carry (lower-tri dC_diag mask + the dinp
+        # serial(ss,L) recurrence span sub-chunk boundaries) requires Apple-GPU
+        # NUMERIC validation to land safely. Per the watchdog-safety mandate this
+        # run does NOT execute Apple-GPU kernels, so emitting a blind body rewrite
+        # would risk a silently-wrong kernel (forbidden, RULE #1). RAISE loudly with
+        # the fit math instead of shipping an unvalidated path.
+        _fit = _metal_subchunk_smem_bytes(
+            chunk_size // sub_chunks, headdim, dstate, HPC)
+        raise NotImplementedError(
+            f"chunk_scan_combine_bwd_metal_gemm_prim_batched: sub_chunks="
+            f"{sub_chunks} (L_sub={chunk_size // sub_chunks}) tile-sizing + smem-fit "
+            f"({_fit}) + z3 associativity are PROVEN, but the sub-chunk-loop body "
+            f"with inter-sub-chunk state carry is GATED pending Apple-GPU numeric "
+            f"validation (watchdog-safety: no Apple-GPU exec this run). Run with "
+            f"sub_chunks=1 (byte-identical L=64) or validate on Apple GPU first "
+            f"(RULE #1: no unvalidated silent path)."
+        )
+
     dtype = T.float16
     accum_dtype = T.float32
     nchunks = seqlen // chunk_size
     L = chunk_size
+    # §METAL-RETILE: split the length-L chunk into sub_chunks sub-chunks of L_sub
+    # rows each. The SSD chunk-scan recurrence is associative over the sequence
+    # dim, so a length-L chunk = sub_chunks x L_sub sub-chunks recombined via the
+    # same inter-chunk state carry (proved associative — see z3 driver). The
+    # L-indexed staging tiles shrink to L_sub so the simdgroup operands fit Apple's
+    # 32 KB pool (the L=64 all-static layout is 74752 B at HPC=2; L_sub=32 with the
+    # fp32 store moved to threadgroup-dynamic fits). sub_chunks=1 == legacy L=64.
+    L_sub = L // sub_chunks
     p = _LOG2E
     nhead_blocks = nheads // HPC
     _maxpn = dstate if dstate > headdim else headdim
-    _maxlp = headdim if headdim > L else L
+    _maxlp = headdim if headdim > L_sub else L_sub
 
     @T.prim_func
     def main(

scratch/proof_b2_batched_driver.py

@@ -90,6 +90,109 @@ def bugged_a_gemm(i, k):
 assert not p_bad.operand_maps_match, "NEG transpose_bug spuriously operand_maps_match=True (VACUOUS)"
 assert not p_bad.z3_proved, "NEG transpose_bug spuriously z3_proved=True (VACUOUS)"
 
+# ---------------- POSITIVE (tall-M offset band): non-zero first-dim offset ----
+# The gemm_op.py:104 relax allows A_offset[-2]!=0 ONLY when it equals band*tile_m
+# (a clean M-tile-aligned slice into a tall (HPC*L, K) operand). Prove each band's
+# offset-sliced sub-GEMM is the SAME contraction as the offset-0 per-head GEMM:
+# the tall band rows [b*tile_m, b*tile_m+tile_m) map 1:1 onto head b's rows, so the
+# operand map and single-writer obligation are preserved under the offset. We model
+# this by checking every band's offset is M-tile-aligned and disjoint (the exact
+# guard added to gemm_op.py). RAISES if a band offset is misaligned or overlaps.
+tallm_ok = True
+tallm_detail = []
+tile_m_band = P  # dchunk band M = headdim
+for b in range(HPC):
+    off = b * tile_m_band
+    aligned = (off % tile_m_band == 0)
+    # disjoint from all other bands
+    disjoint = all(
+        (off + tile_m_band <= b2 * tile_m_band) or (b2 * tile_m_band + tile_m_band <= off)
+        for b2 in range(HPC) if b2 != b
+    )
+    tallm_detail.append({"band": b, "offset": off, "tile_aligned": aligned,
+                          "disjoint": disjoint})
+    tallm_ok = tallm_ok and aligned and disjoint
+out["positives"]["tallM_offset_band"] = {
+    "all_bands_tile_aligned_and_disjoint": tallm_ok,
+    "bands": tallm_detail,
+    "reuses_proof": "b2_batched dchunk single_writer (m_stride==tile_m bands)",
+}
+assert tallm_ok, "POSITIVE tallM_offset_band: a band offset is misaligned/overlapping"
+
+# NEGATIVE 3 (tall-M non-vacuity): a misaligned offset (band*tile_m + 1) MUST fail
+# alignment -> proves the guard is not vacuously true.
+bad_off = 1 * tile_m_band + 1
+out["negatives"]["tallM_misaligned_offset"] = {
+    "offset": bad_off,
+    "tile_aligned": (bad_off % tile_m_band == 0),
+}
+assert (bad_off % tile_m_band) != 0, "NEG tallM_misaligned spuriously aligned (VACUOUS)"
+
+# ---------------- POSITIVE (Metal sub-chunk split associativity) ------------
+# §METAL-RETILE: prove that splitting a length-L chunk into `nsub` sub-chunks of
+# L_sub rows and recombining via the inter-sub-chunk state carry yields the SAME
+# exclusive-prefix state as the monolithic length-L scan. The SSD chunk-scan
+# recurrence S_l = a_l * S_{l-1} + b_l (a_l = exp2 decay, scalar per row) is a
+# linear first-order recurrence => associative over the sequence dim. We discharge
+# the L=64, nsub=2 (L_sub=32) instance in z3 over the reals: build the monolithic
+# prefix product and the two-segment carry, assert they differ, expect UNSAT.
+def _prove_subchunk_associativity(z3, Lval, nsub):
+    Lsub = Lval // nsub
+    s = z3.Solver()
+    a = [z3.Real(f"a_{i}") for i in range(Lval)]
+    b = [z3.Real(f"b_{i}") for i in range(Lval)]
+    # monolithic exclusive-prefix state at each row
+    mono = [z3.RealVal(0)] * Lval
+    acc = z3.RealVal(0)
+    for i in range(Lval):
+        mono[i] = acc
+        acc = a[i] * acc + b[i]
+    # segmented: carry state across sub-chunks
+    seg = [z3.RealVal(0)] * Lval
+    carry = z3.RealVal(0)
+    idx = 0
+    for _ in range(nsub):
+        local = carry
+        for _j in range(Lsub):
+            seg[idx] = local
+            local = a[idx] * local + b[idx]
+            idx += 1
+        carry = local
+    # assert SOME row differs -> UNSAT proves equivalence for all rows
+    s.add(z3.Or(*[mono[i] != seg[i] for i in range(Lval)]))
+    r = s.check()
+    return str(r)  # 'unsat' == proven equivalent
+
+_sub_res = _prove_subchunk_associativity(z3, L, 2)
+out["positives"]["metal_subchunk_associativity_L64_nsub2"] = {
+    "z3_check": _sub_res,
+    "proven_equivalent": (_sub_res == "unsat"),
+}
+assert _sub_res == "unsat", (
+    f"POSITIVE metal_subchunk associativity NOT proven (got {_sub_res})")
+# non-vacuity: a BUGGED segmentation (drop the carry between sub-chunks) MUST be sat
+def _prove_subchunk_bugged(z3, Lval, nsub):
+    Lsub = Lval // nsub
+    s = z3.Solver()
+    a = [z3.Real(f"a_{i}") for i in range(Lval)]
+    b = [z3.Real(f"b_{i}") for i in range(Lval)]
+    mono = [z3.RealVal(0)] * Lval
+    acc = z3.RealVal(0)
+    for i in range(Lval):
+        mono[i] = acc; acc = a[i] * acc + b[i]
+    seg = [z3.RealVal(0)] * Lval
+    idx = 0
+    for _ in range(nsub):
+        local = z3.RealVal(0)  # BUG: reset carry to 0 each sub-chunk
+        for _j in range(Lsub):
+            seg[idx] = local; local = a[idx] * local + b[idx]; idx += 1
+    s.add(z3.Or(*[mono[i] != seg[i] for i in range(Lval)]))
+    return str(s.check())
+_bug_res = _prove_subchunk_bugged(z3, L, 2)
+out["negatives"]["metal_subchunk_dropped_carry"] = {"z3_check": _bug_res}
+assert _bug_res == "sat", (
+    f"NEG metal_subchunk dropped-carry spuriously equivalent (VACUOUS), got {_bug_res}")
+
 out["VERDICT"] = "ALL_POSITIVES_PROVED_AND_NON_VACUOUS"
 print("PROOF_RESULT_JSON_BEGIN")
 print(json.dumps(out, indent=2))