[Bug Fix] Fix DSv4 EAGLE draft CUDA graph replay metadata by hashkanna · Pull Request #25529 · sgl-project/sglang

hashkanna · 2026-05-17T13:41:24Z

Motivation

Fixes an EAGLE v2 draft CUDA graph replay metadata mismatch for DSv4 compressed attention. In the reported failure from #25512, the draft replay path pads raw_bs=3 to graph_bs=4:

EAGLE draft cuda graph replay padding: raw_bs=3 graph_bs=4 raw_out_cache_loc=9 padded_out_cache_loc=12 steps=3 topk=1

The runner padded batch_size, seq_lens, req_pool_indices, and positions, but DSv4 replay metadata still observed the raw out_cache_loc layout. DSv4 compressed attention builds replay metadata from cache locations before graph replay, so the mismatch can produce invalid page metadata and surface later as an async CUDA illegal memory access.

This PR makes the replay metadata contract explicit: graph replay sees the padded graph-bucket cache-location buffer, and DSv4 receives the per-draft-step slice it expects.

Refs #25512.

Modifications

Update EAGLEDraftCudaGraphRunner.replay to pass the padded buffers.out_cache_loc layout into draft attention replay metadata initialization when raw_bs is padded to a captured CUDA graph bucket.
Restore mutated ForwardBatch fields with try/finally so metadata/replay errors do not leave the batch in a padded state.
Keep padded graph lanes zero-filled so they write to reserved cache slot 0.
Update DeepseekV4MultiStepBackend.init_forward_metadata_replay_cuda_graph to split EAGLE's request-major draft cache-location layout into per-step slices before initializing each DSv4 draft-step backend.
Add a loud invariant in DSv4 decode replay metadata if a caller passes a full multi-step out_cache_loc buffer where a per-step slice is expected.
Add registered CPU regression coverage for:
- raw_bs=3 -> graph_bs=4 padded out_cache_loc replay wiring,
- ForwardBatch restoration after metadata init failure,
- DSv4 per-step cache-location slicing,
- rejection of unpadded EAGLE draft cache-location input.

Accuracy Tests

PYTHONPATH=python python -m compileall -q python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py python/sglang/srt/layers/attention/deepseek_v4_backend.py test/registered/unit/spec/test_eagle_draft_dsv4_padding.py
ruff check python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py python/sglang/srt/layers/attention/deepseek_v4_backend.py test/registered/unit/spec/test_eagle_draft_dsv4_padding.py
git diff --check
Focused source-level smoke test for the patched runner replay path and DSv4 multi-step slicing.
L4 CUDA smoke for the padded DSv4 split helper:
- input: [10,11,12,20,21,22,30,31,32,0,0,0]
- output: [[10,20,30,0], [11,21,31,0], [12,22,32,0]]

Speed Tests and Profiling

Not run. This patch does not change attention kernels. It does change DSv4 EAGLE replay metadata initialization to compute metadata per draft step instead of copying step-0 metadata, because DSv4 compressed-attention metadata is step-specific.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Missing run-ci label — add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked — run-ci is required first.

EAGLE draft CUDA graph replay can pad raw batches to a larger graph bucket. Pass the padded out_cache_loc layout into replay metadata init so padding lanes resolve to the reserved cache slot 0. DSv4 compressed attention builds per-step replay metadata from cache locations, so split the request-major EAGLE layout into per-step slices instead of copying step-0 metadata across draft steps. Add CPU-registered regression coverage for the raw_bs=3 graph_bs=4 shape from sgl-project#25512.

gemini-code-assist

Code Review

This pull request improves the integration between EAGLE speculative decoding and the DeepSeek-V4 backend by ensuring correct handling and padding of cache locations. Key changes include the introduction of a helper function to transform request-major cache layouts into step-major slices and updates to the CUDA graph runner to maintain batch state consistency using robust state restoration. Additionally, a comprehensive unit test suite has been added to validate the padding logic and error handling. As there were no review comments provided, I have no further feedback to provide.

… upstream sgl-project#25529) + c_plan barrier (sgl-project#32467) EAGLE draft cuda-graph replay padded batch_size/seq_lens/req_pool_indices/ positions to the captured bucket but left forward_batch.out_cache_loc at the RAW request layout. init_forward_metadata_out_graph then built replay metadata with a padded bs over a raw cache-location buffer — the layout mismatch produces invalid page metadata inside the captured graph and surfaces as an ASYNC illegal memory access at whatever host sync comes first (observed: process_batch_result_decode copy_done.synchronize, alloc_for_decode_prealloc; nondeterministic per TP rank; not reproducible under CUDA_LAUNCH_BLOCKING). This matches the upstream family sgl-project#25512/sgl-project#25529 (same padding signature raw_bs=3 graph_bs=4) and sgl-project#28569 (crash as the running batch shrinks — our mixed load with waves of completing requests triggers exactly the padding boundary; --disable-cuda-graph is immune). Fix: - execute() now swaps forward_batch.out_cache_loc to the zero-filled padded buffers slice (num_tokens * speculative_num_steps) during the padded window, so replay metadata sees the same layout the captured graph does; padded lanes write to reserved cache slot 0. - try/finally restores the raw batch view even when metadata init or replay raises (upstream sgl-project#25529 contract) — a padded leftover poisoned every later consumer of the ForwardBatch. Also ports sgl-project#32467: __syncthreads() between warp-scratch init and the per-warp min/max writes in c_plan.cuh plan_compress_prefill_kernel0 — warp 0 could clobber another warp's slot, misclassifying ragged extend as MTP-uniform and emitting out-of-bounds ragged_id (B300 non-deterministic IMA family sgl-project#33356). Dormant on the GLM path (no callers) but the kernel is shared with DSV4 deployments.

… (port sgl-project#25529 to target path) build_replay_fb_view swapped batch_size/seq_lens/req_pool_indices/positions to the padded capture buffers but passed forward_batch.out_cache_loc RAW (docstring even called it a per-iter unpadded field). DSA target-verify metadata then built a padded-bs view over a raw-layout cache-location buffer -> invalid KV slots written inside the captured verify graph -> async IMA (Xid 31 WRITE faults; surfaces at process_batch_result_decode sync / alloc paths). Same contract bug as the draft runner fixed in the previous commit; upstream sgl-project#25529 established replay metadata must see the padded layout. Registry zero-fills the padding tail (slot 0).

… upstream sgl-project#25529) + c_plan barrier (sgl-project#32467) EAGLE draft cuda-graph replay padded batch_size/seq_lens/req_pool_indices/ positions to the captured bucket but left forward_batch.out_cache_loc at the RAW request layout. init_forward_metadata_out_graph then built replay metadata with a padded bs over a raw cache-location buffer — the layout mismatch produces invalid page metadata inside the captured graph and surfaces as an ASYNC illegal memory access at whatever host sync comes first (observed: process_batch_result_decode copy_done.synchronize, alloc_for_decode_prealloc; nondeterministic per TP rank; not reproducible under CUDA_LAUNCH_BLOCKING). This matches the upstream family sgl-project#25512/sgl-project#25529 (same padding signature raw_bs=3 graph_bs=4) and sgl-project#28569 (crash as the running batch shrinks — our mixed load with waves of completing requests triggers exactly the padding boundary; --disable-cuda-graph is immune). Fix: - execute() now swaps forward_batch.out_cache_loc to the zero-filled padded buffers slice (num_tokens * speculative_num_steps) during the padded window, so replay metadata sees the same layout the captured graph does; padded lanes write to reserved cache slot 0. - try/finally restores the raw batch view even when metadata init or replay raises (upstream sgl-project#25529 contract) — a padded leftover poisoned every later consumer of the ForwardBatch. Also ports sgl-project#32467: __syncthreads() between warp-scratch init and the per-warp min/max writes in c_plan.cuh plan_compress_prefill_kernel0 — warp 0 could clobber another warp's slot, misclassifying ragged extend as MTP-uniform and emitting out-of-bounds ragged_id (B300 non-deterministic IMA family sgl-project#33356). Dormant on the GLM path (no callers) but the kernel is shared with DSV4 deployments.

… (port sgl-project#25529 to target path) build_replay_fb_view swapped batch_size/seq_lens/req_pool_indices/positions to the padded capture buffers but passed forward_batch.out_cache_loc RAW (docstring even called it a per-iter unpadded field). DSA target-verify metadata then built a padded-bs view over a raw-layout cache-location buffer -> invalid KV slots written inside the captured verify graph -> async IMA (Xid 31 WRITE faults; surfaces at process_batch_result_decode sync / alloc paths). Same contract bug as the draft runner fixed in the previous commit; upstream sgl-project#25529 established replay metadata must see the padded layout. Registry zero-fills the padding tail (slot 0).

github-actions · 2026-09-04T01:17:58Z

Thanks @hashkanna. Closing this because it has had no updates in 109 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.

hashkanna requested review from Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, hebiao064, hnyls2002, ispobock and merrymercy as code owners May 17, 2026 13:41

github-actions Bot added the deepseek label May 17, 2026

gemini-code-assist Bot reviewed May 17, 2026

View reviewed changes

PierreLeGuen mentioned this pull request Jul 21, 2026

dsv4-flash: disable EAGLE on gpu26 incident canary nearai/cvm-compose-files#146

Merged

github-actions Bot closed this Sep 4, 2026

Navigation Menu

Uh oh!

[Bug Fix] Fix DSv4 EAGLE draft CUDA graph replay metadata - #25529

[Bug Fix] Fix DSv4 EAGLE draft CUDA graph replay metadata#25529
hashkanna wants to merge 1 commit into
sgl-project:mainfrom
hashkanna:kanna/issue-25512-dsv4-padding

hashkanna commented May 17, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

hashkanna commented May 17, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

github-actions Bot commented Sep 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

hashkanna commented May 17, 2026 •

edited by github-actions Bot

Loading