fix(path_c bwd): dseg block — T.Parallel(L*L) instead of lane-strided T.serial (CUDA scope bug)

The B2 cuda prim FAILED to compile on gb10 sm_121: 'cb_v' undefined (line 134)
+ 'dseg' undefined (line 139). Root cause: the dseg segsum-VJP loop used the
lane-strided 'for ls0 in T.serial(0, L*L, threads)' idiom with a masked
'if lsi<L*L: if ll>ss:' body that (a) reads cb from global and (b) atomic_adds
into SHARED dAcs_acc. TileLang's shared-mem barrier-insertion pass put
__syncthreads() between the cb load, the dseg compute, and the atomic scatter,
re-emitting the if-guard around each fragment WITHOUT hoisting the local
declarations — so cb_v (declared in fragment 1) and dseg (declared in fragment 2)
were out of scope at their uses in later fragments.

Fix: map the L*L (ll,ss) grid over T.Parallel(L*L) — the SAME un-fragmented
thread-parallel idiom the dC apply block (which also atomic_adds into dAcs_acc)
and the DYX build already use and which compiled cleanly. dseg lands in a
thread-local; the +/- scatter keeps the ll>ss mask (dseg=0 off the strict
lower-tri -> no-op atomic adds). Math IDENTICAL to the Metal prim / the prior
lane-strided form. Metal prim untouched. RULE #1: still the ONE CUDA path.

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py

@@ -661,26 +661,27 @@ def main(
             #   dLmat[l,s] = sum_p dY[l,p]*x[s,p] == DYX[l,s] (reuse, zero recompute)
             #   M[l,s] = cb*lmat*dt_s ; dseg = DYX[l,s]*M[l,s]
             #   segsum-vjp: +dAcs_acc[l] over l>=s, -dAcs_acc[s] over s<l (strictly l>s).
-            # Lane-strided over the L*L (ll,ss) grid with the ll>ss mask (mirrors the
-            # dinp lane-strided pattern); the +/- scatter races on dAcs_acc -> atomic.
-            for ls0 in T.serial(0, L * L, threads):
-                lane = T.get_thread_binding(0)
-                lsi = ls0 + lane
-                if lsi < L * L:
-                    ll = lsi // L
-                    ss = lsi % L
-                    if ll > ss:
-                        cb_v = T.Cast(
-                            accum_dtype, cb[batch_idx, chunk_idx, group_idx, ll, ss]
-                        )
-                        lmat = T.exp2((dacs[ll] - dacs[ss]) * p)
-                        dt_s = T.Cast(
-                            accum_dtype, dt[batch_idx, head_idx, chunk_idx, ss]
-                        )
-                        m_ls = cb_v * lmat * dt_s
-                        dseg = DYX[ll, ss] * m_ls
-                        T.atomic_add(dAcs_acc[ll], dseg)
-                        T.atomic_add(dAcs_acc[ss], -dseg)
+            # Map the L*L (ll,ss) grid over T.Parallel (the SAME un-fragmented
+            # thread-parallel idiom the dC apply / DYX build above use — a
+            # lane-strided T.serial here gets its masked body split across the
+            # shared-mem barriers TileLang inserts between the global cb load and the
+            # atomic scatter, hoisting the locals out of scope). The dseg compute
+            # lands in a thread-local; the +/- scatter (ll>ss mask) races on shared
+            # dAcs_acc -> T.atomic_add. dseg=0 off the strict lower-tri (no-op adds).
+            for ls in T.Parallel(L * L):
+                ll = ls // L
+                ss = ls % L
+                dseg = T.alloc_local((1,), accum_dtype)
+                dseg[0] = T.Cast(accum_dtype, 0)
+                if ss < ll:
+                    cb_v = T.Cast(
+                        accum_dtype, cb[batch_idx, chunk_idx, group_idx, ll, ss]
+                    )
+                    lmat = T.exp2((dacs[ll] - dacs[ss]) * p)
+                    dt_s = T.Cast(accum_dtype, dt[batch_idx, head_idx, chunk_idx, ss])
+                    dseg[0] = DYX[ll, ss] * cb_v * lmat * dt_s
+                T.atomic_add(dAcs_acc[ll], dseg[0])
+                T.atomic_add(dAcs_acc[ss], -dseg[0])
             T.sync_threads()
 
             for l in T.Parallel(L):