Commit 2e9ad32
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
- nn
- tests/v4
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
42 | 42 | | |
43 | 43 | | |
44 | 44 | | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
45 | 50 | | |
46 | | - | |
47 | | - | |
| 51 | + | |
48 | 52 | | |
| 53 | + | |
49 | 54 | | |
50 | 55 | | |
51 | 56 | | |
| |||
0 commit comments