Skip to content

[Perf] Hoist DSv4 draft-extend SWA write locs; unify SWA graph buffer naming - #34890

Merged
hnyls2002 merged 5 commits into
mainfrom
lsyin/eagle-verify-war-publish
Aug 18, 2026
Merged

hnyls2002 merged 5 commits into
mainfrom
lsyin/eagle-verify-war-publish

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Aug 15, 2026

Copy link
Copy Markdown
Collaborator

Hoist the DSv4 draft-extend SWA write-loc translate out of the captured graph, and align the remaining backends on the shared buffer conventions.

DSv4 draft-extend

  • Previously the out-graph metadata init left swa_out_cache_loc unset, so every SWA layer's store_cache recorded a full_to_swa_index_mapping[out_cache_loc] translate inside the graph (one per layer per replay).
  • init_cuda_graph_state now pre-allocates cuda_graph_swa_out_cache_loc (draft runner only); the metadata init refills it outside the graph (zeroed tail, [:num_tokens] view) and stamps it, so the captured store_cache takes the cached path -- one translate per step, off the replay.
  • Eager miss (no graph state / oversized batch) stamps a fresh per-step translate instead of falling back to per-layer translation at store time; capture always fits the pre-sized buffer, so it never records a temporary tensor's address.

Cross-backend alignment

  • Rename swa_out_cache_loc_buf -> cuda_graph_swa_out_cache_loc in the flashattention and ascend backends; the SWA write-target graph buffer now has one name across triton / trtllm_mha / flashinfer / aiter / fa3 / ascend / dsv4.
  • XPU: read full_to_swa_index_mapping live from the pool instead of caching the tensor ref at init; drop a redundant pool re-assignment.

CI States

Latest PR Test (Base): 🚫 Run #32104845985
Latest PR Test (Extra): ❌ Run #32104845896

@hnyls2002
hnyls2002 force-pushed the lsyin/eagle-verify-war-publish branch from b2cb13c to 15b2bca Compare August 17, 2026 22:40
@hnyls2002 hnyls2002 changed the title [Spec] Publish EAGLE v2 WAR read-done at verify via draft-extend input pre-staging [Spec] Publish EAGLE v2 shared-read-done at verify with staged draft-extend Aug 17, 2026
@hnyls2002
hnyls2002 changed the base branch from main to lsyin/draft-extend-input-staging August 17, 2026 22:40
@github-actions github-actions Bot added documentation Improvements or additions to documentation quant LLM Quantization amd lora diffusion SGLang Diffusion model-gateway jit-kernel labels Aug 18, 2026
@hnyls2002
hnyls2002 changed the base branch from lsyin/draft-extend-input-staging to main August 18, 2026 05:29
@github-actions github-actions Bot added the npu label Aug 18, 2026
@hnyls2002 hnyls2002 changed the title [Spec] Publish EAGLE v2 shared-read-done at verify with staged draft-extend [Perf] Hoist DSv4 draft-extend SWA write locs; unify SWA graph buffer naming Aug 18, 2026
@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@hnyls2002

Copy link
Copy Markdown
Collaborator Author

/rerun-test test_deepseek_v4_flash_fp4_b200.py test_deepseek_v4_flash_fp8_h200.py test_swa_radix_cache_kl.py

@github-actions

github-actions Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Results for /rerun-test test_deepseek_v4_flash_fp4_b200.py test_deepseek_v4_flash_fp8_h200.py test_swa_radix_cache_kl.py:

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp8_h200.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/radix_cache/test_swa_radix_cache_kl.py

@hnyls2002
hnyls2002 merged commit 0111b29 into main Aug 18, 2026
192 of 267 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

amd deepseek diffusion SGLang Diffusion documentation Improvements or additions to documentation jit-kernel lora model-gateway npu quant LLM Quantization run-ci

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant