perf(glm5next): run the KDA gate projections on a side stream - #582
perf(glm5next): run the KDA gate projections on a side stream#582MadeBy561 wants to merge 1 commit into
Conversation
The two low-rank output-gate projections of the GLM-5.3 KDA layer (g_a_proj then g_b_proj) depend only on the layer input but trailed the large fused in_proj_qkvgfab GEMM on the main stream. Issue them on a side CUDA stream forked before in_proj and joined before the gate states are read, so they overlap the GEMM. Kernels, shapes and reduction orders are unchanged (outputs bitwise identical); only the fork/join edges are new, and CUDA graph capture records them as dependencies. VLLM_GLM53_KDA_GATE_SIDE_STREAM=0 restores the sequential forward. The optional _l2_prefetch_hook of the shared forward is honoured. GLM-5.3-Flash TP4 decode: 0.70 ms of gate GEMMs per verifier step leave the main stream (greedy C1 trace). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
Warning Review limit reachedNext included review available in 14 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Independent validationThe focused SM120 suite passed both tests: the gate projections execute on the Matched serving A/B used four stock-clock NVIDIA RTX PRO 6000 Blackwell
The mixed target-step result is within the observed run-to-run band; output |
|
Superseded by #619. The replacement preserves MadeBy561 as the author of the KDA gate side-stream commit and adds the CUDA graph lifecycle contract required to prevent auxiliary work from crossing uncaptured-warmup and breakable-capture boundaries. The combined source passed repeated TP4 FULL graph startup in no-speculative, MTP3, and DFlash2 modes. |
Purpose
The GLM-5.3 KDA layer's two low-rank output-gate projections (
g_a_projtheng_b_proj) depend only on the layer input but trail the large fusedin_proj_qkvgfabGEMM on the main stream. Run them on a side CUDA stream so they overlap that GEMM. This is part of the R22 serving stack and is filed so the maintainer can evaluate it with the rest; its isolated gain has not been measured (see below).Behavior
Glm5NextLinearAttention.forwardforks a per-device side stream beforein_proj_qkvgfab, runsg_b_proj(g_a_proj(x))there, and joins before the gate states are read. Kernels, shapes and reduction orders are unchanged, so the output is bitwise identical to the shared sequential forward; only the fork/join edges are new, and CUDA graph capture records them as dependencies._l2_prefetch_hookof the shared forward (perf(glm5next): L2 weight prefetch for SM120 decode (+7% C1 steps/s) #576) is honoured at the same point.VLLM_GLM53_KDA_GATE_SIDE_STREAM=0restores it for all configurations.Validation
Tests inside the GLM-5.3 serving image on one RTX PRO 6000 Blackwell (a stub layer with recorded streams):
They check that the gate projections run on a stream other than the main one and every other projection on the main stream, that the result equals the sequential forward bitwise, and that the overlapped forward captures into a CUDA graph and replays equal to eager.
Measurements: in the greedy C1 torch-profiler trace of the serving stack (GLM-5.3-Flash NVFP4 target, MXFP8 DFlash2 K7 draft, TP4) the gate side stream carries 0.70 ms of GEMM time per verifier step that previously sat on the main stream (main stream 6.1 ms, second stream 4.35 ms, prefetch 2.05 ms, gate side stream 0.70 ms). The only end-to-end comparison on record is confounded: the #576 files as-is (CuTe prefetch kernel, no gate side stream) measured 89.7 verifier steps/s at C1 against 91.1 for the overlay stack with it, but that pair also differs in the prefetch kernel implementation. An A/B with
VLLM_GLM53_KDA_GATE_SIDE_STREAM=0on the same boot is the outstanding measurement; the change is filed for completeness of the R22 stack.Precision: identical kernels and inputs; bitwise-equal output by construction and by test. The R10n greedy gates (Estonia 29/30; LAVD 28 exact, 0 near, 1 fail, 1 truncated) ran with this change live.
Merge simulation against #576 completes without conflicts; #495 touches
kda.pyin a different region (its conflicts with the current branch are in files this change does not touch).Generated with Claude Code; the submitter reviewed the change and ran the validation on the listed hardware.