perf(path_c bwd): precompute shared LMAT[L,L] — kill ~65K exp2 recomputes/threadgroup in dinp

MEASURED run #1 (gb10 sm_121, prod): B2 dropped 2484 -> 334.6 ms/call (7.4x),
chain 447.8 ms < 456.1 numpy, ALL 8 grads pass incl dD=2.48e-5 (was 1.40e-3).
GO secured. But B2's dinp block (the verbatim Y_diag transpose) recomputes
lmat=exp2((dacs[l]-dacs[s])*p) in its innermost (sp/32-wave x nn/64 x l-reduction)
nest ~= 65K exp2/threadgroup — now the dominant remaining serial term, plus
dC_diag and dseg each recompute the same exp2.

Optimization (CUDA prim ONLY; Metal prim byte-identical): build the lower-tri
decay LMAT[l,s]=exp2((dacs[l]-dacs[s])*p) (s<=l, else 0) ONCE over L*L threads
(4096 exp2) into a shared tile, then read it from shared in dC_diag (l.630),
dinp (l.666) and dseg (l.696). Drops exp2 count ~16x in dinp alone. +16KB shared
(LMAT) -> ~64.5KB total, within the sm_121 per-block budget. Math IDENTICAL
(LMAT carries the exact same value; all 3 consumers index only the lower-tri
region where it is defined). RULE #1: still the ONE CUDA path.

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py

@@ -511,13 +511,29 @@ def main(
             dAcs_acc = T.alloc_shared((L,), accum_dtype)     # accumulated dA_cumsum grad
             XT = T.alloc_shared((L, headdim), accum_dtype)   # staged x[base+s, head, p]
             DYX = T.alloc_shared((L, L), accum_dtype)        # dyx[l,s] = sum_p dY[l,p]*x[s,p]
+            LMAT = T.alloc_shared((L, L), accum_dtype)       # exp2((dacs[l]-dacs[s])*p), l>=s
 
             # --- load dA_cumsum row, init dA grad accumulator ---
             for l in T.Parallel(L):
                 dacs[l] = T.Cast(accum_dtype, dA_cumsum[batch_idx, head_idx, chunk_idx, l])
                 dAcs_acc[l] = T.Cast(accum_dtype, 0)
             T.sync_threads()
 
+            # --- precompute the lower-tri intra-chunk decay LMAT[l,s] ONCE ---
+            #   lmat[l,s] = exp2((dacs[l]-dacs[s])*p) for s<=l (else 0). The Metal prim
+            #   recomputed this exp2 inside dC_diag (per (l,n,s)), dinp (per (s,p,n,l))
+            #   AND dseg (per (l,s)) — ~65K exp2/threadgroup in dinp alone. Built once
+            #   over L*L threads (4096 exp2), then read from shared everywhere below.
+            for ls in T.Parallel(L * L):
+                ll = ls // L
+                ss = ls % L
+                LMAT[ll, ss] = T.if_then_else(
+                    ss <= ll,
+                    T.exp2((dacs[ll] - dacs[ss]) * p),
+                    T.Cast(accum_dtype, 0),
+                )
+            T.sync_threads()
+
             # --- output/gate + D-skip transpose: dY[l,p], dz, dx, dD ---
             #   (VERBATIM from the Metal prim — already threaded over L*headdim;
             #    dD is fp32-correct, the 1.40e-3 gate miss is an input-precision
@@ -611,7 +627,7 @@ def main(
                 cdiag = T.alloc_local((1,), accum_dtype)
                 cdiag[0] = T.Cast(accum_dtype, 0)
                 for ss in T.serial(0, ll + 1):
-                    lmat = T.exp2((dacs[ll] - dacs[ss]) * p)
+                    lmat = LMAT[ll, ss]
                     dt_s = T.Cast(accum_dtype, dt[batch_idx, head_idx, chunk_idx, ss])
                     b_v = T.Cast(accum_dtype, B[batch_idx, base + ss, group_idx, nn])
                     cdiag[0] = cdiag[0] + lmat * dt_s * DYX[ll, ss] * b_v
@@ -647,7 +663,7 @@ def main(
                         acc = T.alloc_local((1,), accum_dtype)
                         acc[0] = T.Cast(accum_dtype, 0)
                         for ll in T.serial(ss, L):  # l >= s (lower-tri)
-                            lmat = T.exp2((dacs[ll] - dacs[ss]) * p)
+                            lmat = LMAT[ll, ss]  # shared, precomputed (no exp2 here)
                             c_v = T.Cast(
                                 accum_dtype, C[batch_idx, base + ll, group_idx, nn]
                             )
@@ -677,7 +693,7 @@ def main(
                 ll = ls // L
                 ss = ls % L
                 cb_v = T.Cast(accum_dtype, cb[batch_idx, chunk_idx, group_idx, ll, ss])
-                lmat = T.exp2((dacs[ll] - dacs[ss]) * p)
+                lmat = LMAT[ll, ss]  # shared, precomputed (no exp2 here)
                 dt_s = T.Cast(accum_dtype, dt[batch_idx, head_idx, chunk_idx, ss])
                 tri = T.if_then_else(
                     ss < ll, T.Cast(accum_dtype, 1), T.Cast(accum_dtype, 0)