[None][perf] Seed the GVR prior from device-side prefill lengths

tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py

@@ -1609,9 +1609,11 @@ def sparse_attn_indexer(
 
         # Dense skip outputs remain valid priors if a later step enters GVR.
         if has_prefill:
+            # Device lengths: gathering the CUDA selections with the host
+            # seq_lens would force a synchronous H2D copy every layer.
             self.top_k.update_gvr_prior_from_prefill(
                 topk_indices_buffer[:num_ctx_tokens],
-                metadata.seq_lens[:num_contexts],
+                metadata.seq_lens_cuda[:num_contexts],
                 gvr_prior_indices,
                 request_offset=num_generations,
             )

tensorrt_llm/_torch/modules/top_k.py

@@ -467,7 +467,9 @@ def update_gvr_prior_from_prefill(
         Args:
             output_indices: Int32 prefill selections with shape
                 ``[num_prefill_rows, top_k]``.
-            request_lengths: Per-request prefill row counts.
+            request_lengths: Per-request prefill row counts on
+                ``output_indices.device``; a host tensor here makes the row
+                gather a synchronous host-to-device copy.
             gvr_prior_indices: Int32 caller-owned state on
                 ``output_indices.device`` with shape ``[capacity, top_k]``.
                 The slice starting at ``request_offset`` is updated in place.

tests/unittest/_torch/modules/test_top_k.py

@@ -298,14 +298,25 @@ def test_gvr_v2_decode_rejects_output_width_mismatch(monkeypatch) -> None:
     runner.assert_not_called()
 
 
-def test_update_gvr_prior_from_prefill_uses_last_request_rows() -> None:
+@pytest.mark.parametrize(
+    "device",
+    [
+        "cpu",
+        pytest.param(
+            "cuda",
+            marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"),
+        ),
+    ],
+)
+def test_update_gvr_prior_from_prefill_uses_last_request_rows(device) -> None:
     top_k = TopK(2, decode_implementation=TopKImplementation.CUTE_DSL_GVR)
-    prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32)
-    prior_indices = torch.zeros(3, 2, dtype=torch.int32)
+    prefill_indices = torch.tensor([[0, 1], [2, 3], [4, 5]], dtype=torch.int32, device=device)
+    prior_indices = torch.zeros(3, 2, dtype=torch.int32, device=device)
 
+    # Production passes the device seq_lens twin so the row gather stays async.
     top_k.update_gvr_prior_from_prefill(
         prefill_indices,
-        torch.tensor([2, 1], dtype=torch.int32),
+        torch.tensor([2, 1], dtype=torch.int32, device=device),
         prior_indices,
         request_offset=1,
     )