Skip to content

fix graph of ascend - #526

Open
Jingbo-gao wants to merge 1 commit into
InfiniTensor:mainfrom
Jingbo-gao:fix-graph-ascend
Open

fix graph of ascend#526
Jingbo-gao wants to merge 1 commit into
InfiniTensor:mainfrom
Jingbo-gao:fix-graph-ascend

Conversation

@Jingbo-gao

@Jingbo-gao Jingbo-gao commented Aug 4, 2026

Copy link
Copy Markdown

Summary

  • Update csrc/engine/compiler/paged_compiler.cpp to stage capture-time host sequence lengths for Ascend graph compilation.
  • Bind the current CPU total_sequence_lengths metadata to the captured graph input before every graph replay.
  • Validate that replay sequence lengths are a one-dimensional CPU int32 tensor matching the compiled batch size.
  • Limit the new behavior to the Ascend backend; other backends keep their existing compilation and replay paths.

Motivation

Ascend FIA receives actualSeqLengthsKv as host-side metadata. Copying a new total_sequence_lengths tensor 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 model
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Test configuration:

Platform Model Tensor Parallelism Graph Mode Result
Ascend, 8 NPUs FM9G_70B_SFT_MHA 8 Enabled To be confirmed and screenshot attached
Ascend, 8 NPUs FM9G_70B_SFT_MHA 8 Disabled Reference output; screenshot to be attached

Graph-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-graph

Eager 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-attn

Expected validation:

  • Graph replay completes without an FIA parameter-update error.
  • Graph output no longer degenerates into a repeated first token or repeated
    <|im_end|> tokens.
  • Graph and eager outputs are consistent within the expected model precision.

Benchmark / Performance Impact

N/A — this is a correctness fix, not a performance PR.

The implementation reuses CPU sequence-length metadata already available to
PagedCompiler and updates one graph-owned host array per replay. It avoids
introducing a device-to-host sequence-length copy in every attention layer.

Notes for Reviewers

  • The capture-time sequence lengths are initialized to one for every sequence,
    matching the graph capture input initialization.
  • Replay metadata is validated as CPU int32, one-dimensional, and equal to the
    compiled batch size before it is bound to the graph.
  • The binding key is the captured
    graph_input.total_sequence_lengths tensor, not the transient runtime tensor.
    This allows the graph operator to find the host metadata associated with the
    captured device input.
  • The new path is guarded by an Ascend device check and does not change other
    backends.
  • This PR is the InfiniLM-side caller change. It requires the corresponding
    InfiniCore graph host-array staging/update APIs and Ascend FIA task-update
    support.
  • Dedicated update streams, external events, FIA V5, and layout changes are
    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 /retest or /test.


Checklist

Every contributor must verify every item below before requesting review.
Items that depend on the final branch, commits, local test results, screenshots,
or CI are intentionally left unchecked.

Title, Branch, and Commits

  • PR title follows Conventional Commits, for example: fix(ascend): update FIA sequence lengths during graph replay.
  • Branch name follows <type>/xxx-yyyy-zzzz, where <type> matches the PR title.
  • Each commit message follows Conventional Commits.
  • Small PR is a single squashable commit; or every commit is meaningful, well-formed, and independently reviewable.
  • No stray merge commits from main; the branch is rebased cleanly on the current main.
  • No fixup!, squash!, or wip commits remain.
  • N/A — no existing PR, branch, or commit using the legacy issue format.

Scope and Design

  • Changes are minimal and limited to the stated motivation.
  • No dead code, commented-out blocks, debug prints, or unowned TODO items were added.
  • No unrelated formatting churn was introduced.
  • No InfiniLM public API was changed.

General Code Hygiene

  • Comments describe only non-obvious behavior.
  • The modified file ends with a trailing newline.
  • No trailing whitespace, tab/space mixing, or stray BOM was introduced.
  • Identifiers in comments and error messages follow the repository formatting requirements.
  • New comments and error messages are in English.
  • New comments and error messages have been reviewed for capitalization and terminal punctuation.

C++ Specific

  • Code follows the Google C++ Style Guide.
  • Error message wording follows the LLVM Coding Standards.
  • N/A — no constructor initializer list was changed.
  • No raw new or delete was introduced.
  • Changed files were formatted by scripts/format.py.
  • No changes or references to csrc/models/llama_legacy/ were introduced.

Python Specific

  • N/A — no Python files were changed.

Testing

  • For any platform that could not be tested, an explicit reason is provided and a reviewer with access is tagged.
  • Passed the single-request test using examples/test_infer.py, or documented why it was skipped.
  • Passed the offline benchmark using examples/bench.py, and attached the graph/eager comparison screenshot.
  • Passed the sanity test using test/bench/test_benchmark.py, or documented why it was skipped.
  • Passed the service test using python/infinilm/server/inference_server.py and scripts/test_perf.py, or documented why it was skipped.

Build, CI, and Tooling

  • The project builds cleanly from a fresh directory on at least one affected Ascend platform.
  • CI was triggered manually, or /retest was requested.

Documentation

  • N/A — no user-facing command, build flag, or developer workflow was changed.
  • N/A — this PR has no user-visible breaking change.

Security and Safety

  • Confirmed that no secrets, access tokens, internal URLs, customer data, or personal hardware identifiers are included in the final PR.
  • No third-party code was added.
  • No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.

@Jingbo-gao
Jingbo-gao requested a review from a team August 4, 2026 08:06
@github-actions

github-actions Bot commented Aug 4, 2026

Copy link
Copy Markdown

⛔ Only repository members can run retest.

@wooway777

Copy link
Copy Markdown
Collaborator

/test

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] 昇腾平台图模式开启输出异常

2 participants