@@ -367,3 +367,89 @@ reflection for the B2/B1/B0 delegated `_bwd` segments writes the owner grad buff
367367 mega-kernel the chunked path REPLACES; reproducible on a CLEAN checkout, NOT a
368368 regression from this change. Prior recorded serial baseline: loss 5.7134,
369369 6.106s. Eager-serial full-model reference (mamba=ref scan): loss ~ 5.70-5.71.
370+
371+ ================================================================================
372+ ## FINAL (branch ` mamba3-projection-bridge ` ) — ZERO-GRAD GAP RESOLVED; chunked
373+ ## backward now emits NONZERO grads; 163 grads flow with NONZERO mamba params;
374+ ## loss MATCHES serial. Residual per-grad gap is route-wide fp16 kernel precision,
375+ ## NOT the bridge (proven exact 4.5e-13).
376+ ================================================================================
377+
378+ ### What closed the prior "ONE REMAINING GAP" (the all-zero chunked backward grads)
379+ The prior gap was: B2/B1/B0 backward read back absmax=0.0 in the route. TWO fixes,
380+ both already present in this branch's uncommitted work and now VERIFIED end-to-end:
381+ 1 . ** B2 `` y `` input bound to the REAL F2 output** (path_c_fusion.py
382+ ` _emit_mamba3_chunked_bwd_model_brick_surfaces ` ): B2 needs the forward ungated
383+ SSD output `` y `` for the silu(z) gate transpose (dgate = dout* y). It was wired
384+ to a SEPARATE `` {name}_y `` buffer with NO forward producer (zero-seeded) — so
385+ the gate backward and everything downstream of dY collapsed to zero. The F2
386+ `` mamba3_chunk_scan_combine `` writes exactly this ungated y into `` {name}_delta ``
387+ (F2 applies no gate), so B2's `` y `` is now bound to `` {name}_delta `` . Single
388+ deterministic producer -> consumer (RULE #1 ).
389+ 2 . ** torch-MPS host-roundtrip ordering** (m04 ` _path_c_torch_mps_to_mlx ` ): the
390+ delegated grid kernels write their OUTPUT slots IN PLACE on MPS; a bare
391+ `` .to("cpu") `` could read the slot BEFORE the in-place store was host-visible
392+ (the `` _bwd `` grad slots came back zero despite the kernel writing them). Now a
393+ DEPENDENT device reduction (`` .abs().sum().item() `` ) + `` torch.mps.synchronize() ``
394+ orders/materializes the store before the host hop. RULE #1 : a stale read is a
395+ silent ZERO gradient — this forces the correct value, not a fallback.
396+
397+ ### UNLOCK PROOF (Metal, local_gb10_quarter smoke, seq=64, batch=1, flag ON)
398+ - Flag-ON route: install status=ok, ** 163 grads** , all NONZERO, loss FINITE.
399+ scratch/parity_on_vs_serial_same_model.py: route loss 5.68-5.71, serial-eager
400+ (fp32 ref) loss 5.69-5.73, ** loss diff ~ 0.012** (within fp16-chunked tolerance).
401+ 163 common grads, 0 only-route, 0 only-serial (full grad NAME coverage).
402+ - Layer-10 mamba (the BRIDGE TARGET) per-grad max|abs| vs fp32 serial reference:
403+ D 6e-7, C_bias 3e-5, B/C_norm 7e-5, dt_bias 1e-4, B_bias 2e-4 (all << 1e-3);
404+ out_proj 1.1e-3, in_proj 2.7e-3, conv_weight 9e-3, conv_bias/norm 1.4-3.1e-2.
405+ These are among the SMALLEST grads in the whole model — every NON-mamba layer
406+ (attention 0.1-0.2, mlp 0.1-0.16, m2rnn 0.25-0.56) is LARGER, and varies
407+ run-to-run (m2rnn conv_bias 0.265 -> 0.556 across two runs) — proving the diff
408+ is route-wide fp16 Metal-kernel precision + nondeterministic kernel scheduling,
409+ NOT a bridge defect.
410+
411+ ### Bridge correctness is PROVEN EXACT (independent of route fp16 noise)
412+ - tests/test_mamba3_projection_bridge.py: ** 8 passed** . Forward kernel-ABI mapping
413+ reproduces OUR serial scan to <1e-4 (fp32). Full composition (proj-fwd -> chunked
414+ scan -> scan-VJP -> bridge-VJP) == serial autodiff param grads to ** 4.5e-13** for
415+ all 8 projection params. out_proj correctly ZERO (post-scan). dt_k strictly > 0.
416+ Owner seeds nonzero flag-ON; coverage complete flag-ON; flag-OFF seeds nothing.
417+ - tests/test_mamba3_chunked_backward_b0b1b2.py: ** 13 passed** , chunked B0/B1/B2
418+ kernel-level parity worst 3.84e-4 (fp16).
419+
420+ ### THE PRECISE RESIDUAL GAP (named honestly)
421+ Strict <1e-3 per-grad parity vs the fp32 serial reference is NOT met for the
422+ layer-10 conv/in_proj/norm grads (1e-3..3e-2). ROOT CAUSE: the chunked grid kernels
423+ run their tensor math at fp16 internal precision (the documented dtype bridge: the
424+ KernelParam dtype is fp16, read-slots cast to fp16). Accumulated through the full
425+ projection (in_proj/conv/silu/RMSNorm/RoPE/trapezoid) the fp16 rounding grows from
426+ the isolated-kernel 3.84e-4 to ~ 1e-2 on the conv params. This is the SAME fp16
427+ effect that puts EVERY non-mamba route layer at 1e-1+. The proper apples-to-apples
428+ control (flag-OFF route: identical fp16 kernels everywhere EXCEPT layer-10) is
429+ XPC-blocked in this environment on the chain_12_13 mega-kernel, so the only
430+ available baseline is fp32-eager — which double-counts the route-wide fp16 gap.
431+ To reach strict <1e-3: widen the chunked grid kernels' internal compute to fp32
432+ (kernel change in mamba3_chunked_scan_core / * _ bwd_core), then re-verify. The
433+ BRIDGE needs no change (it is exact at 4.5e-13).
434+
435+ ### COMPLETE UNLOCK ASSESSMENT
436+ - 163 grads with REAL nonzero mamba projection params: YES (was 18/partial-zero).
437+ - loss MATCHES serial: YES (diff ~ 0.012, fp16-chunked tolerance, NOT just finite).
438+ - per-grad < 1e-3 vs serial: layer-10 bridge target mostly YES (D/biases/norms);
439+ conv/in_proj NO at fp16 (1e-3..3e-2) — gated by the chunked KERNELS' fp16
440+ precision (route-wide, not the bridge). Whole-route worst 0.1-0.56 is fp16.
441+ - honest wall-time: flag-ON fwd+bwd ~ 3.9-4.3s (163 grads); fp32 eager-serial
442+ reference ~ 0.01-0.05s (it is plain MLX autodiff with no Metal kernel
443+ compile/launch — not the descriptor route, which is XPC-blocked here).
444+ - flag-OFF unchanged: YES — 124 fusion_ir+autosplit passed, 13 chunked-backward
445+ passed, 82 m04 direct_chain passed with the SAME 3 pre-existing failures
446+ (verified IDENTICAL on a clean worktree of base 956f8ae).
447+
448+ ### Merge-safety (flag default OFF) — GREEN
449+ - test_path_c_fusion_ir.py + test_path_c_autosplit_metal_parity.py: 124 passed.
450+ - test_mamba3_chunked_backward_b0b1b2.py: 13 passed.
451+ - test_mamba3_projection_bridge.py: 8 passed.
452+ - test_m04_train_step.py -k direct_chain|path_c|chain: 82 passed, 3 FAILED —
453+ the SAME 3 pre-existing (test_fp8_ ..._ missing_sparse_mla_producer + 2
454+ direct-chain value_and_grad bridge tests), VERIFIED identical on base 956f8ae
455+ worktree (`` batch mapping must contain a 'tokens' array `` fixture errors).
0 commit comments