Skip to content

fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv - #339

Open
dejay2 wants to merge 1 commit into
FlashML-org:mainfrom
dejay2:pr/gdn-graph-capture
Open

fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv#339
dejay2 wants to merge 1 commit into
FlashML-org:mainfrom
dejay2:pr/gdn-graph-capture

Conversation

@dejay2

@dejay2 dejay2 commented Sep 2, 2026

Copy link
Copy Markdown

What

causal_conv1d_varlen only needs the longest request in the batch to size its
triton 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_metadata already computes the per-request lengths on the host, so it
now carries FLAMetadata.max_seq_len, and the three linear-attention ops that call
the conv (qwen3_5_moe/gdn.py, qwen4_exp/gdn.py, glm5_next/kda.py) pass it
down. 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 whole
pipeline 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_kernel installed).

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
  • before this commit: 21 passed
  • after this commit: 26 passed (5 new)

The new tests/kernels/test_causal_conv1d_capture.py pins that the host-metadata
path performs no .item() at all, that the default device-derived path still
works, 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

  • Nothing in kernel/triton/causal_conv1d_triton.py: it already accepts an optional
    max_seq_len and falls back to the device-side max when it is None. This PR
    only supplies the value.
  • The fork branch this comes from also primes the fla chunk-index cache before
    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

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

1 participant