From dff43c56698401c0e0daa002a98c8cef014c94bb Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Wed, 14 Oct 2026 13:39:25 +0000
Subject: [PATCH 1/4] [https://nvbugs/6791733][fix] Use native Q TMA loads in
 PrimTS context

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>
---
 3rdparty/vendor_sources.lock.yaml             |  6 +-
 .../kernels/fmha_context/fmha_kernel.py       | 36 +++----
 .../kernels/fmha_context/fmha_resources.py    | 10 +-
 .../test_prims_ts_attention_backend.py        | 96 +++++++++++++++++++
 4 files changed, 111 insertions(+), 37 deletions(-)

diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml
index f58609e5f6e7..65d9194cac3d 100644
--- a/3rdparty/vendor_sources.lock.yaml
+++ b/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: 6b00497398f491136be91fb0e60ac3f314da53ba
     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:27b127506e2637f9e7f2a0cce0e3a3812597c0fdc688495542fdf4a40040e459
   sana-sol-attn:
     url: https://github.com/NVlabs/Sana.git
     branch: sol-engine
diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py
index c74fbcc1287c..b50fbe1026b2 100644
--- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py
+++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2026 by FlashInfer team.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -3145,31 +3146,16 @@ def __call__(
                     cfg.pv_mma_tiler[1] // cfg.tma_copy_v_iters,
                 )
                 kv_stride_order = stride_order
-        if cutlass.const_expr(cum_seqlen_q is not None):
-            # Q/O use the packed sequence axis as a ragged TMA dimension.
-            # Dense TMA bounds only see the full sum_seqlen tensor and cannot
-            # prevent a partial final tile from crossing into the next packed
-            # sequence. Paged K/V keep native page-pool descriptors. Invalid K
-            # scores are masked before softmax; invalid V rows are overwritten
-            # in SMEM after their TMA stage completes and before PV MMA.
-            tma_q_desc = create_tensor_map_ragged_from_tensor(
-                q_cute,
-                box_dims=q_box_dims,
-                ragged_dim=0,
-                stride_order=stride_order,
-                swizzle=tma_qkv_swizzle,
-                l2_promotion=cuda.TensorMapL2Promotion.none,
-                # Zero-fill partial Q rows and head-dimension fragments.
-                oob_fill=cuda.TensorMapFloatOOBFill.none,
-            )
-        else:
-            tma_q_desc = cuda.create_tensor_map_tiled_from_view(
-                q_cute,
-                box_dims=q_box_dims,
-                stride_order=stride_order,
-                swizzle=tma_qkv_swizzle,
-                l2_promotion=cuda.TensorMapL2Promotion.none,
-            )
+        # Q reductions are independent per row, so a packed tail tile may
+        # read the next request's Q rows. Ragged O stores discard those rows;
+        # the native Q descriptor bounds reads at the end of the allocation.
+        tma_q_desc = cuda.create_tensor_map_tiled_from_view(
+            q_cute,
+            box_dims=q_box_dims,
+            stride_order=stride_order,
+            swizzle=tma_qkv_swizzle,
+            l2_promotion=cuda.TensorMapL2Promotion.none,
+        )
 
         if cutlass.const_expr(cfg.use_paged_kv):
             # Paged-KV path: K/V are compact pool tensors with shape
diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py
index c4d1b5512f92..6233840fb35d 100644
--- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py
+++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py
@@ -1,4 +1,5 @@
 # Copyright (c) 2026 by FlashInfer team.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -1533,9 +1534,6 @@ def tma_load(
         q_seq_offset = (
             seq_coord_q + inst_idx * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m
         )
-        q_seq_extent = Int32(0)
-        if cutlass.const_expr(self.cfg.has_varlen):
-            q_seq_extent = cuseqlen_q + seqlen_q - q_seq_offset
         smem_stage_elements = self.cfg.tma_copy_q_elements
         d_granu_inner = self.cfg.tma_copy_q_granu_inner
 
@@ -1546,12 +1544,6 @@ def tma_load(
                 q_coords = (d_offset, q_head_coord, q_seq_offset, batch_coord)
                 if cutlass.const_expr(self.cfg.has_varlen):
                     q_coords = (d_offset, q_head_coord, q_seq_offset)
-                    q_coords = transform_ragged_coords(
-                        q_coords,
-                        ragged_dim_idx=2,
-                        ragged_box_size=self.cfg.qk_mma_tiler[0],
-                        ragged_extent=q_seq_extent,
-                    )
                 if cutlass.const_expr(self.cfg.two_cta_umma):
                     # Each CTA loads only its own 128 Q rows (multicast mask = own
                     # rank), but as a cta_group::2 copy, so the byte completion
diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
index 2adfac26c11c..960fd02e6647 100644
--- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
+++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
@@ -1310,3 +1310,99 @@ def counted_prims_context(self, *args, **kwargs):
 
     assert calls["fallback"] > 0
     assert calls["prims_context"] == 0
+
+
+@pytest.mark.parametrize("paged", [False, True], ids=["contiguous", "paged"])
+@pytest.mark.parametrize("window_left", [-1, 31], ids=["causal", "sliding-window"])
+def test_prims_ts_context_packed_query_tails(paged: bool, window_left: int) -> None:
+    from tensorrt_llm._torch.attention.backends.prims_ts import (
+        BatchPrefillPagedTSWrapper,
+        BatchPrefillTSWrapper,
+    )
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    num_qo_heads, num_kv_heads, head_dim = 8, 2, 128
+    page_size = 32
+    # Cover a tile crossing into another request, a partial second Q tile,
+    # and a final tile that extends beyond the packed allocation.
+    q_lengths = [3, 130, 1]
+    kv_lengths = [33, 160, 17]
+    q_offsets = [0, 3, 133, 134]
+    kv_offsets = [0, 33, 193, 210]
+    batch_size = len(q_lengths)
+    max_pages = math.ceil(max(kv_lengths) / page_size)
+    query = torch.randn(sum(q_lengths), num_qo_heads, head_dim, device=device, dtype=dtype)
+    k_cache = torch.randn(
+        batch_size * max_pages, num_kv_heads, page_size, head_dim, device=device, dtype=dtype
+    )
+    v_cache = torch.randn_like(k_cache)
+    block_tables = torch.arange(batch_size * max_pages, device=device, dtype=torch.int32).view(
+        batch_size, max_pages
+    )
+    qo_indptr = torch.tensor(q_offsets, device=device, dtype=torch.int32)
+    kv_indptr = torch.tensor(kv_offsets, device=device, dtype=torch.int32)
+    seq_lens = torch.tensor(kv_lengths, device=device, dtype=torch.int32)
+    keys, values = [], []
+    for batch_idx, kv_length in enumerate(kv_lengths):
+        pages = block_tables[batch_idx].long()
+        keys.append(k_cache[pages].transpose(1, 2).reshape(-1, num_kv_heads, head_dim)[:kv_length])
+        values.append(
+            v_cache[pages].transpose(1, 2).reshape(-1, num_kv_heads, head_dim)[:kv_length]
+        )
+    output = torch.empty_like(query)
+    plan_args = {
+        "device": device,
+        "batch_size": batch_size,
+        "max_seq_len_q": max(q_lengths),
+        "max_kv_len": max(kv_lengths),
+        "num_qo_heads": num_qo_heads,
+        "num_kv_heads": num_kv_heads,
+        "head_dim": head_dim,
+        "q_dtype": dtype,
+        "k_dtype": dtype,
+        "v_dtype": dtype,
+        "out_dtype": dtype,
+        "mask_type": "causal",
+        "window_left": window_left,
+    }
+    if paged:
+        wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND")
+        wrapper.plan(**plan_args, page_size=page_size)
+        run = functools.partial(
+            wrapper.run, query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output
+        )
+    else:
+        wrapper = BatchPrefillTSWrapper()
+        wrapper.plan(**plan_args, packed=True)
+        run = functools.partial(
+            wrapper.run, query, torch.cat(keys), torch.cat(values), qo_indptr, kv_indptr, out=output
+        )
+
+    def reference() -> torch.Tensor:
+        outputs = []
+        for batch_idx, (q_length, kv_length) in enumerate(zip(q_lengths, kv_lengths)):
+            q = query[q_offsets[batch_idx] : q_offsets[batch_idx + 1]].float()
+            k = keys[batch_idx].float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1)
+            v = values[batch_idx].float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1)
+            scores = torch.einsum("qhd,khd->hqk", q, k) / math.sqrt(head_dim)
+            q_positions = torch.arange(q_length, device=device) + kv_length - q_length
+            k_positions = torch.arange(kv_length, device=device)
+            allowed = k_positions[None, :] <= q_positions[:, None]
+            if window_left > 0:
+                allowed &= k_positions[None, :] >= q_positions[:, None] - window_left
+            scores.masked_fill_(~allowed, float("-inf"))
+            outputs.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), v))
+        return torch.cat(outputs).to(dtype)
+
+    run()
+    torch.testing.assert_close(output, reference(), atol=3e-2, rtol=3e-3)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run(validate=False)
+    query.neg_()
+    output.fill_(float("nan"))
+    graph.replay()
+    torch.cuda.synchronize()
+    torch.testing.assert_close(output, reference(), atol=3e-2, rtol=3e-3)

From 346ced0cac56c0963d04fd092be562ecf522ae45 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Thu, 15 Oct 2026 07:35:13 +0000
Subject: [PATCH 2/4] [https://nvbugs/6791733][fix] Guard PrimTS ragged Q
 tensor maps on Blackwell

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>
---
 3rdparty/vendor_sources.lock.yaml             |   4 +-
 .../attention/backends/prims_ts/context.py    |   3 +
 .../kernels/fmha_context/fmha_kernel.py       |  75 ++++++-
 .../kernels/fmha_context/fmha_resources.py    |  10 +-
 .../test_prims_ts_attention_backend.py        | 200 ++++++++++++++++++
 5 files changed, 278 insertions(+), 14 deletions(-)

diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml
index 65d9194cac3d..daa1e0d8b70c 100644
--- a/3rdparty/vendor_sources.lock.yaml
+++ b/3rdparty/vendor_sources.lock.yaml
@@ -3,14 +3,14 @@ vendors:
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
     branch: bug-6791733-prims-ts-cuda-graph--source-q-tma-release13
-    commit: 6b00497398f491136be91fb0e60ac3f314da53ba
+    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:27b127506e2637f9e7f2a0cce0e3a3812597c0fdc688495542fdf4a40040e459
+    digest: sha256-tree-v1:c9bcfbbac9db243110b7f9249ec0edd8ba4447ea03a7c89b1c7f6649d6d4bb6e
   sana-sol-attn:
     url: https://github.com/NVlabs/Sana.git
     branch: sol-engine
diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py
index 1163324ec632..9c20e4356d79 100644
--- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py
+++ b/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
diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py
index b50fbe1026b2..9b23b5dd7d38 100644
--- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py
+++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py
@@ -1,5 +1,4 @@
 # Copyright (c) 2026 by FlashInfer team.
-# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -201,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
 
@@ -2744,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__(
@@ -2773,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(
@@ -2816,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
@@ -3146,16 +3166,49 @@ def __call__(
                     cfg.pv_mma_tiler[1] // cfg.tma_copy_v_iters,
                 )
                 kv_stride_order = stride_order
-        # Q reductions are independent per row, so a packed tail tile may
-        # read the next request's Q rows. Ragged O stores discard those rows;
-        # the native Q descriptor bounds reads at the end of the allocation.
-        tma_q_desc = cuda.create_tensor_map_tiled_from_view(
-            q_cute,
-            box_dims=q_box_dims,
-            stride_order=stride_order,
-            swizzle=tma_qkv_swizzle,
-            l2_promotion=cuda.TensorMapL2Promotion.none,
-        )
+        if cutlass.const_expr(cum_seqlen_q is not None):
+            # Q/O use the packed sequence axis as a ragged TMA dimension.
+            # Dense TMA bounds only see the full sum_seqlen tensor and cannot
+            # prevent a partial final tile from crossing into the next packed
+            # sequence. Paged K/V keep native page-pool descriptors. Invalid K
+            # scores are masked before softmax; invalid V rows are overwritten
+            # in SMEM after their TMA stage completes and before PV MMA.
+            tma_q_desc = create_tensor_map_ragged_from_tensor(
+                q_cute,
+                box_dims=q_box_dims,
+                ragged_dim=0,
+                stride_order=stride_order,
+                swizzle=tma_qkv_swizzle,
+                l2_promotion=cuda.TensorMapL2Promotion.none,
+                # 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,
+                box_dims=q_box_dims,
+                stride_order=stride_order,
+                swizzle=tma_qkv_swizzle,
+                l2_promotion=cuda.TensorMapL2Promotion.none,
+            )
 
         if cutlass.const_expr(cfg.use_paged_kv):
             # Paged-KV path: K/V are compact pool tensors with shape
diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py
index 6233840fb35d..c4d1b5512f92 100644
--- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py
+++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py
@@ -1,5 +1,4 @@
 # Copyright (c) 2026 by FlashInfer team.
-# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -1534,6 +1533,9 @@ def tma_load(
         q_seq_offset = (
             seq_coord_q + inst_idx * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m
         )
+        q_seq_extent = Int32(0)
+        if cutlass.const_expr(self.cfg.has_varlen):
+            q_seq_extent = cuseqlen_q + seqlen_q - q_seq_offset
         smem_stage_elements = self.cfg.tma_copy_q_elements
         d_granu_inner = self.cfg.tma_copy_q_granu_inner
 
@@ -1544,6 +1546,12 @@ def tma_load(
                 q_coords = (d_offset, q_head_coord, q_seq_offset, batch_coord)
                 if cutlass.const_expr(self.cfg.has_varlen):
                     q_coords = (d_offset, q_head_coord, q_seq_offset)
+                    q_coords = transform_ragged_coords(
+                        q_coords,
+                        ragged_dim_idx=2,
+                        ragged_box_size=self.cfg.qk_mma_tiler[0],
+                        ragged_extent=q_seq_extent,
+                    )
                 if cutlass.const_expr(self.cfg.two_cta_umma):
                     # Each CTA loads only its own 128 Q rows (multicast mask = own
                     # rank), but as a cta_group::2 copy, so the byte completion
diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
index 960fd02e6647..2fb0b0d649b7 100644
--- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
+++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
@@ -16,6 +16,7 @@
 import functools
 import inspect
 import math
+from collections.abc import Callable
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -1406,3 +1407,202 @@ def reference() -> torch.Tensor:
     graph.replay()
     torch.cuda.synchronize()
     torch.testing.assert_close(output, reference(), atol=3e-2, rtol=3e-3)
+
+
+def _prims_ts_causal_context_reference(
+    queries: list[torch.Tensor], keys: list[torch.Tensor], values: list[torch.Tensor]
+) -> torch.Tensor:
+    outputs = []
+    for query, key, value in zip(queries, keys, values):
+        q_length, num_qo_heads, head_dim = query.shape
+        kv_length, num_kv_heads, _ = key.shape
+        key = key.float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1)
+        value = value.float().repeat_interleave(num_qo_heads // num_kv_heads, dim=1)
+        scores = torch.einsum("qhd,khd->hqk", query.float(), key) / math.sqrt(head_dim)
+        q_positions = torch.arange(q_length, device=query.device) + kv_length - q_length
+        k_positions = torch.arange(kv_length, device=query.device)
+        scores.masked_fill_(k_positions[None, :] > q_positions[:, None], float("-inf"))
+        outputs.append(torch.einsum("hqk,khd->qhd", scores.softmax(-1), value))
+    return torch.cat(outputs)
+
+
+def _check_prims_ts_context_replays(
+    query: torch.Tensor,
+    output: torch.Tensor,
+    run: Callable[..., torch.Tensor],
+    reference: Callable[[], torch.Tensor],
+) -> None:
+    is_fp8 = query.dtype == torch.float8_e4m3fn
+    atol, rtol, relative_l2_limit = (0.13, 0.05, 0.1) if is_fp8 else (0.03, 0.003, 0.02)
+
+    def check() -> None:
+        expected = reference()
+        actual = output.reshape(expected.shape).float()
+        torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+        relative_l2 = torch.linalg.vector_norm(actual - expected) / torch.linalg.vector_norm(
+            expected
+        )
+        assert relative_l2.item() < relative_l2_limit
+
+    run()
+    check()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run(validate=False)
+    output.fill_(float("nan"))
+    graph.replay()
+    torch.cuda.synchronize()
+    check()
+    # copy_ also supports FP8, whose in-place arithmetic is not implemented.
+    query.copy_(-query.float())
+    output.fill_(float("nan"))
+    graph.replay()
+    torch.cuda.synchronize()
+    check()
+
+
+def _check_prims_ts_compact_query(
+    paged: bool,
+    dtype: torch.dtype,
+    head_dim: int,
+    total_q: int,
+    storage_offset_rows: int = 0,
+) -> None:
+    from tensorrt_llm._torch.attention.backends.prims_ts import (
+        BatchPrefillPagedTSWrapper,
+        BatchPrefillTSWrapper,
+    )
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    num_qo_heads, num_kv_heads, page_size = 8, 2, 32
+    q_lengths = [3, total_q - 4, 1]
+    kv_lengths = [33, total_q + 28, 17]
+    q_offsets = [0, 3, total_q - 1, total_q]
+    kv_offsets = [0, 33, total_q + 61, total_q + 78]
+    query_storage = torch.randn(
+        total_q + storage_offset_rows,
+        num_qo_heads,
+        head_dim,
+        device=device,
+        dtype=torch.bfloat16,
+    ).to(dtype)
+    query = query_storage[storage_offset_rows:]
+    assert query.is_contiguous()
+    assert query.storage_offset() == storage_offset_rows * num_qo_heads * head_dim
+    keys = [
+        torch.randn(length, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16).to(dtype)
+        for length in kv_lengths
+    ]
+    values = [
+        torch.randn(length, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16).to(dtype)
+        for length in kv_lengths
+    ]
+    output = torch.empty(query.shape, device=device, dtype=torch.bfloat16)
+    qo_indptr = torch.tensor(q_offsets, device=device, dtype=torch.int32)
+    plan_args = {
+        "device": device,
+        "batch_size": len(q_lengths),
+        "max_seq_len_q": max(q_lengths),
+        "max_kv_len": max(kv_lengths),
+        "num_qo_heads": num_qo_heads,
+        "num_kv_heads": num_kv_heads,
+        "head_dim": head_dim,
+        "q_dtype": dtype,
+        "k_dtype": dtype,
+        "v_dtype": dtype,
+        "out_dtype": torch.bfloat16,
+        "mask_type": "causal",
+    }
+    if paged:
+        max_pages = math.ceil(max(kv_lengths) / page_size)
+        cache_shape = (len(q_lengths) * max_pages, num_kv_heads, page_size, head_dim)
+        k_cache = torch.zeros(cache_shape, device=device, dtype=dtype)
+        v_cache = torch.zeros_like(k_cache)
+        for batch_idx, (key, value) in enumerate(zip(keys, values)):
+            for start in range(0, key.shape[0], page_size):
+                page = batch_idx * max_pages + start // page_size
+                count = min(page_size, key.shape[0] - start)
+                k_cache[page, :, :count].copy_(key[start : start + count].transpose(0, 1))
+                v_cache[page, :, :count].copy_(value[start : start + count].transpose(0, 1))
+        block_tables = torch.arange(
+            len(q_lengths) * max_pages, device=device, dtype=torch.int32
+        ).view(len(q_lengths), max_pages)
+        seq_lens = torch.tensor(kv_lengths, device=device, dtype=torch.int32)
+        wrapper = BatchPrefillPagedTSWrapper(kv_layout="HND")
+        wrapper.plan(**plan_args, page_size=page_size)
+        run = functools.partial(
+            wrapper.run, query, k_cache, v_cache, qo_indptr, block_tables, seq_lens, out=output
+        )
+    else:
+        kv_indptr = torch.tensor(kv_offsets, device=device, dtype=torch.int32)
+        wrapper = BatchPrefillTSWrapper()
+        wrapper.plan(**plan_args, packed=True)
+        run = functools.partial(
+            wrapper.run, query, torch.cat(keys), torch.cat(values), qo_indptr, kv_indptr, out=output
+        )
+
+    def reference() -> torch.Tensor:
+        queries = [query[start:end] for start, end in zip(q_offsets[:-1], q_offsets[1:])]
+        return _prims_ts_causal_context_reference(queries, keys, values)
+
+    _check_prims_ts_context_replays(query, output, run, reference)
+
+
+@pytest.mark.parametrize("paged", [False, True], ids=["contiguous", "paged"])
+@pytest.mark.parametrize("delta", [-1, 0, 1], ids=["below", "at", "above"])
+@pytest.mark.parametrize(
+    "dtype,head_dim,boundary_q",
+    [
+        pytest.param(torch.bfloat16, 128, 64, id="bf16-d128"),
+        pytest.param(torch.bfloat16, 256, 32, id="bf16-d256"),
+        pytest.param(torch.float8_e4m3fn, 128, 128, id="fp8-d128"),
+        pytest.param(torch.float8_e4m3fn, 256, 64, id="fp8-d256"),
+    ],
+)
+def test_prims_ts_context_compact_query_boundary(
+    paged: bool, delta: int, dtype: torch.dtype, head_dim: int, boundary_q: int
+) -> None:
+    # These compact Q tensors span just below, at, and above 128 KiB.
+    _check_prims_ts_compact_query(paged, dtype, head_dim, boundary_q + delta)
+
+
+def test_prims_ts_context_compact_query_storage_offset() -> None:
+    _check_prims_ts_compact_query(
+        paged=False, dtype=torch.bfloat16, head_dim=128, total_q=64, storage_offset_rows=1
+    )
+
+
+def test_prims_ts_context_fixed_query_control() -> None:
+    from tensorrt_llm._torch.attention.backends.prims_ts import BatchPrefillTSWrapper
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    query = torch.randn(2, 3, 8, 128, device=device, dtype=torch.bfloat16)
+    keys = torch.randn(2, 33, 2, 128, device=device, dtype=torch.bfloat16)
+    values = torch.randn_like(keys)
+    output = torch.empty_like(query)
+    wrapper = BatchPrefillTSWrapper()
+    wrapper.plan(
+        device=device,
+        batch_size=2,
+        max_seq_len_q=3,
+        max_kv_len=33,
+        num_qo_heads=8,
+        num_kv_heads=2,
+        head_dim=128,
+        q_dtype=torch.bfloat16,
+        k_dtype=torch.bfloat16,
+        v_dtype=torch.bfloat16,
+        out_dtype=torch.bfloat16,
+        packed=False,
+        mask_type="causal",
+    )
+    run = functools.partial(wrapper.run, query, keys, values, out=output)
+
+    def reference() -> torch.Tensor:
+        return _prims_ts_causal_context_reference(
+            list(query.unbind()), list(keys.unbind()), list(values.unbind())
+        )
+
+    _check_prims_ts_context_replays(query, output, run, reference)

From e957dccee6b35f71223c86c7779cf410b0fc49d9 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Mon, 19 Oct 2026 10:21:50 +0000
Subject: [PATCH 3/4] [https://nvbugs/6791733][test] Keep PrimTS context tests
 within CI time 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>
---
 .../test_lists/test-db/l0_b200.yml            |  5 ++++-
 .../test_lists/test-db/l0_b300.yml            |  5 ++++-
 .../test_prims_ts_attention_backend.py        | 22 ++++++++++++++-----
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml
index ed023f9e176f..14fbe8a6ebb9 100644
--- a/tests/integration/test_lists/test-db/l0_b200.yml
+++ b/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
diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml
index 3a0e50c09bae..924a25b3fc48 100644
--- a/tests/integration/test_lists/test-db/l0_b300.yml
+++ b/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)
diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
index 2fb0b0d649b7..fe0c4597c5f7 100644
--- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
+++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
@@ -1467,6 +1467,7 @@ def _check_prims_ts_compact_query(
     head_dim: int,
     total_q: int,
     storage_offset_rows: int = 0,
+    plan_total_q: int | None = None,
 ) -> None:
     from tensorrt_llm._torch.attention.backends.prims_ts import (
         BatchPrefillPagedTSWrapper,
@@ -1476,6 +1477,10 @@ def _check_prims_ts_compact_query(
     torch.manual_seed(0)
     device = torch.device("cuda")
     num_qo_heads, num_kv_heads, page_size = 8, 2, 32
+    # Equal plan bounds share one compiled kernel across Q sizes, which keeps
+    # JIT time down and makes the kernel choose the descriptor path at run time.
+    plan_total_q = plan_total_q or total_q
+    assert plan_total_q >= total_q
     q_lengths = [3, total_q - 4, 1]
     kv_lengths = [33, total_q + 28, 17]
     q_offsets = [0, 3, total_q - 1, total_q]
@@ -1503,8 +1508,8 @@ def _check_prims_ts_compact_query(
     plan_args = {
         "device": device,
         "batch_size": len(q_lengths),
-        "max_seq_len_q": max(q_lengths),
-        "max_kv_len": max(kv_lengths),
+        "max_seq_len_q": plan_total_q - 4,
+        "max_kv_len": plan_total_q + 28,
         "num_qo_heads": num_qo_heads,
         "num_kv_heads": num_kv_heads,
         "head_dim": head_dim,
@@ -1515,7 +1520,7 @@ def _check_prims_ts_compact_query(
         "mask_type": "causal",
     }
     if paged:
-        max_pages = math.ceil(max(kv_lengths) / page_size)
+        max_pages = math.ceil(plan_args["max_kv_len"] / page_size)
         cache_shape = (len(q_lengths) * max_pages, num_kv_heads, page_size, head_dim)
         k_cache = torch.zeros(cache_shape, device=device, dtype=dtype)
         v_cache = torch.zeros_like(k_cache)
@@ -1564,12 +1569,19 @@ def test_prims_ts_context_compact_query_boundary(
     paged: bool, delta: int, dtype: torch.dtype, head_dim: int, boundary_q: int
 ) -> None:
     # These compact Q tensors span just below, at, and above 128 KiB.
-    _check_prims_ts_compact_query(paged, dtype, head_dim, boundary_q + delta)
+    _check_prims_ts_compact_query(
+        paged, dtype, head_dim, boundary_q + delta, plan_total_q=boundary_q + 1
+    )
 
 
 def test_prims_ts_context_compact_query_storage_offset() -> None:
     _check_prims_ts_compact_query(
-        paged=False, dtype=torch.bfloat16, head_dim=128, total_q=64, storage_offset_rows=1
+        paged=False,
+        dtype=torch.bfloat16,
+        head_dim=128,
+        total_q=64,
+        storage_offset_rows=1,
+        plan_total_q=65,
     )
 
 

From a7d6732b1e451e67e1d5396251395f6a742a7e4f Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Mon, 19 Oct 2026 11:54:38 +0000
Subject: [PATCH 4/4] [https://nvbugs/6791733][test] Cut PrimTS backend test
 JIT time

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>
---
 .../test_prims_ts_attention_backend.py        | 131 ++++++------------
 1 file changed, 41 insertions(+), 90 deletions(-)

diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
index fe0c4597c5f7..878512866e7c 100644
--- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
+++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
@@ -31,6 +31,15 @@
 )
 
 
+@pytest.fixture(autouse=True)
+def _trtllm_backend_only(monkeypatch: pytest.MonkeyPatch) -> None:
+    # PrimTS runs inside the TRTLLM backend. Other backends have their own
+    # tests and would only add JIT time here.
+    import backend_case
+
+    monkeypatch.setattr(backend_case, "BACKENDS_UNDER_TEST", ("TRTLLM",))
+
+
 _QWEN2_7B = {
     "num_heads": 28,
     "num_kv_heads": 4,
@@ -83,9 +92,10 @@
             },
             id="generation",
         ),
+        # The context request matches the context case so both reuse one kernel.
         pytest.param(
             {
-                "seq_lens": [41, 1],
+                "seq_lens": [65, 1],
                 "num_cached_tokens": [0, 63],
                 "num_contexts": 1,
             },
@@ -193,29 +203,6 @@ def run_with_zeroed_v_tail_check(
     assert zeroed_tails == [(0, 1), (1, 5)]
 
 
-def test_prims_ts_fp16_dense_context_with_alternate_shape(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts")
-    case = BackendCase(
-        num_heads=8,
-        num_kv_heads=2,
-        head_dim=256,
-        seq_lens=[67, 33],
-        num_cached_tokens=[0, 0],
-        num_contexts=2,
-        dtype="float16",
-        causal=False,
-        kv_layout="HND",
-        page_size=64,
-        use_kv_cache_manager_v2=True,
-    )
-
-    results = run_case(case)
-
-    assert "TRTLLM" in results
-
-
 def test_prims_ts_uses_compact_preprocessing_and_separate_decode_workspace(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -372,8 +359,11 @@ def test_prims_ts_deepseek_v3_lite_mla_generation(
     run_case(case)
 
 
-@pytest.mark.parametrize("num_heads", [6, 12, 96])
-@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"])
+# 6 and 96 heads select the M8 and repeated-M64 Q tiles; the scales test below
+# covers the M16 tile with 12 heads.
+@pytest.mark.parametrize(
+    "use_kv_cache_manager_v2,num_heads", [(False, 6), (True, 96)], ids=["v1-6", "v2-96"]
+)
 def test_prims_ts_fp8_mla_preprocessing(
     monkeypatch: pytest.MonkeyPatch, num_heads: int, use_kv_cache_manager_v2: bool
 ) -> None:
@@ -437,7 +427,7 @@ def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None:
     assert calls == [0, 1, 0, 1]
 
 
-@pytest.mark.parametrize("num_heads", [6, 12, 96])
+@pytest.mark.parametrize("num_heads", [12, 96])
 @pytest.mark.parametrize("batch_size", [2, 65])
 @pytest.mark.parametrize("max_seq_len", [64, 128], ids=["short-kv", "regular-kv"])
 def test_prims_ts_fp8_mla_scales_and_graph_replay(
@@ -1313,9 +1303,12 @@ def counted_prims_context(self, *args, **kwargs):
     assert calls["prims_context"] == 0
 
 
-@pytest.mark.parametrize("paged", [False, True], ids=["contiguous", "paged"])
-@pytest.mark.parametrize("window_left", [-1, 31], ids=["causal", "sliding-window"])
-def test_prims_ts_context_packed_query_tails(paged: bool, window_left: int) -> None:
+@pytest.mark.parametrize(
+    "window_left,paged",
+    [(-1, False), (-1, True), (31, False)],
+    ids=["causal-contiguous", "causal-paged", "sliding-window-contiguous"],
+)
+def test_prims_ts_context_packed_query_tails(window_left: int, paged: bool) -> None:
     from tensorrt_llm._torch.attention.backends.prims_ts import (
         BatchPrefillPagedTSWrapper,
         BatchPrefillTSWrapper,
@@ -1467,7 +1460,6 @@ def _check_prims_ts_compact_query(
     head_dim: int,
     total_q: int,
     storage_offset_rows: int = 0,
-    plan_total_q: int | None = None,
 ) -> None:
     from tensorrt_llm._torch.attention.backends.prims_ts import (
         BatchPrefillPagedTSWrapper,
@@ -1477,10 +1469,11 @@ def _check_prims_ts_compact_query(
     torch.manual_seed(0)
     device = torch.device("cuda")
     num_qo_heads, num_kv_heads, page_size = 8, 2, 32
-    # Equal plan bounds share one compiled kernel across Q sizes, which keeps
-    # JIT time down and makes the kernel choose the descriptor path at run time.
-    plan_total_q = plan_total_q or total_q
-    assert plan_total_q >= total_q
+    # Plan with the packed-tail test's bounds. Every Q size then reuses one
+    # compiled kernel per dtype, head size, and layout (BF16 D128 reuses the
+    # packed-tail kernels), and the kernel chooses the descriptor path at run time.
+    max_seq_len_q, max_kv_len = 130, 160
+    assert total_q - 4 <= max_seq_len_q and total_q + 28 <= max_kv_len
     q_lengths = [3, total_q - 4, 1]
     kv_lengths = [33, total_q + 28, 17]
     q_offsets = [0, 3, total_q - 1, total_q]
@@ -1508,8 +1501,8 @@ def _check_prims_ts_compact_query(
     plan_args = {
         "device": device,
         "batch_size": len(q_lengths),
-        "max_seq_len_q": plan_total_q - 4,
-        "max_kv_len": plan_total_q + 28,
+        "max_seq_len_q": max_seq_len_q,
+        "max_kv_len": max_kv_len,
         "num_qo_heads": num_qo_heads,
         "num_kv_heads": num_kv_heads,
         "head_dim": head_dim,
@@ -1520,7 +1513,7 @@ def _check_prims_ts_compact_query(
         "mask_type": "causal",
     }
     if paged:
-        max_pages = math.ceil(plan_args["max_kv_len"] / page_size)
+        max_pages = math.ceil(max_kv_len / page_size)
         cache_shape = (len(q_lengths) * max_pages, num_kv_heads, page_size, head_dim)
         k_cache = torch.zeros(cache_shape, device=device, dtype=dtype)
         v_cache = torch.zeros_like(k_cache)
@@ -1554,67 +1547,25 @@ def reference() -> torch.Tensor:
     _check_prims_ts_context_replays(query, output, run, reference)
 
 
-@pytest.mark.parametrize("paged", [False, True], ids=["contiguous", "paged"])
 @pytest.mark.parametrize("delta", [-1, 0, 1], ids=["below", "at", "above"])
 @pytest.mark.parametrize(
-    "dtype,head_dim,boundary_q",
+    "dtype,head_dim,boundary_q,paged",
     [
-        pytest.param(torch.bfloat16, 128, 64, id="bf16-d128"),
-        pytest.param(torch.bfloat16, 256, 32, id="bf16-d256"),
-        pytest.param(torch.float8_e4m3fn, 128, 128, id="fp8-d128"),
-        pytest.param(torch.float8_e4m3fn, 256, 64, id="fp8-d256"),
+        pytest.param(torch.bfloat16, 128, 64, False, id="bf16-d128-contiguous"),
+        pytest.param(torch.bfloat16, 128, 64, True, id="bf16-d128-paged"),
+        pytest.param(torch.bfloat16, 256, 32, True, id="bf16-d256-paged"),
+        pytest.param(torch.float8_e4m3fn, 128, 128, False, id="fp8-d128-contiguous"),
     ],
 )
 def test_prims_ts_context_compact_query_boundary(
-    paged: bool, delta: int, dtype: torch.dtype, head_dim: int, boundary_q: int
+    delta: int, dtype: torch.dtype, head_dim: int, boundary_q: int, paged: bool
 ) -> None:
-    # These compact Q tensors span just below, at, and above 128 KiB.
-    _check_prims_ts_compact_query(
-        paged, dtype, head_dim, boundary_q + delta, plan_total_q=boundary_q + 1
-    )
+    # These compact Q tensors span just below, at, and above 128 KiB. The groups
+    # cover both dtypes, head sizes, and KV layouts with three row thresholds.
+    _check_prims_ts_compact_query(paged, dtype, head_dim, boundary_q + delta)
 
 
 def test_prims_ts_context_compact_query_storage_offset() -> None:
     _check_prims_ts_compact_query(
-        paged=False,
-        dtype=torch.bfloat16,
-        head_dim=128,
-        total_q=64,
-        storage_offset_rows=1,
-        plan_total_q=65,
+        paged=False, dtype=torch.bfloat16, head_dim=128, total_q=64, storage_offset_rows=1
     )
-
-
-def test_prims_ts_context_fixed_query_control() -> None:
-    from tensorrt_llm._torch.attention.backends.prims_ts import BatchPrefillTSWrapper
-
-    torch.manual_seed(0)
-    device = torch.device("cuda")
-    query = torch.randn(2, 3, 8, 128, device=device, dtype=torch.bfloat16)
-    keys = torch.randn(2, 33, 2, 128, device=device, dtype=torch.bfloat16)
-    values = torch.randn_like(keys)
-    output = torch.empty_like(query)
-    wrapper = BatchPrefillTSWrapper()
-    wrapper.plan(
-        device=device,
-        batch_size=2,
-        max_seq_len_q=3,
-        max_kv_len=33,
-        num_qo_heads=8,
-        num_kv_heads=2,
-        head_dim=128,
-        q_dtype=torch.bfloat16,
-        k_dtype=torch.bfloat16,
-        v_dtype=torch.bfloat16,
-        out_dtype=torch.bfloat16,
-        packed=False,
-        mask_type="causal",
-    )
-    run = functools.partial(wrapper.run, query, keys, values, out=output)
-
-    def reference() -> torch.Tensor:
-        return _prims_ts_causal_context_reference(
-            list(query.unbind()), list(keys.unbind()), list(values.unbind())
-        )
-
-    _check_prims_ts_context_replays(query, output, run, reference)