Skip to content

perf(glm5next): run the KDA gate projections on a side stream - #582

Closed
MadeBy561 wants to merge 1 commit into
local-inference-lab:dev/jovian-judgementfrom
MadeBy561:perf/glm53-kda-gate-side-stream
Closed

perf(glm5next): run the KDA gate projections on a side stream#582
MadeBy561 wants to merge 1 commit into
local-inference-lab:dev/jovian-judgementfrom
MadeBy561:perf/glm53-kda-gate-side-stream

Conversation

@MadeBy561

Copy link
Copy Markdown

Purpose

The GLM-5.3 KDA layer's two low-rank output-gate projections (g_a_proj then g_b_proj) depend only on the layer input but trail the large fused in_proj_qkvgfab GEMM 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.forward forks a per-device side stream before in_proj_qkvgfab, runs g_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.
  • The optional _l2_prefetch_hook of the shared forward (perf(glm5next): L2 weight prefetch for SM120 decode (+7% C1 steps/s) #576) is honoured at the same point.
  • Full-rank-gate configurations keep the shared forward. VLLM_GLM53_KDA_GATE_SIDE_STREAM=0 restores 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):

python -m pytest tests/models/test_glm5next_kda_gate_stream.py
2 passed

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=0 on 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.py in 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.

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>
@coderabbitai

coderabbitai Bot commented Sep 2, 2026

Copy link
Copy Markdown

Warning

Review limit reached

Next included review available in 14 minutes.

Check out review usage here.

View limit details

Limit 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.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: 8f9d58bc-3013-4919-b098-c46e0971dac3

📥 Commits

Reviewing files that changed from the base of the PR and between 9c4dd05 and 4df4cf5.

📒 Files selected for processing (2)
  • tests/models/test_glm5next_kda_gate_stream.py
  • vllm/models/glm5next/nvidia/kda.py

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@voipmonitor

Copy link
Copy Markdown

Independent validation

The focused SM120 suite passed both tests: the gate projections execute on the
side stream while the remaining projections stay on the main stream, the
result matches the sequential path bitwise, and CUDA-graph replay matches eager
execution.

Matched serving A/B used four stock-clock NVIDIA RTX PRO 6000 Blackwell
Workstation Edition GPUs, TP4/DCP1, an NVFP4 target, an MXFP8 DFlash2 K7 draft,
15 seconds of warmup, and three 30-second context-zero cells per concurrency.
The two arms differed only by VLLM_GLM53_KDA_GATE_SIDE_STREAM:

Concurrency Sequential steps/s Side-stream steps/s Change Sequential tok/s Side-stream tok/s Change
C1 86.62 86.63 +0.02% 220.83 219.30 -0.70%
C8 288.59 286.88 -0.59% 745.06 752.64 +1.02%
C12 347.60 350.60 +0.86% 917.66 905.57 -1.32%

The mixed target-step result is within the observed run-to-run band; output
throughput additionally depends on accepted length. This qualifies the
side-stream dependency graph and bitwise behavior but does not establish a
material isolated E2E gain. The complete source-locked stack retains the path
because it removes main-stream work without changing arithmetic.

@voipmonitor

Copy link
Copy Markdown

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.

@voipmonitor voipmonitor closed this Sep 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants