[https://nvbugs/6791733][fix] Guard PrimTS ragged Q tensor maps on Blackwell

3rdparty/vendor_sources.lock.yaml

@@ -2,15 +2,15 @@ schema_version: 1
 vendors:
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261012
-    commit: ae312f7cc1723b06fda5168105ade0f328ee47d8
+    branch: bug-6791733-prims-ts-cuda-graph--source-q-tma-release13
+    commit: ae3851ab5f958ba22f154059df5b38321f6369d2
     source: flashinfer/attention/prims_ts
     destination: tensorrt_llm/_torch/attention/backends/prims_ts
     include:
     - '**/*.py'
     patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
     patch_digest: sha256:65cc08e20a9cc7bdc190eb6f3e79124986dced74394c1f3250ae2d100a8ee340
-    digest: sha256-tree-v1:e411b6c6679474d362b166e586df4276f52322d75a35c59c782d6266143e560e
+    digest: sha256-tree-v1:c9bcfbbac9db243110b7f9249ec0edd8ba4447ea03a7c89b1c7f6649d6d4bb6e
   sana-sol-attn:
     url: https://github.com/NVlabs/Sana.git
     branch: sol-engine

tensorrt_llm/_torch/attention/backends/prims_ts/context.py

@@ -328,6 +328,9 @@ def _make_context_kernel(
         enable_skip_correction=True,
         uses_ldtm_stat=uses_ldtm_stat,
         causal_single_kv_tile=(causal_single_kv_tile and not use_paged_kv),
+        use_blackwell_tma_oob_workaround=(
+            torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3))
+        ),
         **paged_kwargs,
     )
     return fmha

tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py

@@ -200,6 +200,22 @@
 from cutlass.experimental import primitives as prims
 
 
+def _compact_q_prefix_min_rows(q: cute.Tensor) -> int | None:
+    """Rows needed for a proven 128 KiB contiguous prefix, or None if unproven."""
+    if len(q.shape) != 3:
+        return None
+    h, d = q.shape[1:]
+    element_bits = q.element_type.width
+    if not all(isinstance(x, int) for x in (h, d, element_bits, *q.stride)):
+        return None
+    if h <= 0 or d <= 0 or element_bits not in (8, 16):
+        return None
+    if tuple(q.stride) != (h * d, d, 1):
+        return None
+    row_bytes = h * d * (element_bits // 8)
+    return (128 * 1024 + row_bytes - 1) // row_bytes
+
+
 def _as_i32(x: int | Int32) -> Int32:
     return Int32(x) if isinstance(x, int) else x
 
@@ -2743,6 +2759,9 @@ class FmhaTs:
         Use the fixed causal one-K/V-tile task domains. The context runner
         enables this only for query-paired, fixed-length inputs whose K/V
         extent fits one 128-token tile (default: False).
+    use_blackwell_tma_oob_workaround : bool, optional
+        Use conservative OOB address generation for ragged Q tensor maps.
+        Enable only for a Blackwell execution device (default: False).
     """
 
     def __init__(
@@ -2772,6 +2791,7 @@ def __init__(
         causal_single_kv_tile: bool = False,
         exhaustive_deadlock_race_check: bool = True,
         d_v: int | None = None,
+        use_blackwell_tma_oob_workaround: bool = False,
     ) -> None:
         """Initialize mode-specific tiling, dtype, and schedule configuration."""
         head_paired = resolve_head_paired_mode(
@@ -2815,6 +2835,7 @@ def __init__(
         self.is_causal = is_causal
         self.is_clc_dynamic = is_clc_dynamic
         self.exhaustive_deadlock_race_check = exhaustive_deadlock_race_check
+        self.use_blackwell_tma_oob_workaround = use_blackwell_tma_oob_workaround
 
         q_dtype = in_qk_dtype or cutlass.Float16
         k_dtype = in_qk_dtype or cutlass.Float16
@@ -3162,6 +3183,24 @@ def __call__(
                 # Zero-fill partial Q rows and head-dimension fragments.
                 oob_fill=cuda.TensorMapFloatOOBFill.none,
             )
+            if cutlass.const_expr(self.use_blackwell_tma_oob_workaround):
+                # Only the real Q layout can establish a mapped prefix;
+                # synthetic ragged dimensions and padded head sizes cannot.
+                minimum_rows = _compact_q_prefix_min_rows(q_cute)
+                correct_oob_addressing = True
+                if cutlass.const_expr(minimum_rows is not None):
+                    correct_oob_addressing = q_cute.shape[0] < minimum_rows
+                if correct_oob_addressing:
+                    # Work around the Blackwell driver TMA descriptor issue documented at
+                    # https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-615-71-09/index.html#known-issues
+                    # Clear Blackwell descriptor bit 85 on the host when a
+                    # contiguous 128 KiB prefix has not been established.
+                    descriptor_word1 = (
+                        cutlass.Pointer(tma_q_desc.get_ptr(), dtype=cutlass.Uint64) + 1
+                    )
+                    descriptor_word1.store(
+                        descriptor_word1.load() & ~cutlass.Uint64(1 << 21)
+                    )
         else:
             tma_q_desc = cuda.create_tensor_map_tiled_from_view(
                 q_cute,

tests/integration/test_lists/test-db/l0_b200.yml

@@ -100,8 +100,11 @@ l0_b200:
   #               by absolute index, and they gate on >= 2 GPUs plus Blackwell. They also
   #               rebind mpi4py's serializer at import time, which applies process-wide on
   #               every stage that collects them, including where the tests themselves skip.
-  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
+  # test_prims_ts_attention_backend.py keeps its own entry so its kernel JIT time does not
+  # count against this entry's timeout.
+  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/test_prims_ts_attention_backend.py
   - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120)
+  - unittest/_torch/attention/test_prims_ts_attention_backend.py
   - unittest/_torch/compilation
   - unittest/_torch/debugger
   - unittest/_torch/peft/test_fp8_lora_grouped_gemm_regressions.py

tests/integration/test_lists/test-db/l0_b300.yml

@@ -24,11 +24,14 @@ l0_b300:
   #               rebind mpi4py's serializer at import time, which applies process-wide on
   #               every stage that collects them, including where the tests themselves skip.
   # The sparse/dsa/test_cute_dsl_* files ignored below keep their own entries (pre-existing).
-  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py
+  # test_prims_ts_attention_backend.py keeps its own entry so its kernel JIT time does not
+  # count against this entry's timeout.
+  - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py --ignore=unittest/_torch/attention/test_prims_ts_attention_backend.py
   - unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py
   - unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py
   - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py TIMEOUT (120)
   - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py
+  - unittest/_torch/attention/test_prims_ts_attention_backend.py
   - unittest/_torch/modeling/test_qsa_runtime_wiring.py
   - unittest/_torch/modeling/test_qwen4_exp_support.py
   - unittest/_torch/thop/parallel TIMEOUT (90)