[https://nvbugs/6791733][fix] Guard PrimTS ragged Q tensor maps on Blackwell by yuxianq · Pull Request #20298 · NVIDIA/TensorRT-LLM

yuxianq · 2026-10-14T15:14:38Z

Dev Engineer Review

The vendor lock updates the flashinfer-prims-ts source revision and tree digest. Source and destination paths, include pattern, patch, and patch digest remain unchanged.

The test-list changes isolate PrimTS kernel JIT time from the broad attention-directory timeout. No unrelated implementation files changed.

QA Engineer Review

test_prims_ts_attention_backend.py adds coverage for packed-query tails, compact-query boundaries near 128 KiB, storage offsets, causal and sliding-window modes, paged KV, and CUDA graph replay. It also adjusts FP8 MLA coverage and removes an obsolete FP16 dense case.

The file is a separate CI entry in both l0_b200.yml and l0_b300.yml. The module contains 16 test functions, including the new packed-query and compact-query boundary tests. Reported hardware validation passed, but supplied CI results are mixed. Coverage verdict: needs follow-up.

Per-File QA Perspective

3rdparty/vendor_sources.lock.yaml — Updates the pinned FlashInfer PrimTS snapshot. Verify reproducible resolution and the Blackwell descriptor fix.
tests/unittest/_torch/attention/test_prims_ts_attention_backend.py — Covers packed-query tails, boundary sizes, storage offsets, masks, paging, dtypes, and CUDA graph replay. The module is listed separately in the B200 and B300 CI test-db files.
tests/integration/test_lists/test-db/l0_b200.yml — Excludes the PrimTS test from the broad attention entry and adds a separate entry. This separates kernel JIT time from the directory timeout.
tests/integration/test_lists/test-db/l0_b300.yml — Applies the same separate-entry treatment for B300 CI. This changes test selection and timeout accounting.

Description

PrimTS context attention can fail with CUDA error: an illegal memory access was encountered on Blackwell when packed (variable-length) Q is small and sits near the end of its allocation. The packed-Q TMA descriptor uses a synthetic ragged layout whose declared extent (2^31 rows) is far larger than the real tensor. For such descriptors the driver sets descriptor bit 85, which lets TMA access up to 128 KiB from the base address and faults when less than 128 KiB is mapped after Q. NVIDIA documents this in the 615.71.09 known issues.

Backport the fix that FlashInfer #6482 merged into FlashInfer main (f57c58b8). On SM100/SM103 execution devices, clear bit 85 on the host after encoding the ragged-Q descriptor, unless the real Q tensor is a compact layout of at least 128 KiB. Ragged Q loads, coordinates, and device code are unchanged; large compact Q and fixed-layout Q keep their original descriptors.

The FlashInfer source pin moves from the PrimTS vendor refresh ae312f7cc172 (#20301) to ae3851ab5f95. Its changes to context.py and fmha_kernel.py are line-for-line the PrimTS source changes in #6482. The existing compatibility patch is unchanged.

This PR briefly switched to native-Q TMA loads (1174214b). It now uses the guard again so release/1.3 matches FlashInfer main.

Vendor Promotion

schema_version: 1
vendors:
  flashinfer-prims-ts:
    upstream_prs:
      - https://github.com/flashinfer-ai/flashinfer/pull/6482

Test Coverage

tests/unittest/_torch/attention/test_prims_ts_attention_backend.py adds 30 cases (67 in the file), each checked against an independent FP32 reference in eager execution, CUDA graph replay, and replay after changing Q (graph outputs are poisoned with NaN before replay):

test_prims_ts_context_packed_query_tails: packed query tails with contiguous and paged KV, causal and head-paired sliding-window masks.
test_prims_ts_context_compact_query_boundary: BF16/FP8 E4M3, D128/D256, contiguous/paged KV, with compact Q just below, at, and above 128 KiB (24 cases). Each dtype/D/layout group is planned with shared bounds, so one compiled kernel handles all three sizes and must choose whether to correct the descriptor at run time.
A contiguous Q view with a nonzero storage offset, and a fixed-layout control.

CI time limit

CI #78475 on the earlier guard head 346ced0 failed on DGX_B200-PyTorch-4: the unittest/_torch/attention entry ran into the one-hour pytest timeout (3581 s) while test_prims_ts_context_compact_query_boundary[fp8-d128-at-contiguous] was running. That entry runs the whole directory as one pytest item. Without this PR it already takes about 2900–3200 s on B200, and another PR's run also timed out there (pipeline 304). With these cases, the B300 entry took 3503 s. The added cases cost about 25 s each, almost all CuTe DSL compilation.

So test_prims_ts_attention_backend.py now has its own entry in l0_b200.yml and l0_b300.yml and is ignored by the directory entry, following the existing test_cute_dsl_gvr_topk_decode.py split. The shared plan bounds above also cut the boundary test from 24 compilations to 8.

Validation

On head e957dcce, the whole file (67 cases) passed in two ordered sessions on B200 with UMD/KMD 615.71.09, the affected driver (134 executions, 762 s). This includes test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata, which failed on the unpatched base. Starting from an empty kernel cache, the 24 boundary cases took 195 s, about 8 s each, compared with about 25 s each in CI #78475.

Earlier, the guard head 346ced0c passed all 67 cases in two ordered sessions on B200 with UMD/KMD 615.71.09 (134 executions). Before the vendor-refresh rebase, the same guard passed 872 executions on B200 and B300 (the original 41 cases repeated in 10 ordered sessions, plus 26 boundary/control cases per GPU).

Tests imported the candidate Python source with retained release CI native libraries (build 10743). The C++ files that release/1.3 changed since that build (MPI progress thread, EPLB allocator, SM107 FMHA cubins) are not exercised by these tests, and this Python-only change needs no native rebuild. Vendor reconstruction passed offline, against the local FlashInfer source, and against the published pin with access required.

Reproduction

On the unpatched release/1.3 base d661fdbf, running the existing test file in file order with session repeats reproduced the CI failure at the 68th execution every time, on two B200 nodes and one B300 node: CUDA error: an illegal memory access was encountered in test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata (second session), with UMD/KMD 615.71.09. No failing driver-version baseline or first fixed driver build is claimed.

Performance

Measured before the rebase against this PR's earlier native-Q revision. Values are warmed CUDA-graph wrapper latencies per call (median of two paired rounds with reversed order), not end-to-end serving throughput.

GPU	Case	Native Q µs	Guard µs	Change
B200	Small packed BF16 D256	8.250	8.046	-2.47%
B200	Paged BF16 D256, many short requests	45.900	46.414	+1.12%
B200	Packed FP8 E4M3 D256	7.578	7.634	+0.74%
B200	Fixed BF16 D128 control	9.010	9.010	approximately 0%
B300	Small packed BF16 D256	7.612	7.437	-2.31%
B300	Paged BF16 D256, many short requests	42.797	43.100	+0.71%
B300	Packed FP8 E4M3 D256	7.054	7.056	+0.04%
B300	Fixed BF16 D128 control	8.072	8.072	-0.01%

All cases use Hq/Hkv=8/2, causal masking and BF16 output. Small-Q uses Q lengths [3,27,1], KV [33,59,17]; paged uses Q [1,3,17,65] and KV [33,33,65,129], each repeated 32 times with 32-token pages; FP8 uses two Q=KV=128 requests; fixed uses two Q=3/KV=33 requests. The guard is host code; relative to the unpatched release source, device code is unchanged.

PR Checklist

Description explains the final change and its scope.
Source and tests follow repository guidance; upstream licenses are preserved.
Relevant regression tests are included.
Public wrapper APIs and dependency pins remain unchanged.
Vendor source and compatibility-patch integrity are verified.
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.

yuxianq · 2026-10-15T07:36:47Z

/bot run

tensorrt-cicd · 2026-10-15T07:44:22Z

PR_Github #78053 [ run ] triggered by Bot. Commit: 30950da Link to invocation

tensorrt-cicd · 2026-10-15T11:10:11Z

PR_Github #78053 [ run ] completed with state FAILURE. Commit: 30950da
/LLM/release-1.3/L0_MergeRequest_PR pipeline #247 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

Pin FlashInfer to the native packed-Q tensor-map fix while preserving the existing TensorRT-LLM compatibility patch. Keep ragged output stores to suppress rows outside each request. Add independent-reference coverage for packed query tails in contiguous and paged attention, including head-paired sliding windows and CUDA graph replay with updated query values. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

…ackwell Replace the native-Q workaround with the bounded FlashInfer correction on the existing release source snapshot. Preserve ragged Q loads and apply the post-encoding correction only on SM100/SM103 when a compact 128 KiB prefix cannot be established. Forward the selected device through all compilation paths and retain the existing vendor compatibility patch. Add independent-reference coverage for descriptor-size boundaries, storage offsets, and fixed-layout control with changed-input CUDA graph replay. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq · 2026-10-16T11:53:22Z

/bot run --disable-fail-fast

yuxianq · 2026-10-16T11:59:02Z

@PerkzZheng, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~

Assignment explanation

Components	kernels/prims-ts (60%); fmha/prims-ts (15%); integration/attention-tests (20%); integration/dependencies (5%)
Tags	algorithm:dense (100%); representation:gqa (100%); phase:prefill (100%); features:cuda-graphs (50%); features:paged-kv (30%); features:quantization (20%)
PR effort	8 points — Despite its small implementation diff, review requires deep reasoning about descriptor-bit semantics, the compact-prefix proof, static versus runtime shape checks, and graph replay behavior. The test matrix adds meaningful review work without constituting a broad kernel rewrite.
Importance	Important — The kernel setup directly modifies ragged-Q TMA address-generation metadata. Correctness depends on the 128 KiB compact-prefix condition and SM100/SM103 gating, making this a sensitive GPU memory-addressing change.
Affinity (0–1)	0.699 component/tag match; 0.524 after PIC preference
Outstanding points at selection	32 → 40 (+8)
Points already reserved on this PR	0
Effective weight	0.5
Completed review points (preceding 168 hours)	17
Combined points after assignment	57
Demand-derived target	49.43 points (soft share)
Additional fair share before assignment	0.43 points
Share explanation	Whole PR / qualification or coverage constraint exceeds a soft share; assignment is allowed.
Selection rule	affinity with 0.3-share guard
Assignment scope (PerkzZheng)	all

Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once.

coderabbitai · 2026-10-16T12:00:09Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 11369378-ae81-4e2a-a277-93476bed30ec

📥 Commits

Reviewing files that changed from the base of the PR and between e957dcc and a7d6732.

📒 Files selected for processing: 1

tests/unittest/_torch/attention/test_prims_ts_attention_backend.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

The flashinfer-prims-ts vendor pin changes. Attention tests expand packed-query and compact-query coverage with causal references and CUDA graph replay. B200 and B300 pre-merge lists register the attention test separately.

Changes

PrimsTS context attention

Layer / File(s)	Summary
Vendor source pin `3rdparty/vendor_sources.lock.yaml`	The vendor entry changes its branch, commit, and tree digest. Its source, destination, include pattern, patch, and patch digest remain unchanged.
Context attention replay coverage `tests/unittest/_torch/attention/test_prims_ts_attention_backend.py`	Tests restrict execution to `TRTLLM`, adjust existing Qwen2 and FP8 MLA coverage, and add reference comparisons for packed and compact queries. Coverage includes contiguous and paged inputs, causal and sliding-window masks, BF16 and FP8 inputs, query-size boundaries, nonzero storage offsets, and CUDA graph replay.
Pre-merge test registration `tests/integration/test_lists/test-db/l0_b200.yml`, `tests/integration/test_lists/test-db/l0_b300.yml`	Both pre-merge lists exclude the PrimsTS attention test from the attention-directory entry and register it separately. The B200 entry has no explicit timeout.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to a7d67

The attention tests remain registered on both B200 and B300. No concrete merge-blocking risk is established.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 16 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the PrimTS ragged-Q tensor-map fix on Blackwell and uses the required NVBugs and fix format.
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 description is complete and follows the required structure. It explains the issue and solution, documents vendor promotion, lists detailed test coverage and validation, describes perf…

Full details: Description check

Explanation

The pull request description is complete and follows the required structure. It explains the issue and solution, documents vendor promotion, lists detailed test coverage and validation, describes performance impact, and includes the PR checklist. The PR title is not included in the provided context, so its template compliance cannot be verified.

Fix all pre-merge checks with AI

✨ 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.}

tensorrt-cicd · 2026-10-16T12:00:33Z

PR_Github #78475 [ run ] triggered by Bot. Commit: 346ced0 Link to invocation

tensorrt-cicd · 2026-10-16T15:29:02Z

PR_Github #78475 [ run ] completed with state SUCCESS. Commit: 346ced0
/LLM/release-1.3/L0_MergeRequest_PR pipeline #303 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

github-actions · 2026-10-19T06:00:47Z

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.

## 📌 Description Repeated Prims-TS context launches can fail in the packed-Q TMA load on Blackwell. Preserve the original five-dimensional ragged Q descriptor, loads, and request-local masking, and apply NVIDIA's documented OOB address-generation correction only when needed by the real Q layout. The context builder gates the correction on the tensor's actual device capability, SM100 or SM103. After descriptor encoding, clear descriptor bit 85 when a contiguous 128 KiB prefix cannot be established from the real compact Q tensor. For compact `(T,H,D)` inputs, compare runtime `T` against `ceil(131072 / (H*D*element_bytes))`; unknown/noncompact layouts conservatively receive the correction. This avoids runtime byte-count multiplication overflow and uses actual D192 FP8 storage rather than its padded staging width. The SM100/SM103 gate describes this backend's scope, not an exhaustive list of affected GPUs. Prims-TS context accepts SM100, SM103, and SM107; its SM107 support can use a `sm_100f` compiler fallback, so the compiler target alone must not enable a Blackwell-specific descriptor edit. SM110/SM120/SM121 are rejected by this context path. NVIDIA's driver notice says Blackwell generally and does not provide a per-SM affected/unaffected matrix; this PR does not establish whether those other Blackwell devices encounter the same issue. The descriptor edit must not be generalized to other architectures without confirming their layout and applicability. Large compact Q tensors retain the original descriptor. Fixed-Q, K/V/O descriptors, GPU scheduling, and public APIs retain their existing behavior. The correction adds no GPU kernel, workspace, or host-device synchronization. The driver issue is consistent with allocation-placement-dependent failures involving overlapping tensor descriptions. Synthetic descriptor dimensions cannot establish the actual mapped prefix. The targeted correction strongly supports this diagnosis; it does not establish the driver's internal address-translation sequence. Retain the four packed-tail regression tests covering contiguous/paged KV, causal/sliding-window attention, and replay after changing Q. ## 🔍 Related Issues - [NVIDIA driver 615.71.09 known issues](https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-615-71-09/index.html#known-issues): Blackwell tensor-map failures and the descriptor correction. - [CUTLASS contiguous-prefix criterion](https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/copy_sm90_desc.hpp#L443). - [Related TensorRT-LLM draft PR](NVIDIA/TensorRT-LLM#20298), which currently contains the earlier native-Q approach. ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. All applicable hooks passed for the affected files using `pre-commit run --files ...`, including mypy, Ruff lint, and formatting. Repository-wide hooks were not run. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). GPU validation used the guarded-candidate source in `a9883a5e8ada0d25754ecb31ae7eb83109c403cc`, based on main `2895577720741ddfedff118202c89584996b8875`. The subsequent copyright-header cleanup changes no executable Python code (AST equality checked); applicable pre-commit hooks pass. GPU tests were not repeated for that comment-only cleanup: - **B200: 80 ordered graph/tail tests passed**: four packed-tail cases followed by four BF16 packed/paged D128/D256 live-offset graph cases, repeated in session order ten times. - **B200: five supplemental boundary/view oracle cases passed**, each checking eager execution, captured replay, and replay after changing Q: BF16 H8/D128 at 63 and 64 rows; FP8 H8/QK192/V128 at 85 and 86 rows; and a contiguous 64-row BF16 view with a nonzero storage offset. - **The guarded candidate passed 54 independent FP32-reference checks in 18 benchmark processes** across B200/B300, covering eager execution, graph replay, and changed-Q replay. Across all comparison variants (baseline, native-Q, unconditional correction, and guarded correction), 40 processes passed 120 checks. Source and harness hashes matched the frozen validation snapshots. - Runtime: PyTorch `2.14.0a0+b2c75dd062.nv26.09`, CUDA 13.4, CUTLASS DSL `4.8.0.dev0`, reported driver `615.71.09`. Reproduce the ordered selection in one process (requires pytest-repeat): ```bash python -m pytest -p no:xdist -v \ tests/attention/test_attention_ts_context.py::test_attention_ts_context_packed_query_tails \ tests/attention/test_attention_ts_context.py::test_attention_ts_context_live_q_offsets_expand_causal_domain_on_graph_replay \ -k 'packed_query_tails or bf16' --count=10 --repeat-scope=session ``` The full 322-case context suite and the full FlashInfer suite have not been run on this guarded candidate. ### Performance Paired measurements use the original main source as baseline, two fresh-process rounds with reversed variant order, 50 eager and 10 graph warmups, and 31 CUDA-event samples per process. Each timed graph contains 256 wrapper calls with preallocated outputs. Every paired process-median delta across five B200 and two B300 configurations was within **±0.02%**. All samples were retained, including one B300 baseline timing outlier; this is a median-latency result, not a tail-latency guarantee. All configurations below use Hq/Hkv=8/2 and BF16 output. Times are medians of the two process medians. | GPU | Configuration | Baseline µs | Guarded correction µs | | --- | --- | ---: | ---: | | B200 | Paged BF16 D256, many short requests | 61.710 | 61.710 | | B200 | Packed BF16 D256, many short requests | 49.754 | 49.758 | | B200 | Packed BF16 D128, tails | 9.361 | 9.361 | | B200 | Fixed BF16 D256 | 47.318 | 47.316 | | B200 | Packed FP8 D256, long aligned requests | 178.111 | 178.097 | | B300 | Paged BF16 D256, many short requests | 59.017 | 59.010 | | B300 | Fixed BF16 D256 | 42.226 | 42.232 | Problem sizes: many-short Q=`(1,3,17,65)` ×32, KV=`(33,33,65,129)` ×32, causal window 31, with nonidentity 32-token pages for paged KV; tails Q=`(3,130,1)`, KV=`(33,160,17)`, causal; fixed 256 requests with Q=1/KV=33, causal; long FP8 two Q=KV=4096 requests, dense. For small packed causal inputs Q=`(3,2)`, KV=`(33,64)`, B200 measured **9.199 → 9.193 µs** for BF16 D128 and **9.892 → 9.809 µs** for FP8 QK192/V128, comparing the earlier native-Q fix with this guarded correction. No uncorrected small-Q baseline was run. These measurements cover warmed CUDA-graph GPU execution. Planning, JIT, capture, host overhead, cold-cache behavior, and end-to-end serving throughput are outside the measurement scope. ## 🔬 Experimental Track  - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.  ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` ## Reviewer Notes Please review the real-Q prefix proof, actual-device architecture gate, and host descriptor update. Generated host IR confirms the runtime row-count guard modifies the descriptor later passed to launch. The writable host `TensorMap` alias works in the tested DSL runtime; cross-version support or an encoder-side helper needs DSL-owner review. This revision replaces the earlier native-Q implementation, which passed broad correctness checks but increased measured graph-operation latency by up to 3.43% in a separate B300 comparison. A separately frozen unconditional descriptor-correction port passed the original 37-case TensorRT-LLM sequence ten times (370 passes), supporting the diagnosis; those results are not validation of this exact guarded source. Keep the workaround scoped to documented Blackwell descriptor layouts and revisit it when a driver fix or supported encoder helper is available.  ## Summary by CodeRabbit * **Bug Fixes** * Improved attention correctness on supported Blackwell GPUs when processing partial query tiles in contiguous and paged attention. * Verified correct results with causal masking, both unwindowed and sliding-window attention, and captured graph replay after query data changes.  --------- Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

…me limits Run test_prims_ts_attention_backend.py as its own B200 and B300 test-list entry. The attention-directory entry already used most of its one-hour pytest timeout, and the added context cases pushed B200 over it. Plan each compact-query boundary group with shared bounds so the below, at, and above sizes reuse one compiled kernel instead of compiling three. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq · 2026-10-19T10:37:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-19T10:44:30Z

PR_Github #78900 [ run ] triggered by Bot. Commit: e957dcc Link to invocation

tensorrt-cicd · 2026-10-19T10:47:13Z

PR_Github #78842 [ run ] completed with state ABORTED. Commit: 1174214

Link to invocation

yuxianq · 2026-10-19T10:49:54Z

Vendor promotion metadata and source attribution are valid. This is not a code-review approval. Promotion still requires the source PR to merge.

Compare PrimTS only against the Vanilla golden; the FlashInfer backend has its own tests and only adds JIT builds here. Reuse compiled kernels where cases differ only in bounds: the mixed Qwen2 case shares the context kernel, and compact-query boundary cases plan with the packed-tail bounds. Drop the FP16 dense case, which repeats a step of the graph-replay test. Cover each FP8 MLA Q tile (M8, M16, repeated M64) once across the two FP8 MLA tests, and keep boundary groups that cover both dtypes, head sizes, and KV layouts. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

brnguyen2

Approving — the comments below are optional touch-ups, not blockers.

This reads as a faithful port of FlashInfer #6482. Clearing word1 bit 21 is the right bit (85), and the clear is const_expr-gated to SM100/SM103, so other architectures compile the original path. The remaining items are mostly about the PR record:

The auto-generated description header contradicts the diff. The CodeRabbit section says "No implementation source file changed" and describes native-Q loading. This PR changes context.py and fmha_kernel.py and uses the bit-85 guard instead. Please regenerate or delete that block so a future reader doesn't take the wrong mechanism from it.
Main counterpart. https://nvbugs/6791733 reproduces on main post-merge, but this PR targets release/1.3 only. Please link the main PR, or the vendor refresh on main that picks up FlashInfer f57c58b8, so the fix isn't release-only.
Vendor pin durability. The new pin points to the topic branch bug-6791733-prims-ts-cuda-graph--source-q-tma-release13 on a personal fork. If that branch is deleted, ae3851ab can become unreachable and the lock can't be re-resolved. Please push the commit to a long-lived branch (like the dated trtllm-prims-ts-dev-* branches) or tag it.
Tests don't check which path ran. The boundary tests (below/at/above 128 KiB) only check numerical correctness. If _compact_q_prefix_min_rows always returned None (for example, because the fake compact tensor's strides arrive as dynamic values instead of Python ints), every case would still pass, and large compact Q would silently lose the original descriptor. A small unit test of _compact_q_prefix_min_rows against the fake Q tensor the context compiler builds would catch that.

tensorrt-cicd · 2026-10-19T16:24:50Z

PR_Github #78900 [ run ] completed with state SUCCESS. Commit: e957dcc
/LLM/release-1.3/L0_MergeRequest_PR pipeline #348 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

yuxianq · 2026-10-20T04:02:00Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-20T04:10:29Z

PR_Github #79134 [ run ] triggered by Bot. Commit: a7d6732 Link to invocation

tensorrt-cicd · 2026-10-20T08:09:18Z

PR_Github #79134 [ run ] completed with state FAILURE. Commit: a7d6732
/LLM/release-1.3/L0_MergeRequest_PR pipeline #387 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

yuxianq · 2026-10-20T08:14:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-20T08:23:06Z

PR_Github #79211 [ run ] triggered by Bot. Commit: a7d6732 Link to invocation

tensorrt-cicd · 2026-10-20T11:38:47Z

PR_Github #79211 [ run ] completed with state FAILURE. Commit: a7d6732
/LLM/release-1.3/L0_MergeRequest_PR pipeline #399 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

yuxianq · 2026-10-21T03:36:32Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-21T03:43:14Z

PR_Github #79501 [ run ] triggered by Bot. Commit: a7d6732 Link to invocation

tensorrt-cicd · 2026-10-21T05:17:58Z

PR_Github #79501 [ run ] completed with state SUCCESS. Commit: a7d6732
/LLM/release-1.3/L0_MergeRequest_PR pipeline #442 completed with status: 'SUCCESS'

CI Report

Link to invocation

----- BEGIN VENDOR PROMOTION V1 ----- { "attribution": { "6b00497398f491136be91fb0e60ac3f314da53ba": { "method": "same-edits", "upstream_base": "2895577720741ddfedff118202c89584996b8875", "upstream_head": "14770b419ec3261c2afc55e401051c92bb7f3fd9" }, "ae3851ab5f958ba22f154059df5b38321f6369d2": { "method": "source-series-aggregate", "upstream_base": "2895577720741ddfedff118202c89584996b8875", "upstream_head": "14770b419ec3261c2afc55e401051c92bb7f3fd9" } }, "attribution_scope": { "include": [ "**/*.py" ], "source": "flashinfer/attention/prims_ts" }, "canonical_branch": "trtllm-prims-ts-dev-20261012", "canonical_url": "https://github.com/yuxianq/flashinfer.git", "changes": [ { "commits": [ "6b00497398f491136be91fb0e60ac3f314da53ba", "ae3851ab5f958ba22f154059df5b38321f6369d2" ], "upstream_base": "2895577720741ddfedff118202c89584996b8875", "upstream_head": "14770b419ec3261c2afc55e401051c92bb7f3fd9", "upstream_merge": "f57c58b898cc0f96ddea03066a3fd66ad201337a", "upstream_pr": "flashinfer-ai/flashinfer#6482" } ], "consumer_base": "release/1.3", "consumer_merge": "b053431a1a62303d5c3b7a3d9d5af7f82c7c76d0", "consumer_pr": "#20298", "contributor_metadata": { "metadata_digest": "95942a2db5f1373b658ac1a2010629f282614a3ca905b68fc64f054d39c8129e", "pr_author": "yuxianq", "source_head": "a7d6732b1e451e67e1d5396251395f6a742a7e4f" }, "kind": "promotion", "previous": "ae312f7cc1723b06fda5168105ade0f328ee47d8", "promoted": "ae3851ab5f958ba22f154059df5b38321f6369d2", "source_branch": "bug-6791733-prims-ts-cuda-graph--source-q-tma-release13", "source_url": "https://github.com/yuxianq/flashinfer.git", "upstream_base": "ae312f7cc1723b06fda5168105ade0f328ee47d8", "upstream_branch": "main", "upstream_observed_main": "a3f2525e39dffc98dcad4ce50706641a580358f0", "upstream_repository": "flashinfer-ai/flashinfer", "vendor": "flashinfer-prims-ts", "version": 1 } ----- END VENDOR PROMOTION V1 ----- Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

…fer-ai#6482)  ## 📌 Description Repeated Prims-TS context launches can fail in the packed-Q TMA load on Blackwell. Preserve the original five-dimensional ragged Q descriptor, loads, and request-local masking, and apply NVIDIA's documented OOB address-generation correction only when needed by the real Q layout. The context builder gates the correction on the tensor's actual device capability, SM100 or SM103. After descriptor encoding, clear descriptor bit 85 when a contiguous 128 KiB prefix cannot be established from the real compact Q tensor. For compact `(T,H,D)` inputs, compare runtime `T` against `ceil(131072 / (H*D*element_bytes))`; unknown/noncompact layouts conservatively receive the correction. This avoids runtime byte-count multiplication overflow and uses actual D192 FP8 storage rather than its padded staging width. The SM100/SM103 gate describes this backend's scope, not an exhaustive list of affected GPUs. Prims-TS context accepts SM100, SM103, and SM107; its SM107 support can use a `sm_100f` compiler fallback, so the compiler target alone must not enable a Blackwell-specific descriptor edit. SM110/SM120/SM121 are rejected by this context path. NVIDIA's driver notice says Blackwell generally and does not provide a per-SM affected/unaffected matrix; this PR does not establish whether those other Blackwell devices encounter the same issue. The descriptor edit must not be generalized to other architectures without confirming their layout and applicability. Large compact Q tensors retain the original descriptor. Fixed-Q, K/V/O descriptors, GPU scheduling, and public APIs retain their existing behavior. The correction adds no GPU kernel, workspace, or host-device synchronization. The driver issue is consistent with allocation-placement-dependent failures involving overlapping tensor descriptions. Synthetic descriptor dimensions cannot establish the actual mapped prefix. The targeted correction strongly supports this diagnosis; it does not establish the driver's internal address-translation sequence. Retain the four packed-tail regression tests covering contiguous/paged KV, causal/sliding-window attention, and replay after changing Q. ## 🔍 Related Issues - [NVIDIA driver 615.71.09 known issues](https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-615-71-09/index.html#known-issues): Blackwell tensor-map failures and the descriptor correction. - [CUTLASS contiguous-prefix criterion](https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/copy_sm90_desc.hpp#L443). - [Related TensorRT-LLM draft PR](NVIDIA/TensorRT-LLM#20298), which currently contains the earlier native-Q approach. ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. All applicable hooks passed for the affected files using `pre-commit run --files ...`, including mypy, Ruff lint, and formatting. Repository-wide hooks were not run. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). GPU validation used the guarded-candidate source in `a9883a5e8ada0d25754ecb31ae7eb83109c403cc`, based on main `2895577720741ddfedff118202c89584996b8875`. The subsequent copyright-header cleanup changes no executable Python code (AST equality checked); applicable pre-commit hooks pass. GPU tests were not repeated for that comment-only cleanup: - **B200: 80 ordered graph/tail tests passed**: four packed-tail cases followed by four BF16 packed/paged D128/D256 live-offset graph cases, repeated in session order ten times. - **B200: five supplemental boundary/view oracle cases passed**, each checking eager execution, captured replay, and replay after changing Q: BF16 H8/D128 at 63 and 64 rows; FP8 H8/QK192/V128 at 85 and 86 rows; and a contiguous 64-row BF16 view with a nonzero storage offset. - **The guarded candidate passed 54 independent FP32-reference checks in 18 benchmark processes** across B200/B300, covering eager execution, graph replay, and changed-Q replay. Across all comparison variants (baseline, native-Q, unconditional correction, and guarded correction), 40 processes passed 120 checks. Source and harness hashes matched the frozen validation snapshots. - Runtime: PyTorch `2.14.0a0+b2c75dd062.nv26.09`, CUDA 13.4, CUTLASS DSL `4.8.0.dev0`, reported driver `615.71.09`. Reproduce the ordered selection in one process (requires pytest-repeat): ```bash python -m pytest -p no:xdist -v \ tests/attention/test_attention_ts_context.py::test_attention_ts_context_packed_query_tails \ tests/attention/test_attention_ts_context.py::test_attention_ts_context_live_q_offsets_expand_causal_domain_on_graph_replay \ -k 'packed_query_tails or bf16' --count=10 --repeat-scope=session ``` The full 322-case context suite and the full FlashInfer suite have not been run on this guarded candidate. ### Performance Paired measurements use the original main source as baseline, two fresh-process rounds with reversed variant order, 50 eager and 10 graph warmups, and 31 CUDA-event samples per process. Each timed graph contains 256 wrapper calls with preallocated outputs. Every paired process-median delta across five B200 and two B300 configurations was within **±0.02%**. All samples were retained, including one B300 baseline timing outlier; this is a median-latency result, not a tail-latency guarantee. All configurations below use Hq/Hkv=8/2 and BF16 output. Times are medians of the two process medians. | GPU | Configuration | Baseline µs | Guarded correction µs | | --- | --- | ---: | ---: | | B200 | Paged BF16 D256, many short requests | 61.710 | 61.710 | | B200 | Packed BF16 D256, many short requests | 49.754 | 49.758 | | B200 | Packed BF16 D128, tails | 9.361 | 9.361 | | B200 | Fixed BF16 D256 | 47.318 | 47.316 | | B200 | Packed FP8 D256, long aligned requests | 178.111 | 178.097 | | B300 | Paged BF16 D256, many short requests | 59.017 | 59.010 | | B300 | Fixed BF16 D256 | 42.226 | 42.232 | Problem sizes: many-short Q=`(1,3,17,65)` ×32, KV=`(33,33,65,129)` ×32, causal window 31, with nonidentity 32-token pages for paged KV; tails Q=`(3,130,1)`, KV=`(33,160,17)`, causal; fixed 256 requests with Q=1/KV=33, causal; long FP8 two Q=KV=4096 requests, dense. For small packed causal inputs Q=`(3,2)`, KV=`(33,64)`, B200 measured **9.199 → 9.193 µs** for BF16 D128 and **9.892 → 9.809 µs** for FP8 QK192/V128, comparing the earlier native-Q fix with this guarded correction. No uncorrected small-Q baseline was run. These measurements cover warmed CUDA-graph GPU execution. Planning, JIT, capture, host overhead, cold-cache behavior, and end-to-end serving throughput are outside the measurement scope. ## 🔬 Experimental Track  - [ ] This PR is **experimental**: it adds or changes code under `flashinfer/experimental/` and/or an `@flashinfer_experimental_api`. Tracking issue: # - [ ] The tracking issue names an owner, the reason for the experimental path, and a graduation plan with a target release. - [ ] Core changes are limited to a thin entry point (signature, shared validation, feature-gate check, backend selection, handoff). - [ ] Tests live in `tests/experimental/` and were validated on the intended hardware; a runnable example is included. - [ ] Nothing is registered in `flashinfer/aot.py`, and no experimental backend is reachable from `backend="auto"` without `FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1`. (Calling an `@flashinfer_experimental_api` or naming a backend explicitly is itself the opt-in and needs no environment variable.) - [ ] **Test scope declared below.** The experimental CI lane runs exactly these targets, so keep them as narrow as the change allows.  ```experimental-tests # One target per line: a directory or a file. (A pytest ::selector is not # supported -- the sharding runner cannot consume one.) Must be under # tests/experimental/ and must exist. Delete these comment lines and add yours, e.g. # # tests/experimental/test_my_backend.py # tests/experimental/my_backend/ # # Declaring the whole tree (tests/experimental/) is allowed but means every # experimental PR pays for every other feature's tests, in every matrix cell. ``` ## Reviewer Notes Please review the real-Q prefix proof, actual-device architecture gate, and host descriptor update. Generated host IR confirms the runtime row-count guard modifies the descriptor later passed to launch. The writable host `TensorMap` alias works in the tested DSL runtime; cross-version support or an encoder-side helper needs DSL-owner review. This revision replaces the earlier native-Q implementation, which passed broad correctness checks but increased measured graph-operation latency by up to 3.43% in a separate B300 comparison. A separately frozen unconditional descriptor-correction port passed the original 37-case TensorRT-LLM sequence ten times (370 passes), supporting the diagnosis; those results are not validation of this exact guarded source. Keep the workaround scoped to documented Blackwell descriptor layouts and revisit it when a driver fix or supported encoder helper is available.  ## Summary by CodeRabbit * **Bug Fixes** * Improved attention correctness on supported Blackwell GPUs when processing partial query tiles in contiguous and paged attention. * Verified correct results with causal masking, both unwindowed and sliding-window attention, and captured graph replay after query data changes.  --------- Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> (cherry picked from commit f57c58b)

github-actions Bot assigned yuxianq Oct 14, 2026

yuxianq mentioned this pull request Oct 14, 2026

[NVBUG-6791733][fix] Guard ragged Q tensor maps on Blackwell flashinfer-ai/flashinfer#6482

Merged

2 of 11 tasks

yuxianq changed the title ~~[https://nvbugs/6791733][fix] Use native Q TMA loads in PrimTS context~~ [https://nvbugs/6791733][fix] Guard PrimTS ragged Q tensor maps on Blackwell Oct 15, 2026

yuxianq added 2 commits October 16, 2026 11:51

yuxianq force-pushed the bug/6791733-prims-ts-q-tma-release13 branch from 30950da to 346ced0 Compare October 16, 2026 11:52

yuxianq marked this pull request as ready for review October 16, 2026 11:53

yuxianq requested review from a team as code owners October 16, 2026 11:53

yuxianq requested review from heyuhhh, lfr-0531, mzweilz and yuanjingx87 October 16, 2026 11:53

yuxianq requested review from PerkzZheng and removed request for heyuhhh and lfr-0531 October 16, 2026 11:59

PerkzZheng reviewed Oct 19, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py Resolved

mzweilz approved these changes Oct 19, 2026

View reviewed changes

PerkzZheng approved these changes Oct 19, 2026

View reviewed changes

trtllm-agent added the ci: full pre-merge approved label Oct 19, 2026

yuxianq force-pushed the bug/6791733-prims-ts-q-tma-release13 branch from 1174214 to e957dcc Compare October 19, 2026 10:37

yuxianq requested a review from a team as a code owner October 19, 2026 10:37

yuxianq requested review from BowenFu and schetlur-nv October 19, 2026 10:37

yuxianq changed the title ~~[https://nvbugs/6791733][fix] Use native Q TMA loads in PrimTS context~~ [https://nvbugs/6791733][fix] Guard PrimTS ragged Q tensor maps on Blackwell Oct 19, 2026

brnguyen2 approved these changes Oct 19, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py Resolved

BowenFu approved these changes Oct 20, 2026

View reviewed changes

yuxianq merged commit b053431 into NVIDIA:release/1.3 Oct 21, 2026
1 check passed

yuxianq mentioned this pull request Oct 21, 2026

[None][chore] promote flashinfer-prims-ts source from #20298 #20731

Merged

3 tasks done

yunruis mentioned this pull request Oct 21, 2026

[None][feat] Add the PrimTS native sparse MLA FMHA for DeepSeek-V4 generation and enable it by default on SM100 #20209

Merged

1 task done

Conversation

yuxianq commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Per-File QA Perspective

Description

Vendor Promotion

Test Coverage

CI time limit

Validation

Reproduction

Performance

PR Checklist

GitHub Bot Help

Uh oh!

yuxianq commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

yuxianq commented Oct 16, 2026

Uh oh!

yuxianq commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

yuxianq commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

yuxianq commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

brnguyen2 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Oct 19, 2026

Uh oh!

yuxianq commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

yuxianq commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

tensorrt-cicd commented Oct 20, 2026

Uh oh!

yuxianq commented Oct 21, 2026

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

tensorrt-cicd commented Oct 21, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

yuxianq commented Oct 14, 2026 •

edited

Loading

yuxianq commented Oct 16, 2026 •

edited

Loading

coderabbitai Bot commented Oct 16, 2026 •

edited

Loading

yuxianq commented Oct 19, 2026 •

edited

Loading