feat(path_c B2): batched large-tile tensor-core GEMM twin (CUDA mma.sync + Metal simdgroup)

NEW prims (flag-gated, the ONE path when on, RAISE on failure — RULE #1; serial/§27
prims kept byte-identical):
  * chunk_scan_combine_bwd_cuda_prim_gemm_batched (sm_121, fragment-C):
    grid (batch*nchunks, nheads//HEADS_PER_CTA); each CTA owns HEADS_PER_CTA heads and
    issues the 4 GEMM-able contractions (DYX, dC_off, dC_diag masked, dchunk_states
    transpose_A) as a HEAD-LOOP over a SHARED tall-M operand band (M = HEADS_PER_CTA*64),
    amortizing ldmatrix/staging/sync over the band — the P1/Tri-Dao lever the §27
    single-64-tile twin (0.749x) lacked. dinp 3-index term + dz/dx/dD + dseg stay
    threaded (documented scope boundary). HEADS_PER_CTA default 4.
  * chunk_scan_combine_bwd_metal_gemm_prim_batched (simdgroup, C-in-frag): batches
    DYX+dchunk_states only (Apple 32KB pool); HEADS_PER_CTA default 2.

Flags: CPPMEGA_PATH_C_B2_GEMM_BATCHED (cuda) / CPPMEGA_PATH_C_METAL_GEMM_BATCHED (metal),
HEADS_PER_CTA via CPPMEGA_PATH_C_B2_HEADS_PER_CTA / CPPMEGA_PATH_C_METAL_HEADS_PER_CTA.
Mutually exclusive with the other B2 flags (RAISE). Per-backend smem budget gate RAISES
over-budget (gb10 99KB / Apple 32KB), no silent re-tile.

z3 PROOF WIRED (RULE #1, fail-closed): _gemm_rewrite_proof.py gains
b2_batched_dchunk/dcoff/dcdiag_contraction + b2_batched_tiling (m_blocks=HEADS_PER_CTA,
m_stride=tile_m per-head-band single-writer disjointness). build dispatcher calls
require_gemm_rewrite_proof before emitting the batched prim — RAISES GemmRewriteNotProven
on sat/unknown/disabled. Verified locally: 3 rewrites discharge z3_proved=True, single_writer
True; non-vacuous (overlap-band -> single_writer False, transpose bug -> maps False);
fail-closed when z3 disabled.

probe_chunked_backward_cuda_gb10.py: B2-BATCHED-AB arm (vs v1 + vs single-tile gemm,
HEADS_PER_CTA-parametrized) for the gb10 measure (bs1 isolates tiling; bs4 = full recipe).

8-grad parity + tile microbench is the gb10/Test step (local tilelang import is broken;
no heavy local python per 80GB guard). MEASURED win is EXTRAPOLATION until gb10 runs.

Author: apstenku123

cppmega_mlx/nn/_tilelang/_gemm_rewrite_proof.py

@@ -547,6 +547,103 @@ def b2_dinp_contraction(z3, *, chunk_size: int, headdim: int):
     )
 
 
+def b2_batched_dchunk_contraction(z3, *, chunk_size: int, headdim: int, dstate: int):
+    """B2 BATCHED ``dchunk_states[p,n] = sum_l (decay[l]*dY[l,p])*C[l,n]`` (transpose_A).
+
+    The batched-large-tile prim stacks ``HEADS_PER_CTA`` heads along M; each head is
+    a DISJOINT 64-row band of the tall opA tile and is GEMM'd against its own C
+    operand (block-diagonal). This descriptor proves the PER-HEAD contraction (the
+    operand maps + the transpose_A address arithmetic + the decay scale fold); the
+    per-band single-writer disjointness is proved by the GemmTiling
+    ``m_blocks=HEADS_PER_CTA, m_stride=tile_m`` passed alongside (each band owns a
+    contiguous 64-row slice, no overlap). Output rows = p (headdim), cols = n
+    (dstate), reduction k = l (chunk). The decay ``sd[l]=exp2(dacs[l]*p)`` is a
+    function of the reduction index l only, folded into the dY operand BEFORE the
+    GEMM — modeled as an uninterpreted ``scale(l)`` that the GEMM folds identically.
+    """
+
+    scale = z3.Function("b2_sd", z3.IntSort(), z3.RealSort())
+    return GemmContraction(
+        name="B2_batched_dchunk_states",
+        m_extent=headdim,
+        n_extent=dstate,
+        k_extent=chunk_size,
+        a_addr_serial=lambda i, k: k * headdim + i,  # dY[l, p] (transpose_A: k=l rows)
+        a_addr_gemm=lambda i, k: k * headdim + i,
+        b_addr_serial=lambda k, j: k * dstate + j,  # C[l, n]
+        b_addr_gemm=lambda k, j: k * dstate + j,
+        scale_serial=lambda i, k, j: scale(k),
+        scale_gemm=lambda i, k, j: scale(k),
+    )
+
+
+def b2_batched_dcoff_contraction(z3, *, chunk_size: int, headdim: int, dstate: int):
+    """B2 BATCHED ``dC_off[l,n] = sum_p dY[l,p]*prev_states[p,n]`` (dense, un-transposed).
+
+    Per-head sub-GEMM inside the head-band. Output rows = l (chunk), cols = n
+    (dstate), reduction k = p (headdim). No mask, no fold on the GEMM itself (the
+    per-l state_decay is applied post-GEMM in the dC store). Address maps encode the
+    un-transposed ``dY[l,p] @ prev_states[p,n]``.
+    """
+
+    return GemmContraction(
+        name="B2_batched_dC_off",
+        m_extent=chunk_size,
+        n_extent=dstate,
+        k_extent=headdim,
+        a_addr_serial=lambda i, k: i * headdim + k,  # dY[l, p]
+        a_addr_gemm=lambda i, k: i * headdim + k,
+        b_addr_serial=lambda k, j: k * dstate + j,  # prev_states[p, n]
+        b_addr_gemm=lambda k, j: k * dstate + j,
+    )
+
+
+def b2_batched_dcdiag_contraction(z3, *, chunk_size: int, dstate: int):
+    """B2 BATCHED ``dC_diag[l,n] = sum_{s<=l} M[l,s]*B[s,n]`` (lower-tri masked).
+
+    Per-head sub-GEMM. The causal mask ``keep(i,j,k)=(k<=i)`` (s<=l) is applied to
+    the A-operand fragment BEFORE the GEMM — the GEMM must fold the IDENTICAL
+    predicate. Output rows = l (chunk), cols = n (dstate), reduction k = s (chunk).
+    """
+
+    return GemmContraction(
+        name="B2_batched_dC_diag_lowertri",
+        m_extent=chunk_size,
+        n_extent=dstate,
+        k_extent=chunk_size,
+        a_addr_serial=lambda i, k: i * chunk_size + k,  # M[l, s]
+        a_addr_gemm=lambda i, k: i * chunk_size + k,
+        b_addr_serial=lambda k, j: k * dstate + j,  # B[s, n]
+        b_addr_gemm=lambda k, j: k * dstate + j,
+        mask_serial=lambda i, j, k: k <= i,
+        mask_gemm=lambda i, j, k: k <= i,
+    )
+
+
+def b2_batched_tiling(*, tile_m: int, tile_n: int, tile_k: int, heads_per_cta: int) -> "GemmTiling":
+    """The per-head-band tiling for the batched B2 GEMMs.
+
+    M is decomposed into ``heads_per_cta`` DISJOINT bands of ``tile_m`` rows each
+    (``m_blocks=heads_per_cta, m_stride=tile_m`` so band b owns rows
+    ``[b*tile_m, b*tile_m+tile_m)`` — contiguous, non-overlapping). The
+    single-writer obligation proves no two head-bands write the same output row
+    (block-diagonal disjointness, the load-bearing race-freedom property of the
+    head-accumulation tiling). N/K are single-tile (k_steps=1, n_blocks=1) at the
+    prod 64 dims.
+    """
+
+    return GemmTiling(
+        tile_m=tile_m,
+        tile_n=tile_n,
+        tile_k=tile_k,
+        m_blocks=heads_per_cta,
+        n_blocks=1,
+        k_steps=1,
+        m_stride=tile_m,
+        n_stride=tile_n,
+    )
+
+
 def f0_summary_contraction(z3, *, chunk_size: int, headdim: int, dstate: int):
     """F0 ``summary_states[p,n] = sum_l decay[l]*dt[l]*x[l,p]*B[l,n]``.