Repository navigation
[Perf][GDN] Reuse prefill lengths and fuse FlashInfer gate conversion - #60105
YangXu1990uiuc wants to merge 1 commit into
Conversation
Use the existing post-convolution gate exponential only for the resolved FlashInfer prefill route and prepare its int64 sequence offsets once per batch. Keep log-gate defaults for standalone adapter calls and other providers. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Yang Xu <yanxu@nvidia.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
Overview
Reduce repeated host work in the existing FlashInfer GDN prefill route: let the fused post-convolution kernel produce the gate format FlashInfer consumes, and prepare its int64 sequence offsets once in batch metadata. This is a Draft prototype; model-level correctness and serving measurements are still pending.
Claims
torch.explaunch from each GDN prefill layer by enabling the producer's existingoutput_g_expoption.backend="flashinfer"; this PR does not enable automatic cuDNN dispatch.Validation
On candidate
6ed041d0e52b316e9c9bbad8e3ebf4508dc10575, based on mainc32495d58ff580b518cf4c5e5d6b69adc4d38351:72 passed. This was a CPU-only run with CUDA devices hidden and
vllm.utils.torch_utils.PIN_MEMORY=Falseset in the test runner; the installed CUDA platform would otherwise attempt pinned allocation. The runner did not alter production source. Coverage includes both gate formats, int32 conversion/int64 reuse, mixed decode/prefill and speculative classification, padding, and metadata reuse between cache groups.All local pre-commit hooks on the four changed files and
git diff --checkpassed. The remote pre-commit workflow currently stops at the repository author/label authorization gate before running code checks; this is not a remote lint pass.On a reserved, full-power B200, the real fused producer plus the exact parent/candidate FI adapter passed 4/4 eager and side-stream capture/replay cases. The offline harness extracts the adapter functions unchanged from source and calls real Triton and FlashInfer kernels; it does not import or run a full vLLM model. Both arms use the same FlashInfer checkout (
243c8aade55391d075aeb8fbc306120532e0e651), native FI GDN route, Torch 2.13.0+cu130, CUDA runtime 13.0, and one pinned CPU thread. Inputs are BF16, Hq=8, Hv=16, D=128.CPU active time during host submission, measured with
thread_time_ns, in microseconds. Entries show median [min–max] across six measurements, each averaging five calls, interleaved in three ABBA blocks after warmup:The ragged lengths are
[1, 63, 129, 511, 1024, 1537, 2048, 2879]. CPU timing excludes synchronization; replay is a separate CUDA-event measurement and must not be added to CPU time. These are per-layer producer-plus-adapter measurements: they exclude metadata construction/update, convolution, state-pool gather/scatter, and the rest of the model. In particular, the candidate's once-per-batch int64 conversion is outside this timer; it has not disappeared from the whole request. Captured lengths stay fixed.Both output and final state passed comparison against the parent before/after timing and after changing runtime inputs. Captured outputs/state were poisoned before replay, and the input state remained unchanged. The checks use
rtol=1e-2, atol=5e-3; the largest observed O/state absolute differences were 0.000245/0.001061. Q/K/V/beta were bit-identical and both gate implementations passed an FP64 gate reference check.Model evaluation, a serving run with
VLLM_GPU_SYNC_CHECK=error, and TTFT measurements are pending. No model-level accuracy equivalence or serving speedup is claimed, and this should remain Draft until that validation and human review are complete.Details
The producer already supports outputting
exp(g)directly. Both real execution and kernel warmup now select that option only when the resolved GDN prefill backend is FlashInfer, and pass an append-only privateg_is_expflag to its adapter. The flag defaults to false for existing callers. Triton and CuTeDSL continue receiving log gates.The metadata builder casts the selected, rebased prefill offsets after constructing the existing chunk metadata. The returned metadata owns the resulting tensor and shares it across layers/cache groups. Conv and recurrent decode retain their original int32 offsets. This adds one device-side conversion per metadata build and removes repeated conversions per layer; there is no new GPU-to-CPU readback.
The fused Triton exponential can differ slightly from the separate Torch exponential, so numerical validation must allow normal floating-point rounding rather than demand bitwise equality. Output allocation and state-update behavior are unchanged.
Current GDN graph policy supports full graphs for uniform decode, while prefill attention runs outside the retained graph through attention splitting or eager breaks. The batch-owned conversion is consistent with that policy. If full-graph GDN prefill is introduced later, its metadata must be staged into persistent graph-owned buffers.
Duplicate-work searches found no open PR implementing these two changes. #57831 addresses 64-bit token-offset arithmetic rather than per-layer metadata conversion. #60050 adds GDN checkpoints and is orthogonal, but needs an integration adjustment: its new
_warmup_flashinfer_prefill_variantscalls the FI adapter directly with the produced gates and must forwardg_is_exp=True(or the matching explicit flag) when these changes are combined. Otherwise that warmup would exponentiate twice.AI assistance was used for implementation, review, and validation. This Draft does not attest that human review or model evaluation has been completed.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.