fix graph of ascend - #526
Open
Jingbo-gao wants to merge 1 commit into
Open
Conversation
|
⛔ Only repository members can run |
Collaborator
|
/test |
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.
Summary
csrc/engine/compiler/paged_compiler.cppto stage capture-time host sequence lengths for Ascend graph compilation.total_sequence_lengthsmetadata to the captured graph input before every graph replay.int32tensor matching the compiled batch size.Motivation
Ascend FIA receives
actualSeqLengthsKvas host-side metadata. Copying a newtotal_sequence_lengthstensor into the captured device input does not, by itself, refresh the host parameter used by the captured FIA task.Without an explicit host binding, graph replay can continue using the sequence lengths from graph capture. This gives FIA an incorrect effective KV context length during decode and can produce repeated or otherwise incorrect output,
while eager execution remains correct.
This change supplies the current CPU sequence-length metadata to the graph before replay. It also stages valid initial values before capture so the captured FIA task can be initialized without performing a device-to-host copy
inside every attention layer.
Closes #525
Type of Change
feat— new feature / new modelfix— bug fixperf— performance improvement (no behavioral change)refactor— code restructuring without behavior changetest— adding or fixing tests onlydocs— documentation onlybuild/ci— build system or CI configurationchore— tooling, formatting, or other non-code changesTest Results of Involved Models on Supported Platforms (Please attach screenshots)
Test configuration:
FM9G_70B_SFT_MHAFM9G_70B_SFT_MHAGraph-mode command:
ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ python examples/bench.py \ --device ascend \ --model=/home/g00893922/FM9G_70B_SFT_MHA/ \ --tp=8 \ --input-len=32,32 \ --output-len=256 \ --batch-size=16 \ --enable-paged-attn \ --attn=flash-attn \ --enable-graphEager reference command:
ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ python examples/bench.py \ --device ascend \ --model=/home/g00893922/FM9G_70B_SFT_MHA/ \ --tp=8 \ --input-len=32,32 \ --output-len=256 \ --batch-size=16 \ --enable-paged-attn \ --attn=flash-attnExpected validation:
<|im_end|>tokens.Benchmark / Performance Impact
N/A — this is a correctness fix, not a performance PR.
The implementation reuses CPU sequence-length metadata already available to
PagedCompilerand updates one graph-owned host array per replay. It avoidsintroducing a device-to-host sequence-length copy in every attention layer.
Notes for Reviewers
matching the graph capture input initialization.
int32, one-dimensional, and equal to thecompiled batch size before it is bound to the graph.
graph_input.total_sequence_lengthstensor, not the transient runtime tensor.This allows the graph operator to find the host metadata associated with the
captured device input.
backends.
InfiniCore graph host-array staging/update APIs and Ascend FIA task-update
support.
intentionally outside the scope of this PR.
CI / ChatOps
CI does not run automatically on pull requests. Trigger it manually from the
Actions tab using the CI workflow on this PR branch, or ask a maintainer to
comment
/retestor/test.Checklist
Title, Branch, and Commits
fix(ascend): update FIA sequence lengths during graph replay.<type>/xxx-yyyy-zzzz, where<type>matches the PR title.main; the branch is rebased cleanly on the currentmain.fixup!,squash!, orwipcommits remain.Scope and Design
TODOitems were added.General Code Hygiene
C++ Specific
newordeletewas introduced.scripts/format.py.csrc/models/llama_legacy/were introduced.Python Specific
Testing
examples/test_infer.py, or documented why it was skipped.examples/bench.py, and attached the graph/eager comparison screenshot.test/bench/test_benchmark.py, or documented why it was skipped.python/infinilm/server/inference_server.pyandscripts/test_perf.py, or documented why it was skipped.Build, CI, and Tooling
/retestwas requested.Documentation
Security and Safety