Repository navigation
Conversation
yuxianq
added
the
api-compatible
|
@yuxianq — vendor promotion needs author action: The pinned temporary source commit Fix: Fetch The bot will recheck automatically on a later poll. |
|
@yuxianq — the PrimTS upstream refresh #20155 has merged. Please update this source-update PR before merging it. The new canonical source branch is
The publishing bot now targets the new dated canonical branch. After this source-update PR merges, its lock-only promotion will target that branch. The old |
44788c6 to
a29bcb6
Compare
|
/bot run |
|
PR_Github #77487 [ run ] triggered by Bot. Commit: |
|
PR_Github #77487 [ run ] completed with state |
Vendor validation:
|
Vendor validation:
|
Vendor validation:
|
Vendor validation:
|
Vendor validation:
|
a29bcb6 to
50de5d1
Compare
|
/bot run |
Vendor validation:
|
|
PR_Github #81384 [ run ] triggered by Bot. Commit: |
|
PR_Github #81384 [ run ] completed with state
|
Vendor validation:
|
Vendor validation:
|
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
…rguments Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
540e862 to
7382d0f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #81702 [ run ] triggered by Bot. Commit: |
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Architecture and KV-cache support policy 3rdparty/vendor_sources.lock.yaml, tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py, tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py, tests/unittest/_torch/attention/test_prims_ts_fmha.py, tests/unittest/_torch/attention/test_cute_dsl_mla.py, tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py |
The vendor pin changes. PrimsTS availability adds SM107 for standard attention, while MLA remains limited to SM100 and SM103. Paged-cache validation permits supported FP8 E4M3 requests and checks their scales. Tests cover availability, cache modes, and page-size eligibility. |
FP8 preprocessing and wrapper execution tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py, tests/unittest/_torch/attention/test_prims_ts_fmha.py, tests/unittest/_torch/attention/test_prims_ts_attention_backend.py |
Context and generation wrappers pass FP8 flags and preprocessing scales, reinterpret processed queries as FP8, and use FP8 dtype for decode planning and workspace sizing. Tests cover wrapper inputs, scales, decode output, and graph replay. |
Backend integration coverage tests/unittest/_torch/attention/backend_case.py, tests/unittest/_torch/attention/test_prims_ts_attention_backend.py |
Backend test metadata now uses phase-specific prompt lengths. Integration tests add native and FP8 KV-cache cases and apply separate architecture gates for general attention and MLA. |
Priority:
Estimated code review effort: 4 (Complex) | ~45 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant PrimsTSFmha
participant FP8Preprocessing
participant PrimTSFMHA
PrimsTSFmha->>FP8Preprocessing: Preprocess queries and return BMM1 and BMM2 scales
PrimsTSFmha->>PrimTSFMHA: Submit FP8 query and preprocessing scales
PrimTSFMHA-->>PrimsTSFmha: Return attention output
Suggested reviewers: heyuhhh, yuanjingx87
Merge Risk: 
7382d
-
No actionable merge-blocking issue remains after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 27.27% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Title check The title names the main change: FP8 KV cache support for the non-MLA PrimsTS attention backend. The changeset matches this. It adds FP8 E4M3 KV-cache handling for standard PrimsTSFmha context, decode… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Description check The pull request includes the required Description, Test Coverage, and PR Checklist sections. It clearly explains the FP8 KV-cache changes, implementation scope, vendor updates, validation results, an… Full details: Title check
-
Explanation
The title names the main change: FP8 KV cache support for the non-MLA PrimsTS attention backend. The changeset matches this. It adds FP8 E4M3 KV-cache handling for standard PrimsTSFmha context, decode, and mixed batches, extends SM107 support for standard FMHA, and keeps MLA on SM100/SM103. The title follows the required [None][feat] format and is concise.
Full details: Description check
-
Explanation
The pull request includes the required Description, Test Coverage, and PR Checklist sections. It clearly explains the FP8 KV-cache changes, implementation scope, vendor updates, validation results, and known benchmark limitations. The checklist is present and marked as reviewed.
-
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 27.27% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 22 functions across 7 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
@farazkh80, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~ Assignment explanation
Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once. |
|
PR_Github #81702 [ run ] completed with state |
farazkh80
left a comment
farazkh80
left a comment
There was a problem hiding this comment.
Reviewed from attention perspective. One high item (H1, FP8 V-tail), the rest are questions.
| params.sequence_lengths, | ||
| params.batch_size, | ||
| ) | ||
| use_fp8 = QuantMode(attn.quant_mode).has_fp8_kv_cache() |
There was a problem hiding this comment.
zeroPagedVCacheTail only runs when T == TCache, so FP8 caches never get a zeroed V tail, but _get_or_plan_context_wrapper still sets paged_v_tail_is_zero=True and skips the kernel's own clear. A stale NaN byte in the tail would then leak through P·V.
Should that be False for FP8 caches, plus an fp8-kv case in test_prims_ts_context_receives_zeroed_v_tail?
| if mask_type not in (AttentionMaskType.causal, AttentionMaskType.padding): | ||
| return False, f"attention mask type {mask_type} is not supported." | ||
| policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta) | ||
| policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta, allow_fp8_kv_cache=True) |
There was a problem hiding this comment.
FP8 KV non-MLA layers used to fall through to flashinfer_trtllm_gen (fallback on SM107); with +prims_ts they now run PrimTS, and the description says no perf benchmarks were run.
Do you have decode latency numbers against trtllm-gen FP8 at small and large batch?
| ], | ||
| ) | ||
| @pytest.mark.parametrize("dtype", ["bfloat16", "float16"], ids=["bf16", "fp16"]) | ||
| @pytest.mark.parametrize("kv_dtype", [None, "fp8"], ids=["native-kv", "fp8-kv"]) |
There was a problem hiding this comment.
The FP8 backend cases use a 0.4–0.43 atol against an unquantized golden, loose enough to hide drift from E4M3 Q and P, and test_prims_ts_bfloat16 only covers BF16 KV.
Could we add an FP8-KV variant of that accuracy test to l0_b200?
| supports_per_tensor_fp8 = True | ||
|
|
||
| SUPPORTED_SMS = frozenset({100, 103}) | ||
| SUPPORTED_SMS = frozenset({100, 103, 107}) |
There was a problem hiding this comment.
PrimsTSBlockSparseFmha inherits SUPPORTED_SMS and isn't opt-in, so this also enables block-sparse PrimTS on SM107, moving VisualGen VSA off its dense fine stage onto kernels whose tests skip SM107.
Intended, or should block-sparse keep {100, 103}?
Dev Engineer Review
PrimsTSFmhaadds FP8 E4M3 KV-cache support for context, decode, and mixed batches. MLA remains limited to SM100 and SM103; standard attention adds SM107.QA Engineer Review
The changes add coverage for FP8 cache support, scale handling, graph replay, SM107 routing, and block-sparse rejection. Changed test functions are not detailed in the supplied evidence. Whether these tests appear in the CI or manual-QA lists could not be established. Coverage verdict: needs follow-up.
Per-File QA Perspective
3rdparty/vendor_sources.lock.yaml— Verify the pinned FlashInfer commit and tree digest match the intended vendor source. Source attribution and promotion status are not established by this file summary.tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py— Verify MLA page-size and KV-cache dtype selection usesSUPPORTED_MLA_SMS, so adding SM107 to standard attention does not change MLA routing.tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py— Verify FP8 E4M3 validation, preprocessing scales, decode scale modes, cache-plan identity, and device-scale graph replay. Check standard SM107 routing and confirm MLA remains restricted to SM100 and SM103.tests/unittest/_torch/attention/backend_case.py— Verify metadata uses context query lengths and generation cached-plus-new token counts for self-attention, and sequence lengths for cross-attention. Test-list status was not established.tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py— Covers rejection of FP8 KV cache for paged block-sparse generation. Test-list status was not established.tests/unittest/_torch/attention/test_cute_dsl_mla.py— Covers SM107 page-size exclusivity for 128-token BF16 pages. Test-list status was not established.tests/unittest/_torch/attention/test_prims_ts_attention_backend.py— Covers FP8 MLA preprocessing and graph replay, wrapper metadata and workspace behavior, and FP8 decode with device-resident scales across split-KV modes. Test-list status was not established.tests/unittest/_torch/attention/test_prims_ts_fmha.py— Covers standard-attention FP8 cache support across phases and cache managers, preprocessing scale forwarding, rejection paths, and dtype-dependent decode workspace sizing. Test-list status was not established.Description
Enable FP8 E4M3 KV cache for standard non-MLA
PrimsTSFmhacontext, decode, and mixed batches. Fused QKV preprocessing quantizes Q using the shared Q/K/V scale. The kernels consume FP8 Q/K/V and write native BF16/FP16 output, without an intermediate FP8 output or output dequantization pass.Decode reuses
bmm1_scaleandbmm2_scalefor host floats or float32 CUDA tensors. The plan selects the scale argument types before compilation, and_DecodeCompileSpecincludes the scale argument mode, so both in-memory and persistent caches distinguish host/device-scale plans. The standalone reducer only merges already-scaled partial states and takes no scale arguments. Tensor scales stay on the GPU and can change in place during CUDA graph replay. The TRT-LLM forward path usesvalidate=False; optional wrapper metadata validation can read sequence lengths/page IDs on the host, but never reads device scale values.Standard non-MLA PrimTS supports SM100, SM103, and SM107. The adapter relies on the pinned CUTLASS DSL 4.8.0 dependency and retains the active compiler capability check. MLA keeps its existing SM100/SM103 restriction. The shared KV policy helper retains
allow_fp8_kv_cache=Falsefor paged block-sparse attention, which does not support FP8 KV cache; standard PrimTS explicitly opts in.The paired upstream contribution is FlashInfer #8021, ported to current upstream APIs. The vendor lock pins the FlashInfer scale API changes on top of
trtllm-prims-ts-dev-20261012ataa709061230a, retaining the TensorRT-LLM compatibility patch and original headers on existing FlashInfer files.Test Coverage
Rebased
7382d0fb5920onto78145f88ea03:BuildInfo.check_cuda_architecture_supported; build 67209 lackskv_cache_manager_v2.ReusePolicy). The requested/bot run --disable-fail-fastCI run builds the matching runtime.B300 validation of the compile-spec/reducer cleanup at
540e86299e7d:B300 validation of the review fixes at
27c01417838f, with PyTorch 26.09 and CUTLASS DSL 4.8.0:Tensor.item()andTensor.__float__()forbidden. CUDA graph capture/replay passed with those scalar-read guards.Earlier Rubin validation at
ffeb74563528on Hecate SM107, ARM64, PyTorch 26.09, CUDA 13.4, and CUTLASS DSL 4.8.0 (job904798, exit 0): 294 tests passed (227 adapter checks and 67 non-MLA backend cases), with zero failures/errors/skips. A separate persistent-cache reload check also passed.Model-quality and performance benchmarks have not been run.
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.