Skip to content

[None][fix] Fix access to new tokens in sampler. #7958

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
merged 6 commits into from
Oct 2, 2025
Merged

[None][fix] Fix access to new tokens in sampler. #7958

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cd02b6b
Fix access to new tokens.
dcampora Sep 24, 2025
8ee9543
Formatting.
dcampora Sep 24, 2025
3f50a09
Write back to tensor in sample rejection cases.
dcampora Sep 24, 2025
020d9a7
Fix add_new_tokens.
dcampora Sep 30, 2025
7cafec5
Fix formatting.
dcampora Sep 30, 2025
a5e200e
Merge branch 'main' into user/dcampora/fix_sampler_draft_bug
dcampora Oct 2, 2025
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
26 changes: 16 additions & 10 deletions tensorrt_llm/_torch/pyexecutor/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1068,7 +1068,8 @@ def _tree_sampling_batch(self, requests: list[LlmRequest],
return new_draft_tokens_host

def _process_draft_tokens_rejection_sampling(
self, request: LlmRequest, new_tokens: torch.Tensor) -> int:
self, request: LlmRequest, new_tokens: list[list[list[int]]],
new_tokens_tensor: torch.Tensor) -> int:
sampling_strategy = _request_strategy(request)
generator = self.get_generator(request.py_draft_logits.device)
_, draft_probs = sample(sampling_strategy,
Expand All @@ -1089,28 +1090,32 @@ def _process_draft_tokens_rejection_sampling(
num_accepted = num_initially_accepted
for i in range(num_accepted):
new_token = request.py_draft_tokens[i]
new_tokens[i, request.seq_slot, self.BEAM] = new_token
new_token = add_token(request, new_tokens, beam=self.BEAM, step=i)
new_tokens_tensor[i, request.seq_slot, self.BEAM] = new_token
request.add_new_token(new_token, self.BEAM)
stop = self._handle_stop_criteria(request, new_token)
if stop:
num_accepted = i + 1
return num_accepted
if sample_last:
new_token = sample_rejected(draft_probs, target_probs, generator,
num_accepted)
new_tokens[num_accepted, request.seq_slot, self.BEAM] = new_token
new_token = add_token(request,
new_tokens,
beam=self.BEAM,
step=num_accepted)
new_tokens_tensor[num_accepted, request.seq_slot,
self.BEAM] = new_token
request.add_new_token(new_token, self.BEAM)
else:
new_token = add_token(request,
new_tokens,
beam=self.BEAM,
step=num_accepted)
stop = self._handle_stop_criteria(request, new_token)

return num_accepted

def process_draft_tokens(
self,
request: LlmRequest,
new_tokens: torch.Tensor,
new_tokens: list[list[list[int]]],
new_tokens_tensor: torch.Tensor,
resource_manager: Optional[ResourceManager] = None) -> int:
if request.py_draft_logits is None:
spec_tree_manager = self.get_spec_tree_manager(resource_manager)
Expand All @@ -1125,7 +1130,7 @@ def process_draft_tokens(
return num_accepted
else:
return self._process_draft_tokens_rejection_sampling(
request, new_tokens)
request, new_tokens, new_tokens_tensor)

@override
def update_requests(
Expand All @@ -1151,6 +1156,7 @@ def update_requests(
continue
processed = 1
num_accepted = self.process_draft_tokens(req, new_tokens,
state.host.new_tokens,
resource_manager)
if get_draft_token_length(req) > 0:
req.py_num_accepted_draft_tokens = num_accepted
Expand Down