Support TreeUpdateSolve kernel#576
Conversation
Shape KScaleReuse FragmentBV16 Speedup ━━━━━━━━━ ━━━━━━━━━━━━━ ━━━━━━━━━━━━━━ ━━━━━━━━━ B1 T33 146.854 36.916 3.98x ───────── ───────────── ────────────── ───────── B1 T49 224.354 34.250 6.55x ───────── ───────────── ────────────── ───────── B1 T64 276.479 35.604 7.77x ───────── ───────────── ────────────── ───────── B1 T128 708.708 148.979 4.76x ───────── ───────────── ────────────── ───────── B1 T256 1137.688 537.812 2.12x ───────── ───────────── ────────────── ───────── B1 T512 2027.250 1816.646 1.12x
….6x faster
- BuildTreeGram now emits kh0 (reuses its k loads), packed block-pair A
[B*HV, NB, ceil(NB/2), 16, 32] f32, and compact a_inv [B*HV, NB, 16, 16]
- TreeUpdateSolve (Gdn prefix dropped) is a pure block TRSM: RHS is one kh0
fragment load; history consumed one pair tile per MMA (K=32) + tail
- use_h0 function constant compiles the kh0 path out when no initial state
- BV trimmed to {16, 32}; BT template axis removed
- M5 Max, MXU BV32: B1 T64 23.5->15 us, B1 T256 126->77, B8 T512 3914->2762
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
cc @ry2009 |
ry2009
left a comment
ry2009
left a comment
There was a problem hiding this comment.
went through the whole thing against the cuda reference (cpu ref, both metal kernels, tests). semantics line up exactly with the validated triton path -- same rhs factorization, same block substitution, ragged guards all trace clean. the pair-packed A + kh0-in-gram is the same optimization we just landed on the cuda side (we hoisted kh0 to its own parallel kernel instead bc our gram is per key head -- same tradeoff, opposite layout, so the two backends now cross-validate each other, which is exactly what the reference relationship is for). left inline asks -- the main ones are a real-scale oracle test, kh0 dtype, and the per-value-head gram duplication question. nice work, the 97-token ragged test case is a good catch.
| const bool use_h0 SPECIALIZE, | ||
| const ThreadContext thread_context, | ||
| const uint batch_idx GROUPS(batch_size), | ||
| const uint value_head_idx GROUPS(value_heads), |
There was a problem hiding this comment.
this makes the gram per value head, so the k-gram tile + A storage get duplicated G x (G=2 on qwen3.6 shapes) vs the old per-key-head gram. it's the price of fusing kh0 in -- we hit the same fork on the cuda side and went the other way (per-key-head gram + a separate parallel kh0 kernel) to keep A deduped. did you measure fused-per-hv vs per-key-head + standalone kh0? if not there might be another 10-20% hiding behind the 1.6x
There was a problem hiding this comment.
Interesting, will try to bench, now for T=64, B=1 case, all kernels on M5 max are ~10us, not sure adding a new kernel with dispatch overhead etc. will make separation worthy.
Let me try.
There was a problem hiding this comment.
For MXU, the fused one is always better
There was a problem hiding this comment.
For simdgroup on M5 Max, split wins starting from T >= 128
There was a problem hiding this comment.
We can of course have two implementations (fused and split) but are we really using T >= 128.
If not, I'd prefer keeping fused one, so binary size of uzu stay lean.
| let beta = (0..scalar_len).map(|i| 0.25 + ((i as f32 * 0.013).sin() + 1.0) * 0.2).collect::<Vec<_>>(); | ||
| // Packed [B * HV, NB, ceil(NB/2), BT, 2*BT] block-pair tiles; strictly-lower | ||
| // blocks get values, everything else zero (the kernel never reads it). | ||
| let a_f32 = (0..a_len) |
There was a problem hiding this comment.
these synthetic A entries are ~0.01 with the cpu port of the same algorithm as the oracle -- great for catching indexing/layout bugs (it would), but it never stresses conditioning. real A sits at ||A||~1 with l2-normed keys + softplus decays, and the deep (I+A)^-1 cascades are where precision actually dies -- that's what forced tf32x3 on the cuda side. can we add one end-to-end case against a true sequential-recurrence oracle at real magnitudes? i'll send you the exact input recipe from our validation kit (random_tree + normalized keys + softplus decay + h0 ~0.3 scale), ports in ~30 min
There was a problem hiding this comment.
Is it in sglang-mirai/test/registered/spec/utils/test_gdn_chunk_tree_verify.py?
There was a problem hiding this comment.
Bumped the synthetic A to ~0.05 so it's off the toy scale.
On the full independent oracle: I ported your recipe and it passes, but I'm going to hold off landing it. It's ~350 LoC for a chain that isn't even wired into the model path yet, and traces will give us the real end-to-end check at integration anyway (against the actual model, not my re-derivation of it). Mostly a maintainability call. Keeping your recipe for when that lands.
One thing I'm curious about though: we solve (I+A)U=b with exact block forward-substitution in f32, not a Neumann series, so does ‖A‖~1 actually bite us the way it forced tf32x3 on your side?
3 participants
Uh oh!
There was an error while loading. Please reload this page.