[Perf] Hoist DSv4 draft-extend SWA write locs; unify SWA graph buffer naming - #34890
Merged
Merged
Conversation
hnyls2002
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
hebiao064,
ispobock and
merrymercy
as code owners
August 15, 2026 00:26
hnyls2002
force-pushed
the
lsyin/eagle-verify-war-publish
branch
from
August 17, 2026 22:40
b2cb13c to
15b2bca
Compare
hnyls2002
changed the base branch from
main
to
lsyin/draft-extend-input-staging
August 17, 2026 22:40
hnyls2002
force-pushed
the
lsyin/eagle-verify-war-publish
branch
from
August 18, 2026 05:28
1bf408a to
b8eb703
Compare
hnyls2002
requested review from
JustinTong0323,
alphabetc1,
ch-wan,
fzyzcjy,
hanming-lu,
hzh0425,
kpham-sgl,
mickqian,
pyc96,
slin1237,
sundar24295s,
yctseng0211,
yhyang201,
yizhang2077 and
yuan-luo
as code owners
August 18, 2026 05:28
hnyls2002
requested review from
AgainstEntropy,
HydraQYH,
Kangyan-Zhou,
bingxche,
celve,
ping1jing2,
wisclmy0611 and
zijiexia
as code owners
August 18, 2026 05:28
hnyls2002
changed the base branch from
lsyin/draft-extend-input-staging
to
main
August 18, 2026 05:29
Collaborator
Author
|
/tag-and-rerun-ci |
Collaborator
Author
|
/rerun-test test_deepseek_v4_flash_fp4_b200.py test_deepseek_v4_flash_fp8_h200.py test_swa_radix_cache_kl.py |
Contributor
|
Results for 🚀 🚀 🚀 |
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
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.
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
swa_out_cache_locunset, so every SWA layer'sstore_cacherecorded afull_to_swa_index_mapping[out_cache_loc]translate inside the graph (one per layer per replay).init_cuda_graph_statenow pre-allocatescuda_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 capturedstore_cachetakes the cached path -- one translate per step, off the replay.Cross-backend alignment
swa_out_cache_loc_buf->cuda_graph_swa_out_cache_locin the flashattention and ascend backends; the SWA write-target graph buffer now has one name across triton / trtllm_mha / flashinfer / aiter / fa3 / ascend / dsv4.full_to_swa_index_mappinglive 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