diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py index 0dd8a6173806..848baca80a96 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py +++ b/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, ) diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index 2f99e392b6c1..6f323d839bf9 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/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. diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index ae6cda33d9df..daed3f24c9ba 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/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, )