Skip to content

[https://nvbugs/5661877][fix] fix test regression in TestBatchedSampling::test_samples #9215

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
ixlmar merged 1 commit into from
Nov 19, 2025
Merged

[https://nvbugs/5661877][fix] fix test regression in TestBatchedSampling::test_samples #9215

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
34 changes: 28 additions & 6 deletions tensorrt_llm/_torch/pyexecutor/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -278,11 +278,14 @@ def _group_requests_by_strategy_key(
group_dict: dict[tuple[GenericStrategyKeyType, bool], tuple[list[int], list[Strategy]]] = (
defaultdict(lambda: ([], []))
)

for req_index, req in enumerate(requests):
strategy = _request_strategy(req, vocab_size=vocab_size)
# In the overlap path, py_draft_logits is not updated yet,
# so we use get_draft_token_length() for the checking.
speculation_needs_probs = get_draft_token_length(req) > 0 and strategy is not GREEDY
speculation_needs_probs = (
# NB: This criterion needs to be consistent with the gating of rejection sampling in
# process_draft_tokens.
TorchSampler._speculation_could_use_rejection_sampling(req, strategy)
)
strategy_key = strategy_to_key(strategy, speculation_needs_probs)
group_dict_entry = group_dict[(strategy_key, speculation_needs_probs)]
group_dict_entry[0].append(req_index)
Expand Down Expand Up @@ -1026,6 +1029,17 @@ def _process_draft_tokens_rejection_sampling(

return num_accepted

@staticmethod
def _speculation_could_use_rejection_sampling(
request: LlmRequest, strategy: Optional[Strategy] = None
) -> bool:
if strategy is None:
strategy = _request_strategy(
request,
vocab_size=2**31, # vocab_size does not affect greediness
)
return get_draft_token_length(request) > 0 and strategy != GREEDY

def process_draft_tokens(
self,
request: LlmRequest,
Expand All @@ -1034,9 +1048,17 @@ def process_draft_tokens(
finish_reasons: FinishReasonsList,
resource_manager: Optional[ResourceManager] = None,
) -> int:
if (
_request_strategy(request, vocab_size=2**31) == GREEDY
or request.py_draft_logits is None
if not (
self._speculation_could_use_rejection_sampling(request)
# NB: '_speculation_could_use_rejection_sampling' is called in sample_async, which precludes
# inspection of .py_draft_logits, because it is not set yet when the overlap path
# is used.
#
# OTOH, some drafters (e.g. NGram) do not provide draft logits, precluding rejection
# sampling. The current solution accepts that .py_target_probs may sometimes be
# computed, even though .py_draft_logits may never be set and the target probs
# may ultimately not be required.
and request.py_draft_logits is not None
):
spec_tree_manager = self.get_spec_tree_manager(resource_manager)
if spec_tree_manager is not None:
Expand Down
1 change: 0 additions & 1 deletion tests/integration/test_lists/waives.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,6 @@ disaggregated/test_auto_scaling.py::test_disagg_server_restart[etcd-round_robin]
disaggregated/test_auto_scaling.py::test_worker_restart[http-load_balancing] SKIP (https://nvbugs/5649010)
disaggregated/test_auto_scaling.py::test_worker_restart[http-round_robin] SKIP (https://nvbugs/5649010)
disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/5661926)
unittest/_torch/sampler/test_torch_sampler.py::TestBatchedSampling SKIP (https://nvbugs/5661877)
test_e2e.py::test_ptp_quickstart_multimodal[mistral-small-3.1-24b-instruct-Mistral-Small-3.1-24B-Instruct-2503-mixture_text_image-True] SKIP (https://nvbugs/5568836)
test_e2e.py::test_trtllm_multimodal_benchmark_serving SKIP (https://nvbugs/5647825)
unittest/_torch/modules/test_fused_moe.py::test_fused_moe_alltoall_fp4[MNNVL] SKIP (https://nvbugs/5664904)
Expand Down
87 changes: 28 additions & 59 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 @@ -945,15 +945,11 @@ def seq_slot_assignment(
def mock_requests(
self,
sampling_params_list: list[SamplingParams],
with_draft_logits: bool,
vocab_size: int,
seq_slot_assignment: tuple[list[int], int],
draft_lens: list[int],
) -> ScheduledRequests:
return self._build_mock_requests(
sampling_params_list=sampling_params_list,
with_draft_logits=with_draft_logits,
vocab_size=vocab_size,
seq_slot_assignment=seq_slot_assignment,
draft_lens=draft_lens,
)
Expand All @@ -962,8 +958,6 @@ def _build_mock_requests(
self,
sampling_params_list: list[SamplingParams],
*,
with_draft_logits: bool,
vocab_size: int,
seq_slot_assignment: tuple[list[int], int],
draft_lens: list[int],
) -> ScheduledRequests:
Expand All @@ -975,21 +969,9 @@ def __init__(
self,
sampling_params_list: list[SamplingParams],
*,
with_draft_logits: bool,
draft_lens: list[int],
):
self._sampling_params_list = sampling_params_list
self._with_draft_logits = with_draft_logits

def _attach_draft_logits(req: LlmRequest) -> LlmRequest:
draft_len = len(req.py_draft_tokens)
if draft_len and with_draft_logits:
req.py_draft_logits = torch.testing.make_tensor( # type: ignore
(draft_len, vocab_size),
dtype=torch.float32,
device="cuda",
)
return req

# NB:
# - stop words are tested in test_write_finish_reasons
Expand All @@ -999,24 +981,22 @@ def _attach_draft_logits(req: LlmRequest) -> LlmRequest:
# - py_return_log_probs is tested elsewhere
# - code paths gated by py_return_context_logits tested in test_select_generated_logits
self._gen_requests = [
_attach_draft_logits(
LlmRequest(
request_id=seq_slot,
max_new_tokens=(2 * draft_len), # not used by tested code
input_tokens=[12], # not used by tested code
sampling_config=SamplingConfig(sampling_params._get_sampling_config()),
seq_slot=seq_slot,
is_streaming=False, # not relevant for tested code
draft_tokens=( # 'len(.py_draft_tokens)' is inspected by get_draft_token_length
torch.testing.make_tensor(
(draft_len,),
dtype=torch.int32,
device="cpu",
).tolist()
if draft_len
else None
),
)
LlmRequest(
request_id=seq_slot,
max_new_tokens=(2 * draft_len), # not used by tested code
input_tokens=[12], # not used by tested code
sampling_config=SamplingConfig(sampling_params._get_sampling_config()),
seq_slot=seq_slot,
is_streaming=False, # not relevant for tested code
draft_tokens=( # 'len(.py_draft_tokens)' is inspected by get_draft_token_length
torch.testing.make_tensor(
(draft_len,),
dtype=torch.int32,
device="cpu",
).tolist()
if draft_len
else None
),
)
for sampling_params, seq_slot, draft_len in zip(
sampling_params_list, seq_slots, draft_lens
Expand All @@ -1040,9 +1020,7 @@ def all_requests(self) -> list[LlmRequest]:
with torch.inference_mode(True):
return cast(
ScheduledRequests,
ScheduledRequestsMock(
sampling_params_list, with_draft_logits=with_draft_logits, draft_lens=draft_lens
),
ScheduledRequestsMock(sampling_params_list, draft_lens=draft_lens),
)

@pytest.fixture(scope="function")
Expand Down Expand Up @@ -1184,20 +1162,17 @@ def test_backend_selection(
"max_draft_len",
"draft_lens",
"sampling_params_list",
"with_draft_logits",
"params_label",
"allow_zero_draft_len",
"vocab_size",
),
[
# NB: with_draft_logits=True and non-zero draft len ensures that
# LlmRequest.py_target_probs is set.
# NB: non-zero draft len ensures that LlmRequest.py_target_probs is set.
pytest.param(
use_flashinfer,
3,
[3] * len(sampling_params_list),
sampling_params_list,
True,
params_label,
False,
vocab_size,
Expand Down Expand Up @@ -1225,7 +1200,6 @@ def test_probs(
allow_zero_draft_len: bool, # used by fixtures
sampling_params_list: list[SamplingParams],
seq_slot_assignment: tuple[list[int], int],
with_draft_logits: bool,
):
"""Validate probabilities returned by sample_async.

Expand Down Expand Up @@ -1255,9 +1229,7 @@ def _uut_provider(is_warmup: bool) -> Generator[Callable[[], None], None, None]:
# requests.
uut_mock_requests = self._build_mock_requests(
sampling_params_list=sampling_params_list,
vocab_size=vocab_size,
seq_slot_assignment=seq_slot_assignment,
with_draft_logits=with_draft_logits,
draft_lens=draft_lens,
)
else:
Expand Down Expand Up @@ -1427,11 +1399,8 @@ def _compute_probs(
)
mock_requests_with_probs = self._build_mock_requests(
sampling_params_list=sampling_params_list,
vocab_size=vocab_size,
seq_slot_assignment=seq_slot_assignment,
# NB: with_draft_logits=True and non-zero draft len ensures that
# LlmRequest.py_target_probs is set.
with_draft_logits=True,
# NB: non-zero draft len ensures that LlmRequest.py_target_probs is set.
draft_lens=([draft_len_with_probs] * len(sampling_params_list)),
)
# zero-pad logits to draft_len_with_probs
Expand Down Expand Up @@ -1818,6 +1787,12 @@ def _validate_token_frequencies(

# Perform G-test (asymptotically approximated by Pearson's chi-square test) to
# check that sampled tokens are consistent with the expected probs.
#
# NB: Need to use FP64 to avoid negative test statistic values.
test_token_counts_ma = test_token_counts_ma.astype(np.float64)
test_expected_counts_ma = test_expected_counts_ma.astype(np.float64)
test_expected_counts_ma /= test_expected_counts_ma.sum(axis=-1, keepdims=True)
test_expected_counts_ma *= num_samples
test_result = power_divergence(
f_obs=test_token_counts_ma,
f_exp=test_expected_counts_ma,
Expand Down Expand Up @@ -1847,7 +1822,6 @@ def _validate_token_frequencies(
"use_flashinfer",
"max_draft_len",
"sampling_params_list",
"with_draft_logits",
"allow_zero_draft_len",
"bypass_sampling",
"vocab_size",
Expand All @@ -1857,7 +1831,6 @@ def _validate_token_frequencies(
use_flashinfer,
max_draft_len,
sampling_params_list,
with_draft_logits,
allow_zero_draft_len,
# Run full sampling test only for uniform batches, with/without probs, but skip
# sampling statistics when varying draft lens etc. to validate batch handling:
Expand All @@ -1868,22 +1841,20 @@ def _validate_token_frequencies(
id=(
f"{'FlashInfer' if use_flashinfer else 'Torch'}"
f"-draft_len={0 if allow_zero_draft_len else 1}..{max_draft_len}"
f"-return_probs={with_draft_logits}-{params_label}"
f"-{params_label}"
),
)
# https://stackoverflow.com/a/75421799, does not work with nested loops
for (
use_flashinfer,
is_mixed,
with_draft_logits,
max_draft_len,
allow_zero_draft_len,
_build_test_cases,
vocab_size,
) in product(
[False, True],
[False, True],
[True, False],
[0, 3],
[False, True],
[_build_test_cases],
Expand All @@ -1895,8 +1866,7 @@ def _validate_token_frequencies(
include_uniform=(not is_mixed),
include_mixed=is_mixed,
)
if (allow_zero_draft_len or max_draft_len > 0)
and (not with_draft_logits or max_draft_len > 0)
if allow_zero_draft_len or max_draft_len > 0
],
)
def test_samples(
Expand All @@ -1908,7 +1878,6 @@ def test_samples(
vocab_size: int,
sampling_params_list: list[SamplingParams],
seq_slot_assignment: tuple[list[int], int],
with_draft_logits: bool,
max_draft_len: int,
use_flashinfer: bool,
allow_zero_draft_len: bool, # used by fixtures
Expand Down Expand Up @@ -2038,7 +2007,7 @@ def _uut(res=res):
probs = probs[: (draft_len + 1)]

# check probs are returned only when needed
should_return_probs = draft_len and with_draft_logits
should_return_probs = bool(draft_len)
assert (
hasattr(req, "py_target_probs") and req.py_target_probs is not None
) == should_return_probs
Expand Down