Skip to content

Commit 2e9ad32

Browse files
committed
feat(v4): GDN bwd K!=V + KDA bwd V>32 doc/coverage + FP8 MoE training parity
GDN Path B backward (linear_attention_path_b_bwd.py): - Lift the head_k_dim == head_v_dim constraint. Per-thread register state is now length K (loop dim); thread grid x-axis is V (reduction source). Threadgroup padded to 32 * ceil(max(K,V)/32) lanes so simd_sum covers a full simdgroup and the reduce-store loop addresses every K-output even when K > tg_size. - State_hist workspace generalised from [B,H,T+1,Dh,Dh] to [B,H,T+1,K,V]; output shapes correct for asymmetric heads. - Cap raised from "max == 32 single simd OR head_dim<=256" to "max(K,V) <= 256" (multi-simdgroup path was already in place). KDA Path B backward (kda_path_b_bwd.py): - Stale docstring claimed V>32 unsupported; the multi-simdgroup shared-mem path actually handles V up to 256 already. Doc fixed and a V=64 parity test + V=128 finite-grad test added so the constraint is enforced by tests, not by stale prose. FP8 MoE (moe_fp8.py + test_moe_fp8.py): - Add a production training-loop parity test: bf16 master weight, re-quantize to fp8 each step for the forward, take grads through the bf16 reference, 5 SGD steps must monotonically decrease the fp8 loss; per-step fp8-vs-bf16 drift stays within fp8 tolerance. - Add a re-quantization determinism test so requantize(W)==quantize(W) bit-for-bit (prerequisite for stable training). - Drop the "training-time parity is not yet covered" caveat from the module docstring; point at the test that now covers it. Suite: 349 passed / 0 skipped (was 343).
1 parent f0390c5 commit 2e9ad32

6 files changed

Lines changed: 362 additions & 178 deletions

File tree

cppmega_v4/_tilelang/kda_path_b_bwd.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,15 @@
4242
dS_{t-1}[i,j] = dS_decayed[i,j] * decay_t[i]
4343
dg_t[i] = ddecay[i] * decay_t[i] (per-K)
4444
45+
Threadgroup layout: pad to ``32 * ceil(V / 32)`` lanes so simd_sum spans
46+
one full simdgroup; ``V > 32`` rides the multi-simdgroup path with
47+
threadgroup-shared-memory cross-simdgroup reductions (replaces the
48+
atomic_fetch_add path that serialised at V>=64). Cap is ``V <= 256``.
49+
4550
Falls back to ``mx.grad`` through ``naive_recurrent_kda`` for:
46-
- ``V > 32`` (multi-simdgroup not yet implemented)
47-
- ``initial_state`` provided
51+
- ``V > 256`` (per-thread register pressure exceeds Apple GPU limits)
4852
- ``HV % H != 0``
53+
- ``initial_state`` provided
4954
- any future shape outside the kernel's domain.
5055
"""
5156

0 commit comments

Comments
 (0)