[Fix] Use int64 seq_lens across all CUDA graph runners and backends - #27840
Merged
Merged
Conversation
The eager forward path routes through the CUDA-graph input-buffer registry, which mirrored `seq_lens` / `seq_lens_cpu` as int32 (the cuda-graph buffer convention). The canonical `ForwardBatch.seq_lens` is int64, and the `intel_amx` CPU decode kernel requires int64 `req_lens`, so on CPU the registry downcast crashed with `decode: expect req_lens to be int64, got Int` (the Xeon CI failure #27549 worked around with a `device == "cuda"` gate without addressing the root cause). Make `seq_lens` / `seq_lens_cpu` int64 uniformly so every backend receives the canonical dtype (GPU backends convert to int32 internally where needed): - int64 buffer allocations: `DecodeInputBuffers.create`, `_allocate_decode_buffers`, `build_decode_registry` slots, and the speculative-decode runners (eagle draft / draft-extend / frozen-kv-mtp, eagle_info placeholders). - `flashinfer_mla` replay: downcast `kv_len_arr_cpu` to int32 before `wrapper.plan` (the capture path already converts). - Remove the `device == "cuda"` gate on the three `_eager_fb_view` call sites in `model_runner.py`; the eager path now runs through the registry on all devices. CPU unit tests (test_cuda_graph_buffer_registry.py) pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ch-wan
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
hebiao064,
hnyls2002,
ispobock,
kpham-sgl,
merrymercy and
pyc96
as code owners
June 10, 2026 22:26
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
Collaborator
Author
|
/tag-and-rerun-ci |
kpham-sgl
approved these changes
Jun 10, 2026
kpham-sgl
left a comment
kpham-sgl
left a comment
Collaborator
There was a problem hiding this comment.
Needed this change to fuse some memcpys / eager ops in metadata preparation path :D
With seq_lens now int64 in the cuda-graph decode buffers, the DSV4 decode metadata builder derived swa_topk_lengths / raw_positions / page indices from an int64 seq_lens_casual, and the sparse-decode kernel rejects non-int32 topk_length (`RuntimeError: topk_length must have dtype torch::kInt32`). Normalize seq_lens_casual to int32 at the single make_core_attn_metadata choke point (covers both the direct-pass decode callers and the expand_* paths), restoring the prior int32 invariant. Same fix for the HIP radix backend. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ch-wan
force-pushed
the
cheng/seq-lens-int64
branch
from
June 10, 2026 23:33
59bbfd8 to
77e1914
Compare
5 tasks
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Aug 24, 2026
jakki-amd
pushed a commit
to jakki-amd/sglang
that referenced
this pull request
Sep 9, 2026
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
2 participants
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.
Motivation
PR #27407 routed the eager forward path (
forward_decode/forward_extend/forward_idle) through the CUDA-graph input-buffer registry. That registry mirrorsseq_lens/seq_lens_cpuas int32 (the long-standing cuda-graph buffer convention), but the canonicalForwardBatch.seq_lensis int64, and theintel_amxCPU decode kernel strictly requires int64req_lens. On CPU the registry therefore handed the kernel a downcast tensor and crashed:This is the Xeon CI failure. #27549 worked around it with a
device == "cuda"gate on the three eager call sites, which fixed the symptom on CPU but left the int32 downcast in place (a silent dtype change on the CUDA eager path vs.main) and kept non-CUDA devices off the unified eager path.This PR fixes the root cause:
seq_lens/seq_lens_cpuare int64 uniformly across all runners, so every backend receives the canonical dtype. GPU backends already convert to int32 internally where their kernels need it (verified by auditing every attention backend), so the int32 buffer dtype was incidental, not a kernel requirement.Modifications
seq_lens/seq_lens_cpu:DecodeInputBuffers.create(runner_utils/buffers.py)_allocate_decode_buffers(runner/decode_cuda_graph_runner.py)build_decode_registryslots (cuda_graph_buffer_registry.py)eagle_draft_cuda_graph_runner,eagle_draft_extend_cuda_graph_runner,frozen_kv_mtp_cuda_graph_runner, and theeagle_infoplaceholders.flashinfer_mla: the cuda-graph replay path passedkv_len_arr_cpu(=seq_lens_cpu[:bs]) straight intowrapper.plan()while the capture path converts to int32 — added the matching.to(torch.int32).device == "cuda"gate on the three_eager_fb_viewcall sites inmodel_runner.py(back toif not self.server_args.enable_pdmux:). With seq_lens int64 everywhere the eager path runs through the registry on all devices.Backends that assert int32 (
cutlass_mla, DSV4 sparse-prefill,trtllm_mla,tokenspeed) are always fed a converted int32 at the call site, never the raw buffer, so they are unaffected.Accuracy Tests
CPU unit tests pass:
test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py(39 passed). GPU validation of the flashinfer / flashinfer_mla (DeepSeek-MLA) / trtllm / eagle / dflash paths should run in CI.Speed Tests and Profiling
N/A —
seq_lens/seq_lens_cpuare small bs-sized buffers; int64 doubles only those few bytes, and backends already materialize int32 derived tensors where their kernels need them.🤖 Generated with Claude Code
CI States
Latest PR Test (Base): 🚫 Run #27313253351
Latest PR Test (Extra): ❌ Run #27313253134