tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py (2)

126-159: Consider documenting tensor shapes in docstring.

Per coding guidelines, documentation of Tensor-like arguments should include expected dimensions. The current docstring is minimal.

Suggested docstring improvement

 def update_paged_kv_cache(
     k: torch.Tensor,
     v: torch.Tensor,
     batch_indices: torch.Tensor,
     positions: torch.Tensor,
     kv_cache: torch.Tensor,
     kv_indices: torch.Tensor,
     kv_indptr: torch.Tensor,
 ) -> None:
-    """Update the combined paged KV cache with new K, V tensors."""
+    """Update the combined paged KV cache with new K, V tensors.
+
+    Args:
+        k: Key tensor [num_tokens, n_kv_heads, head_dim]
+        v: Value tensor [num_tokens, n_kv_heads, head_dim]
+        batch_indices: Sequence index per token [num_tokens]
+        positions: Position within sequence per token [num_tokens]
+        kv_cache: Paged cache [num_blocks, 2, n_kv_heads, page_size, head_dim]
+        kv_indices: Physical page indices (flattened)
+        kv_indptr: Cumulative page counts [num_seqs + 1]
+    """

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py`
around lines 126 - 159, The update_paged_kv_cache docstring is too
minimal—expand it to document the expected tensor shapes and meanings for each
argument (e.g., k: (num_tokens, n_kv_heads, head_dim), v: (num_tokens,
n_kv_heads, head_dim), batch_indices: (num_tokens,) int64, positions:
(num_tokens,) int64, kv_cache: (batch, n_kv_heads, num_pages, page_size,
head_dim) or the actual layout used, kv_indices: (batch, num_kv_heads,
num_pages) int32, kv_indptr: (batch+1,) or whatever indexing format is used),
and note that the function updates kv_cache in-place and returns None; include
any assumptions about contiguous memory/strides (cache_stride_block,
cache_stride_kv, cache_stride_head, cache_stride_token) so callers know layout
requirements and axis ordering when invoking update_paged_kv_cache.

---

849-866: GPU sync for multi-sequence batches may impact latency.

Line 853 uses .item() which synchronizes CPU-GPU. While documented as necessary for variable-length sequences, this could add latency in multi-sequence prefill scenarios. The comment explains the rationale well.

Also, the SDPA path condition at line 865 (kv_indices.shape[0] == total_expected_pages) restricts usage to batches where all sequences have identical page counts, which may rarely be true in practice.

Consider tracking this as a future optimization opportunity if multi-sequence prefill latency becomes a bottleneck.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_paged_attention.py`
around lines 849 - 866, The code forces a GPU sync via q_lens.max().item() when
computing max_q_len and then strictly requires kv_indices.shape[0] ==
total_expected_pages, which is brittle; change max_q_len to remain a tensor
(e.g., max_q_len = q_lens.max()) and propagate tensor arithmetic for max_pages
and total_expected_pages so no .item() is called, then update the SDPA predicate
in use_sdpa (the condition built from max_q_len, max_pages and kv_indices) to
operate on tensors and relax the kv_indices check (e.g., allow
kv_indices.shape[0] >= total_expected_pages or remove the exact-equality
requirement) so multi-sequence batches don't unnecessarily miss the SDPA path;
adjust any downstream uses of max_q_len/max_pages to handle tensor vs int
accordingly (referencing qo_indptr, max_q_len, max_pages, total_expected_pages,
use_sdpa, kv_indices).

tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_paged_attention.py (1)

220-299: Context kernel tests lack direct PyTorch reference comparison.

Unlike TestTritonPagedDecodeKernel, this class only tests for shape/NaN/Inf sanity but doesn't include a direct correctness test against PyTorch SDPA reference. The TestFlashInferComparison.test_prefill_vs_flashinfer provides correctness validation but requires FlashInfer to be installed.

Consider adding a PyTorch SDPA reference test similar to test_decode_kernel_vs_pytorch_reference for cases where FlashInfer is unavailable.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_triton_paged_attention.py`
around lines 220 - 299, Add a new test method (e.g.,
test_context_kernel_vs_pytorch_reference) that mirrors test_context_kernel_basic
but computes a direct PyTorch SDPA reference and compares it to
triton_paged_context; reconstruct the full per-batch per-head KV sequence from
kv_cache using kv_indices, kv_indptr, and seq_len_with_cache (or reassemble k/v
from the original k and v tensors used to populate the cache), cast tensors to
float32 for the reference, compute scaled dot-product attention per head (use
torch.matmul and softmax with sm_scale) to produce a reference output shaped
like q, then compare triton_paged_context(...) to the reference with
torch.allclose (use reasonable atol/rtol) and keep existing shape/NaN/Inf
checks; reference the functions triton_paged_context, update_paged_kv_cache, and
create_paged_kv_cache so the test locates the same fixtures.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes