fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv - #339
Open
dejay2 wants to merge 1 commit into
Open
fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv#339dejay2 wants to merge 1 commit into
dejay2 wants to merge 1 commit into
Conversation
causal_conv1d_varlen sized its triton launch grid from the longest
request in the batch, and the only place that number existed was on the
device: the triton fallback fell back to int(seq_lens.max().item()), a
D2H sync. Every prefill therefore paid a full pipeline stall to read
back a number the scheduler already knew, and a sync is illegal inside
a stream capture, so the prefill forward of every GDN/KDA model was
uncapturable.
build_fla_metadata computes the per-request lengths on the host, so
carry the max there (FLAMetadata.max_seq_len) and thread it down through
the three linear-attention ops (qwen3_5_moe, qwen4_exp, glm5_next) into
the kernel wrapper. The kwarg is optional and the device-derived path is
unchanged when it is omitted, so no other caller has to change.
Tested on an RTX 5090 (triton fallback path, no sgl_kernel):
python -m pytest -q tests/kernels/test_causal_conv1d_capture.py \
tests/models/qwen4_exp/test_gdn.py \
tests/models/test_glm5_next_kda_snapshot.py \
tests/models/test_glm5_next_kda_op.py \
tests/kvcache/test_linear_state_pool_alloc.py
26 passed (21 before this change, 5 new).
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
causal_conv1d_varlenonly needs the longest request in the batch to size itstriton launch grid, but the only place that number lived was on the device, so the
triton fallback derived it with
int(seq_lens.max().item()).build_fla_metadataalready computes the per-request lengths on the host, so itnow carries
FLAMetadata.max_seq_len, and the three linear-attention ops that callthe conv (
qwen3_5_moe/gdn.py,qwen4_exp/gdn.py,glm5_next/kda.py) pass itdown. The new kwarg is optional; with it omitted the kernel wrapper derives the
value on device exactly as before.
Why
int(seq_lens.max().item())is a device-to-host sync on every prefill: the wholepipeline stalls to read back a number the scheduler computed on the host in the
first place. It is also illegal inside a CUDA stream capture, so its presence alone
makes the prefill forward of every GDN/KDA model uncapturable. Passing the host
value removes both problems without touching the kernel.
How it was tested
Windows 11, RTX 5090, triton fallback path (no
sgl_kernelinstalled).The new
tests/kernels/test_causal_conv1d_capture.pypins that the host-metadatapath performs no
.item()at all, that the default device-derived path stillworks, that both produce bit-identical output and conv-state updates, and that the
call captures into and replays from a
torch.cuda.CUDAGraph.What is NOT included
kernel/triton/causal_conv1d_triton.py: it already accepts an optionalmax_seq_lenand falls back to the device-side max when it isNone. This PRonly supplies the value.
graph capture. That helper has no caller outside the fork's speculative-decoding
graph runner, so it is left out here.
🤖 Generated with Claude Code
https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK