[#22016][perf] Convert only the batch's slots in SpecSampler.update_requests

tensorrt_llm/_torch/speculative/spec_sampler_base.py

@@ -41,7 +41,6 @@
     SampleState,
     SampleStateTensors,
     TorchSampler,
-    add_token,
     int_tensor,
 )
 from ..pyexecutor.sampler.penalties import has_occurrence_penalty
@@ -433,11 +432,11 @@ def bind_spec_step_fold(self, spec_step_fold: SpecStepFold) -> None:
     def _request_common_handling(
         self,
         request: LlmRequest,
-        next_draft_tokens: list[list[int]],
+        next_draft_tokens: list[int],
         runtime_draft_len: Optional[int],
     ) -> None:
         """Common handling for both context and generation requests."""
-        request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len]
+        request.py_draft_tokens = next_draft_tokens[:runtime_draft_len]
         request.py_decoding_iter += 1
 
     def update_requests(
@@ -456,21 +455,28 @@ def update_requests(
         assert isinstance(state, SampleStateSpec)
 
         state.sampler_event.synchronize()
-        new_tokens = state.host.new_tokens.tolist()
-        new_tokens_lens_list = state.host.new_tokens_lens.tolist()
+        beam_idx = DEFAULT_BEAM_IDX
+        # The host copies hold every preallocated slot; convert only the batch's rows.
+        active_requests = [
+            (req_idx, req)
+            for req_idx, req in enumerate(state.requests)
+            if req.state != LlmRequestState.GENERATION_COMPLETE
+        ]
+        active_slots = torch.tensor(
+            [req.py_seq_slot for _, req in active_requests], dtype=torch.long
+        )
+        new_tokens_rows = state.host.new_tokens[:, active_slots, beam_idx].T.tolist()
+        new_tokens_lens_list = state.host.new_tokens_lens[active_slots].tolist()
         token_windows = (
             state.host.verify_lens.tolist() if state.host.verify_lens is not None else None
         )
-        next_draft_tokens_list = state.host.next_draft_tokens.tolist()
-        beam_idx = DEFAULT_BEAM_IDX
+        next_draft_tokens_rows = state.host.next_draft_tokens[active_slots].tolist()
         runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len)
 
         if token_windows is not None:
-            for request_index, request in enumerate(state.requests):
-                if request.state == LlmRequestState.GENERATION_COMPLETE:
-                    continue
+            for row, (request_index, request) in enumerate(active_requests):
                 token_window = token_windows[request_index]
-                num_new_tokens = new_tokens_lens_list[request.py_seq_slot]
+                num_new_tokens = new_tokens_lens_list[row]
                 max_new_tokens = min(token_window, self.max_accepted_path_len)
                 if state.draft_lens is not None:
                     max_new_tokens = min(max_new_tokens, state.draft_lens[request_index] + 1)
@@ -482,19 +488,18 @@ def update_requests(
                         "the executed window, original proposals, and token storage"
                     )
 
-        for req_idx, req in enumerate(state.requests):
-            if req.state == LlmRequestState.GENERATION_COMPLETE:
-                continue
-            num_new_tokens = new_tokens_lens_list[req.py_seq_slot]
-            # new_tokens is sized to this bound, and add_token indexes a plain
+        for row, (req_idx, req) in enumerate(active_requests):
+            num_new_tokens = new_tokens_lens_list[row]
+            # new_tokens is sized to this bound, and each row is a plain
             # host-side list, so a violation would otherwise surface as an
             # opaque IndexError.
             assert num_new_tokens <= self.max_accepted_path_len, (
                 f"accepted {num_new_tokens} tokens in one step, but new_tokens is "
                 f"sized for {self.max_accepted_path_len}"
             )
             for i in range(num_new_tokens):
-                new_token = add_token(req, new_tokens, beam_idx=beam_idx, step=i)
+                new_token = new_tokens_rows[row][i]
+                req.add_new_token(new_token, beam_idx)
                 if handle_stop_criteria(
                     req, new_token, max_seq_len=self.max_seq_len, beam_idx=beam_idx
                 ):
@@ -511,7 +516,7 @@ def update_requests(
             drafted_len = state.draft_lens[req_idx] if state.draft_lens is not None else 0
             req.py_num_draft_tokens_verified = min(drafted_len, verified_len)
             req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens
-            self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len)
+            self._request_common_handling(req, next_draft_tokens_rows[row], runtime_draft_len)
 
     def _scatter_to_stores(
         self, outputs: dict[str, torch.Tensor], num_skip: int, slots: list[int]

tests/unittest/_torch/speculative/hw_agnostic/test_spec_sampler_update_requests.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Host-side request updates of the one-model speculative sampler."""
+
+import types
+from typing import Optional
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
+from tensorrt_llm._torch.speculative import spec_sampler_base as sampler_ops
+from tensorrt_llm._torch.speculative.spec_sampler_base import (
+    SampleStateSpec,
+    SampleStateTensorsSpec,
+    SpecSampler,
+)
+
+pytestmark = pytest.mark.cpu_only
+
+NUM_SLOTS = 16
+PATH_LEN = 4
+DRAFT_LEN = 3
+
+
+def _request(slot: Optional[int], state=LlmRequestState.GENERATION_IN_PROGRESS):
+    request = types.SimpleNamespace(
+        py_seq_slot=slot,
+        state=state,
+        py_draft_tokens=[99] * DRAFT_LEN,
+        py_decoding_iter=0,
+        tokens=[],
+    )
+    request.add_new_token = lambda token, beam_idx: request.tokens.append(token)
+    return request
+
+
+@pytest.mark.parametrize("verify_lens", [None, [3, 2, 1, 2]])
+def test_update_requests_reads_each_request_slot_row(monkeypatch, verify_lens):
+    """Rows of scattered slots in a larger store must reach the request that owns them."""
+    monkeypatch.setattr(sampler_ops, "handle_stop_criteria", lambda *args, **kwargs: False)
+    sampler = SpecSampler.__new__(SpecSampler)
+    sampler.draft_len = DRAFT_LEN
+    sampler.max_accepted_path_len = PATH_LEN
+    sampler.max_seq_len = 128
+
+    new_tokens = torch.arange(PATH_LEN * NUM_SLOTS, dtype=torch.int32).reshape(
+        PATH_LEN, NUM_SLOTS, 1
+    )
+    new_tokens_lens = torch.full((NUM_SLOTS,), PATH_LEN + 1, dtype=torch.int32)
+    new_tokens_lens[[11, 0, 7]] = torch.tensor([3, 1, 2], dtype=torch.int32)
+    next_draft_tokens = 1000 + torch.arange(NUM_SLOTS * DRAFT_LEN, dtype=torch.int32).reshape(
+        NUM_SLOTS, DRAFT_LEN
+    )
+    host = SampleStateTensorsSpec(
+        new_tokens=new_tokens,
+        new_tokens_lens=new_tokens_lens,
+        next_draft_tokens=next_draft_tokens,
+        verify_lens=None if verify_lens is None else torch.tensor(verify_lens, dtype=torch.int32),
+    )
+    # The finished request between the others shifts list positions against slot rows.
+    requests = [
+        _request(11),
+        _request(None, LlmRequestState.GENERATION_COMPLETE),
+        _request(0),
+        _request(7),
+    ]
+    draft_lens = [3, 3, 0, 2]
+    runtime_draft_len = 2
+    state = SampleStateSpec(
+        requests=requests,
+        device=host,
+        host=host,
+        sampler_event=types.SimpleNamespace(synchronize=lambda: None),
+        runtime_draft_len=runtime_draft_len,
+        draft_lens=draft_lens,
+    )
+
+    sampler.update_requests(state)
+
+    finished = requests[1]
+    assert finished.tokens == []
+    assert finished.py_draft_tokens == [99] * DRAFT_LEN
+    assert finished.py_decoding_iter == 0
+    for request_index in (0, 2, 3):
+        request = requests[request_index]
+        slot = request.py_seq_slot
+        num_new_tokens = int(new_tokens_lens[slot])
+        verified_len = runtime_draft_len if verify_lens is None else verify_lens[request_index] - 1
+        assert request.tokens == new_tokens[:num_new_tokens, slot, 0].tolist()
+        assert request.py_num_accepted_draft_tokens == num_new_tokens - 1
+        assert request.py_num_draft_tokens_verified == min(draft_lens[request_index], verified_len)
+        assert request.py_rewind_len == verified_len - (num_new_tokens - 1)
+        assert request.py_draft_tokens == next_draft_tokens[slot, :runtime_draft_len].tolist()
+        assert request.py_decoding_iter == 1