Skip to content

[TRTLLM-14155][fix] Revert host-side greedy stop checks from #15920 #16163

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chzblych merged 1 commit into from
Jul 9, 2026
Merged

[TRTLLM-14155][fix] Revert host-side greedy stop checks from #15920 #16163

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 34 additions & 87 deletions tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1113,8 +1113,6 @@ def finish_reasons_list(self) -> FinishReasonsList:
@dataclass(kw_only=True)
class SampleStateTorch(SampleState[SampleStateTensorsHostTorch, SampleStateTensors]):
beam_history_builders: list[BeamHistoryBuilder | None] | None = None
use_host_stop_criteria: bool = False
"""Whether update_requests should evaluate end-ID and length limits on the host."""


@dataclass(kw_only=True, frozen=True)
Expand Down Expand Up @@ -2454,21 +2452,6 @@ def _can_use_fast_greedy_path(self, requests: list[LlmRequest]) -> bool:
return False
return True

def _can_use_host_stop_criteria(self, requests: list[LlmRequest]) -> bool:
"""Check whether stop criteria can be evaluated from the host token copy.

The non-speculative, single-beam path already copies sampled tokens to
the host. When no request has stop words, checking end IDs and length
limits on the host avoids the device finish-reason kernels and their
additional D2H copy.
"""
return (
bool(requests)
and self.max_tokens == 1
and self.max_beam_width == 1
and all(not req.py_is_draft and not req.py_stop_words_list for req in requests)
)

@staticmethod
def _meet_max_token_stop_criteria(
request: LlmRequest, max_seq_len: int, beam_idx: int = DEFAULT_BEAM_IDX
Expand Down Expand Up @@ -3736,23 +3719,13 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None:
req.py_rewind_len = 0
else:
processed = 1
if state.use_host_stop_criteria:
new_token = add_token(req, new_tokens_list, beam_idx=DEFAULT_BEAM_IDX)
self._handle_stop_criteria(
req,
new_token,
max_seq_len=self.max_seq_len,
beam_idx=DEFAULT_BEAM_IDX,
)
num_accepted = 0
else:
num_accepted = self.process_draft_tokens(
req,
new_tokens_tensor=new_tokens,
new_tokens_list=new_tokens_list,
finish_reasons=finish_reasons,
resource_manager=resource_manager,
)
num_accepted = self.process_draft_tokens(
req,
new_tokens_tensor=new_tokens,
new_tokens_list=new_tokens_list,
finish_reasons=finish_reasons,
resource_manager=resource_manager,
)
if (actual_draft_len := get_draft_token_length(req)) > 0:
req.py_num_accepted_draft_tokens = num_accepted
req.py_rewind_len = actual_draft_len - num_accepted
Expand Down Expand Up @@ -3811,10 +3784,10 @@ def sample_async(
(
requests,
seq_slots_host,
seq_lens_host,
Comment thread
chzblych marked this conversation as resolved.
seq_slots_cuda,
seq_lens_cuda,
new_tokens_host,
use_host_stop_criteria,
) = self._process_requests(
scheduled_requests,
model_outputs,
Expand All @@ -3828,8 +3801,7 @@ def sample_async(
# Forwarded to _record_sampler_event so SamplerEvent.synchronize
# awaits any side-stream D2H copies host-side.
side_stream_event: torch.cuda.Event | None = None
if requests and not use_host_stop_criteria:
assert seq_lens_cuda is not None
if requests:
beam_search_store = self.store.beam_search_store
assert self._use_beam_search == (beam_search_store is not None)
# Prepare stop word handling
Expand Down Expand Up @@ -3895,7 +3867,6 @@ def sample_async(
),
sampler_event=sampler_event,
beam_history_builders=beam_history_builders,
use_host_stop_criteria=use_host_stop_criteria,
)

@staticmethod
Expand Down Expand Up @@ -4679,12 +4650,7 @@ def _process_requests(
new_tokens_cuda: torch.Tensor,
num_context_logits_prefix_sum: list[int],
) -> tuple[
list[LlmRequest],
torch.Tensor,
torch.Tensor,
torch.Tensor | None,
torch.Tensor,
bool,
list[LlmRequest], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
]:
raw_logits_cuda = model_outputs["logits"]

Expand All @@ -4697,36 +4663,25 @@ def _process_requests(
if return_log_probs:
self._prepare_log_probs(sampling_requests)

use_fast_greedy_path = self._can_use_fast_greedy_path(sampling_requests)
use_host_stop_criteria = use_fast_greedy_path and self._can_use_host_stop_criteria(
sampling_requests
)

seq_slots_host = torch.tensor(
[r.py_seq_slot for r in sampling_requests],
dtype=torch.int32,
pin_memory=prefer_pinned(),
)

# necessary for beam search and max_length checks
seq_lens_host = (
None
if use_host_stop_criteria
else torch.tensor(
[r.max_beam_num_tokens for r in sampling_requests],
dtype=torch.int32,
pin_memory=prefer_pinned(),
)
seq_lens_host = torch.tensor(
[r.max_beam_num_tokens for r in sampling_requests],
dtype=torch.int32,
pin_memory=prefer_pinned(),
)

# Cast seq_slots / seq_lens to CUDA exactly once. The fast host stop-
# criteria path only needs seq_slots for scattering sampled tokens;
# all other paths also consume seq_lens in device-side finish handling
# or beam-search metadata.
# Cast seq_slots / seq_lens to CUDA exactly once; consumed by both
# the per-group beam-search metadata builder and the finish-reasons
# handler in sample_async. int64 is required for the index_*_ ops
# downstream.
seq_slots_cuda = seq_slots_host.to(device="cuda", dtype=torch.int64, non_blocking=True)
seq_lens_cuda = (
None if seq_lens_host is None else seq_lens_host.to(device="cuda", non_blocking=True)
)
seq_lens_cuda = seq_lens_host.to(device="cuda", non_blocking=True)

# Handle embedding bias
self._apply_embedding_bias(
Expand All @@ -4741,25 +4696,19 @@ def _process_requests(
)

# Fast path for greedy sampling
if use_fast_greedy_path:
if use_host_stop_criteria:
# There is exactly one token and one beam per request, so the
# linearized destination indices are the sequence slots.
batch_dest_indices_cuda = seq_slots_cuda
else:
# Compute destination indices on CPU (same pattern as
# _unbatch_sampling_results).
batch_destination_indexer = _UnpackedStepIndexer(
seq_slots=seq_slots_host,
num_steps=sampling_requests_metadata.req_num_generated_tokens,
steps_dim_size=new_tokens_cuda.size(0),
slots_dim_size=new_tokens_cuda.size(1),
dim_order=_UnpackedStepIndexer.DimOrder.STEP_MAJOR,
index_dtype=torch.int64,
)
batch_dest_indices_cuda = batch_destination_indexer[:].to(
new_tokens_cuda.device, non_blocking=True
)
if self._can_use_fast_greedy_path(sampling_requests):
# Compute destination indices on CPU (same pattern as _unbatch_sampling_results)
batch_destination_indexer = _UnpackedStepIndexer(
seq_slots=seq_slots_host,
num_steps=sampling_requests_metadata.req_num_generated_tokens,
steps_dim_size=new_tokens_cuda.size(0),
slots_dim_size=new_tokens_cuda.size(1),
dim_order=_UnpackedStepIndexer.DimOrder.STEP_MAJOR,
index_dtype=torch.int64,
)
batch_dest_indices_cuda = batch_destination_indexer[:].to(
new_tokens_cuda.device, non_blocking=True
)

# Get d2t tensor if present
d2t = model_outputs.get("d2t", None)
Expand All @@ -4777,10 +4726,10 @@ def _process_requests(
return (
sampling_requests,
seq_slots_host,
seq_lens_host,
seq_slots_cuda,
seq_lens_cuda,
new_tokens_host,
use_host_stop_criteria,
)

# Indexer for accessing tokens in 'logits_cuda', corresponding to the
Expand All @@ -4793,8 +4742,6 @@ def _process_requests(
)

# Perform sampling in batches
assert seq_lens_host is not None
assert seq_lens_cuda is not None
batched_sampling_result = self._sample_batched_by_strategy(
logits_cuda,
sampling_requests,
Expand Down Expand Up @@ -4832,10 +4779,10 @@ def _process_requests(
return (
sampling_requests,
seq_slots_host,
seq_lens_host,
seq_slots_cuda,
seq_lens_cuda,
new_tokens_host,
use_host_stop_criteria,
)

@override
Expand Down
101 changes: 0 additions & 101 deletions tests/unittest/_torch/sampler/test_torch_sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -866,107 +866,6 @@ def test_write_finish_reasons(cls):

run_test_with_warmup(uut_provider, max_sync_s=0.5)

@pytest.mark.parametrize(
"sampled_token,end_id,max_new_tokens,expected_reason",
[
pytest.param(42, 99, 5, None, id="not-finished"),
pytest.param(99, 99, 5, FinishReason.END_ID, id="end-id"),
pytest.param(42, 99, 1, FinishReason.LENGTH, id="length"),
pytest.param(
99,
99,
1,
FinishReason.END_ID,
id="end-id-precedes-length",
),
],
)
def test_host_stop_criteria_fast_path(
self,
mocker,
sampled_token: int,
end_id: int,
max_new_tokens: int,
expected_reason: FinishReason | None,
):
sampler = TorchSampler(
TorchSampler.Args(
max_seq_len=20,
max_draft_len=0,
max_total_draft_tokens=0,
max_num_sequences=1,
max_beam_width=1,
)
)
request = LlmRequest(
request_id=0,
seq_slot=0,
input_tokens=[1, 2],
max_new_tokens=max_new_tokens,
end_id=end_id,
sampling_config=SamplingConfig(),
is_streaming=False,
)

setup_requests = ScheduledRequests()
setup_requests.context_requests_last_chunk = [request]
sampler.setup_sampler_step(setup_requests)

scheduled_requests = ScheduledRequests()
scheduled_requests.generation_requests = [request]
logits = torch.full((1, 128), -1.0, dtype=torch.float32, device="cuda")
logits[0, sampled_token] = 1.0

finish_by = mocker.spy(request, "finish_by")
write_finish_reasons = mocker.patch.object(
sampler._finish_reasons_handler,
"write_finish_reasons",
side_effect=AssertionError("device finish reasons should be skipped"),
)

state = sampler.sample_async(
scheduled_requests,
model_outputs={"logits": logits},
num_context_logits_prefix_sum=[0],
)
assert state.use_host_stop_criteria
assert state.host is not None
assert state.host.finish_reasons is None

sampler.update_requests(state)

write_finish_reasons.assert_not_called()
assert request.get_tokens(0)[-1] == sampled_token
if expected_reason is None:
finish_by.assert_not_called()
else:
finish_by.assert_called_once_with(expected_reason, 0)

@pytest.mark.parametrize(
"sample_request",
[
pytest.param(
SimpleNamespace(py_is_draft=False, py_stop_words_list=[[42], [1]]),
id="stop-words",
),
pytest.param(
SimpleNamespace(py_is_draft=True, py_stop_words_list=None),
id="draft-request",
),
],
)
def test_host_stop_criteria_fast_path_fallback(self, sample_request):
sampler = TorchSampler(
TorchSampler.Args(
max_seq_len=20,
max_draft_len=0,
max_total_draft_tokens=0,
max_num_sequences=1,
max_beam_width=1,
)
)
assert not sampler._can_use_host_stop_criteria([cast(LlmRequest, sample_request)])

@classmethod
def test_are_stop_words_isnt_called_when_no_stop_words(cls, monkeypatch: pytest.MonkeyPatch):
"""We don't want to call are_stop_words when there are no stop words because it's expensive"""
Expand Down
Loading