[None][perf] Avoid paged MSA K/V materialization during prefill - #16857
[None][perf] Avoid paged MSA K/V materialization during prefill#16857peihu-nv wants to merge 10 commits into
Conversation
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #61667 [ run ] triggered by Bot. Commit: |
|
PR_Github #61667 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62023 [ run ] triggered by Bot. Commit: |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughChangesThe PR replaces MSA submodule packaging with FetchContent-based staging. It updates wheel extraction and runtime imports. It also adds paged HND KV preparation for MSA sparse attention. MSA integration
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant BuildSystem
participant FetchContent
participant WheelBuilder
participant Runtime
participant MSAKernel
BuildSystem->>FetchContent: fetch pinned msa source
FetchContent->>FetchContent: apply validated patch
WheelBuilder->>FetchContent: stage fmha_sm100 files
WheelBuilder->>Runtime: package fmha_sm100
Runtime->>MSAKernel: import packaged module
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #62023 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62108 [ run ] triggered by Bot. Commit: |
|
PR_Github #62108 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62216 [ run ] triggered by Bot. Commit: |
Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
scripts/build_wheel.py (1)
1147-1158: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winBuild the wheel once with the complete arguments and environment.
python -m buildruns twice. The first invocation usesextra_wheel_build_argsbut not the MYPYC environment; the second uses the environment but omitsextra_wheel_build_args, soEXTRA_WHEEL_BUILD_ARGScan be lost while build time is doubled. Move the environment setup before one invocation and passextra_wheel_build_argsto it.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/build_wheel.py` around lines 1147 - 1158, Update the wheel-build flow to perform a single build_run invocation: move the mypyc environment setup before it, and include both extra_wheel_build_args and plat_name_arg in that command. Remove the redundant first invocation so the complete arguments and environment are applied together.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py`:
- Around line 26-40: Update require_msa_module() to add a precise return
annotation using types.ModuleType, importing the types module as needed.
Preserve the existing deferred import, error handling, and returned fmha_sm100
module.
---
Outside diff comments:
In `@scripts/build_wheel.py`:
- Around line 1147-1158: Update the wheel-build flow to perform a single
build_run invocation: move the mypyc environment setup before it, and include
both extra_wheel_build_args and plat_name_arg in that command. Remove the
redundant first invocation so the complete arguments and environment are applied
together.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0f67641d-5e71-48b9-bf96-04a9ec35438a
📒 Files selected for processing (18)
.gitignore.gitmodules3rdparty/CMakeLists.txt3rdparty/MSA3rdparty/fetch_content.jsoncpp/CMakeLists.txtjenkins/Build.groovyjenkins/BuildDockerImage.groovyjenkins/L0_MergeRequest.groovyjenkins/L0_Test.groovyjenkins/TensorRT_LLM_PLC.groovyjenkins/runPerfSanityTriage.groovyscripts/attribution/scan/metadata/msa.ymlscripts/build_wheel.pysetup.pytensorrt_llm/_torch/attention_backend/fmha/msa_sparse_gqa.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_availability.pytensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_utils.py
💤 Files with no reviewable changes (2)
- .gitmodules
- 3rdparty/MSA
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
PR_Github #62216 [ run ] completed with state |
|
PR_Github #62254 [ run ] triggered by Bot. Commit: |
|
PR_Github #62254 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62540 [ run ] triggered by Bot. Commit: |
|
PR_Github #62540 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62644 [ run ] triggered by Bot. Commit: |
|
PR_Github #62644 [ run ] completed with state |
pengbowang-nv
left a comment
pengbowang-nv
left a comment
There was a problem hiding this comment.
Attention Part LGTM
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
I looked at this in depth. The kernel-side change itself is solid: I checked the pinned MSA source and confirmed the compiled kernels mark all non-leading strides dynamic (to_cute_tensor uses mark_layout_dynamic(leading_dim=ndim-1)), so the stride-agnostic compile cache is safe and the outer page stride is honored at runtime. The 16B checks in _prepare_paged_hnd_input match the from_dlpack(assumed_align=16) contract, and _prepare_paged_kv_for_tma at the pin is validation-only, so the view isn't re-materialized downstream. Nice perf win, well evidenced.
My concerns are with the scope and packaging around it:
- The PR description doesn't mention the CI changes at all: every
checkoutSourcecall site flipped, the submodule removed, precompiled-wheel extraction changed. For a change this invasive the description should cover them, and it should cite a JIRA ticket rather than[None]. - The FetchContent migration pulls in a second full cutlass tree and adds a full-tree staging copy on every build. That's the small-file pattern that already thrashes networked filesystems (TRTLLM-14628). Details inline.
- Pre-migration refs become unbuildable through the fixed-branch pipelines, which also matters for bisection. Details inline.
One doc leftover: docs/source/installation/build-from-source.md still tells users to run git submodule update --init --recursive. Since this PR removes the last submodule, drop that line here too.
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
…-paged-kv-main Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> # Conflicts: # tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #64077 [ run ] triggered by Bot. Commit: |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The zero-copy guard is conservative and well tested (stride-contract units here, bit-exact strided-vs-packed end-to-end test inside the MSA patch), and the FetchContent migration keeps the exact submodule pin with a reviewable in-tree patch. A few non-blocking points:
- Ticket: this is a nontrivial perf feature plus a build-system migration under a
[None]tag — it deserves a JIRA ticket for tracking. - Upstreaming: what's the plan for landing
msa_strided_paged_kv.patchin the MSA repo itself so the pin can eventually advance and the patch retire? Worth noting in the PR or a follow-up ticket. - Undescribed fix: the consolidation of the duplicated
-m buildinvocation inscripts/build_wheel.pyfixes a real pre-existing bug — the second (shipped) build droppedEXTRA_WHEEL_BUILD_ARGS— but the description doesn't mention it. Worth a line so it isn't lost if this PR gets split or reverted. - CI coverage: for the record, the new unit tests do run in pre-merge CI —
l0_b200.ymlenrollsunittest/_torch/attentionat directory level and B200 is SM100 — so the auto-generated QA "needs follow-up" note is resolved.
Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com>
|
Thanks again, @brnguyen2 ! Addressed the three follow-ups in |
|
PR_Github #64077 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64115 [ run ] triggered by Bot. Commit: |
|
PR_Github #64115 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64145 [ run ] triggered by Bot. Commit: |
|
PR_Github #64145 [ run ] completed with state |
5 participants
Summary
fmha_sm100package validation.Dev Engineer Review
require_msa_module()uses direct package imports with a consistent return type.fmha_sm100files are available in all supported wheel build paths.QA Engineer Review
test_msa_paged_kv_preserves_tma_compatible_outer_stride().test_msa_paged_hnd_input_materializes_unaligned_outer_stride().tests/integration/test_lists/were identified.test-db/orqa/coverage.Description
The MSA sparse-prefill path currently calls
.contiguous()on paged K/Vviews before FMHA. For a valid HND cache view, the tensor is already packed
within each page and satisfies the SM100 TMA requirements; only the outer page
stride is non-contiguous because the cache pool is coalesced. Materializing
the full tensor is unnecessary in that case.
This change adds a guarded zero-copy path to the pinned MSA dependency:
inner row/head strides, base address, and outer page alignment satisfy the
kernel contract;
layouts;
TRT-LLM-owned patch during population, and packages the patched
fmha_sm100module in the TensorRT-LLM wheel;availability checks, and precompiled-wheel extraction for the packaged
module.
Performance
The matched GB200 1P1D CTX qualification used 8K input tokens, one output
token, concurrency 64, and 1,000 requests. Only the MSA sparse-prefill path
changed.
_forward_stepNsight Systems also showed that the candidate removed 3,648 FP8 copy kernels
and 54.293 GiB of copied data per prefill.
Build design
MSA remains pinned at
e2ebe7656649f619af0ad1d457b534283034655e, but isnow declared in
3rdparty/fetch_content.jsoninstead of.gitmodules. CMakeapplies
msa_strided_paged_kv.patchwhile populating the source. The wheelbuilder then stages the patched
fmha_sm100package for setuptools; invalidor stale patch state fails the build instead of being silently ignored.
FetchContent downloads MSA from public GitLab during CMake configure, so
uncached builds require outbound GitLab access. The nested MSA CUTLASS submodule
is disabled; wheel staging reuses the existing TRT-LLM CUTLASS dependency.
While updating wheel staging, this also removes a duplicate
python -m buildinvocation. The remaining invocation preserves
EXTRA_WHEEL_BUILD_ARGS,including the platform argument already merged into that value.
Test Coverage
retaining an aligned outer-page stride, materializing an unsupported token
stride, and bit-identical output and LSE for strided versus packed paged K/V.
requests with no generation-side performance claim.
packaging fixes.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.