[None][perf] Avoid paged MSA K/V materialization during prefill

[peihu-nv]

@coderabbitai summary

Description

The MSA sparse-prefill path currently calls .contiguous() on paged K/V
views 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:

• preserves a four-dimensional HND paged-cache view when its page size,
  inner row/head strides, base address, and outer page alignment satisfy the
  kernel contract;
• falls back to the existing contiguous materialization for unsupported
  layouts;
• stages a disposable copy of the pinned 3rdparty/MSA source during wheel
  construction, applies the downstream patch there, and packages that staged
  source through TRTLLM_MSA_PACKAGE_DIR;
• leaves the checked-out MSA submodule clean.

This draft intentionally scopes activation to TensorRT-LLM's staged wheel
build. Source/editable-import activation can be reviewed independently.

Performance

The original matched GB300 CTX-only qualification used 8K prompts:

Metric	Baseline	Candidate	Change
Median prefill latency	660.792 ms	532.481 ms	-19.42%
Request throughput	1.0000x	1.2557x	+25.57%

Nsight Systems also showed that the candidate removed 3,648 FP8 copy kernels
and 54.293 GiB of copied data per prefill.

Build design

The MSA submodule remains pinned at e2ebe7656649f619af0ad1d457b534283034655e.
The wheel builder copies it into a disposable build directory, applies the
TRT-owned patch, and packages that staged source. This follows the downstream
patch model used for DeepEP while keeping the checked-out submodule clean.

Test Coverage

• Added staged-wheel build tests in test_build_wheel.py.
• The MSA patch includes coverage for the aligned zero-copy path, fallback
  behavior, and strided-versus-packed numerical equivalence.
• Original staging tests: 2/2 passed.
• Current-main changed-file pre-commit checks pass.
• Current-main staged-wheel GPU qualification remains pending; the PR is
  therefore opened as a draft.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

[peihu-nv]

/bot run

[tensorrt-cicd]

PR_Github #61667 [ run ] triggered by Bot. Commit: 334db27 Link to invocation

[tensorrt-cicd]

PR_Github #61667 [ run ] completed with state SUCCESS. Commit: 334db27
/LLM/main/L0_MergeRequest_PR pipeline #49873 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation