From 06e5020ba665b73de3436ac04505692eb3d7fb4b Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Jun 2026 05:12:44 -0500 Subject: [PATCH 1/2] Add logprobs_mode (raw/processed) to inference Signed-off-by: Teodor-Dumitru Ene --- megatron/core/inference/config.py | 16 +++ .../inference/contexts/dynamic_context.py | 49 +++++++- megatron/core/inference/sampling/base.py | 18 ++- .../inference/sampling/flashinfer_sampling.py | 17 +++ .../core/inference/sampling/torch_sampling.py | 112 ++++++++++++------ .../text_generation_controller.py | 1 + megatron/inference/utils.py | 1 + megatron/training/arguments.py | 6 + 8 files changed, 182 insertions(+), 38 deletions(-) diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index 46d87baee97..991fe0bdf71 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -344,6 +344,9 @@ class InferenceConfig: sampling_backend: Literal['torch', 'flashinfer'] = 'torch' """Which sampling kernels to use during inference.""" + logprobs_mode: Literal['raw_logprobs', 'processed_logprobs'] = 'raw_logprobs' + """Whether returned log-probs are modified by the sampling parameters or not.""" + request_metadata_types: Optional[List[Tuple[str, torch.dtype]]] = None """ A list of the per-request metadata types to track. Each entry is a tuple @@ -387,6 +390,19 @@ def __post_init__(self, verbose: bool): f"got {self.prefix_caching_routing_alpha}" ) + if self.logprobs_mode not in ("raw_logprobs", "processed_logprobs"): + raise ValueError( + f"Unsupported logprobs_mode {self.logprobs_mode!r}. " + "Supported modes: raw_logprobs, processed_logprobs." + ) + + # The speculative log-probs path does not yet apply processed-logprobs. + if self.logprobs_mode == "processed_logprobs" and self.num_speculative_tokens > 0: + raise ValueError( + "logprobs_mode='processed_logprobs' is not yet supported with speculative decoding " + "(num_speculative_tokens > 0)." + ) + if self.sampling_backend == 'flashinfer': try: import flashinfer # noqa: F401 diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index e2524a9834a..24c8e56d08a 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -22,6 +22,7 @@ PrefixCachingEvictionPolicy, ) from megatron.core.inference.inference_request import DynamicInferenceRequest +from megatron.core.inference.sampling.base import Sampling from megatron.core.inference.sampling_params import SamplingParams from megatron.core.inference.unified_memory import ( UnifiedMemoryUnsupportedError, @@ -3772,8 +3773,38 @@ def update_requests( "evict_request_ids": evict_request_ids, } + def _processed_log_probs( + self, + logits: Tensor, + n_active: int, + active_query_lengths: Optional[Tensor], + sampling: Optional[Sampling], + ) -> Tensor: + """Sample the logprobs if desired.""" + if self.config.logprobs_mode == "raw_logprobs": + return F.log_softmax(logits, dim=-1) + + assert sampling is not None, "processed_logprobs requires a sampling backend" + + # Map each logits row to its active request. + request_idx = torch.arange(n_active, device=logits.device) + row_to_request = ( + request_idx + if active_query_lengths is None + else request_idx.repeat_interleave(active_query_lengths) + ) + md = self.active_request_metadata + temperature = md["temperature"][:n_active].to(logits.device, torch.float32)[row_to_request] + top_k = md["top_k"][:n_active].to(logits.device, torch.long)[row_to_request] + top_p = md["top_p"][:n_active].to(logits.device, torch.float32)[row_to_request] + return sampling.log_probs_kernel(logits, temperature, top_k, top_p) + def calculate_log_probs( - self, logits: Tensor, new_tokens: Tensor, only_last_token_logits: Optional[bool] = False + self, + logits: Tensor, + new_tokens: Tensor, + only_last_token_logits: Optional[bool] = False, + sampling: Optional[Sampling] = None, ) -> Tuple[List[List[float]], Tensor]: """Calculate log probs for all active requests and return them. @@ -3783,6 +3814,7 @@ def calculate_log_probs( logits (Tensor): Raw model output logits with shape [1, sequence_length, vocab_size]. new_tokens (Tensor): The newly sampled tokens. only_last_token_logits (bool): If set, the logits are from only the last token in each request + sampling (Optional[Sampling]): Backend used to optionally modify log-probs. Returns: List of lists where each inner list contains log probs for a request in the @@ -3792,14 +3824,19 @@ def calculate_log_probs( # Calculate log_probs (sequence_length x vocab_size) logits_squeezed = logits.squeeze(0).float() + n_active = self.total_request_count - self.paused_request_count if only_last_token_logits or self.is_decode_only(): seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device) - log_probs = F.log_softmax(logits_squeezed[seq_idx], dim=-1) + log_probs = self._processed_log_probs( + logits_squeezed[seq_idx], + n_active, + None, + sampling, + ) selected_log_probs = log_probs[seq_idx, new_tokens] return [[lp] for lp in selected_log_probs.tolist()], log_probs - log_probs = F.log_softmax(logits_squeezed, dim=-1) # Get the selected token ids for all tokens. # We shift the active token window left by one to remove the first prompt token for # prefill requests and then set the token ids explicitly for the newly generated tokens. @@ -3823,13 +3860,17 @@ def calculate_log_probs( # # active_token_ids[new_token_idx] = new_tokens # : [ 52 | 12 | 16 3 | 12 72 24 88 86 ] - n_active = self.total_request_count - self.paused_request_count active_token_ids = self.gpu_view.token_to_input_ids[: self.active_token_count].roll(-1, 0) active_query_lengths = self.gpu_view.request_query_lengths[:n_active] new_token_idx = active_query_lengths.cumsum(0) - 1 active_token_ids[new_token_idx] = new_tokens + # Compute (possibly processed) log-probs over all active-token rows. + log_probs = self._processed_log_probs( + logits_squeezed, n_active, active_query_lengths, sampling + ) + # Extract the log probs for only the selected tokens. # (sequence_length x vocab_size) -> (sequence_length) seq_idx = torch.arange(self.active_token_count, device=log_probs.device) diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py index 8aa4c416c27..dceebb060a8 100644 --- a/megatron/core/inference/sampling/base.py +++ b/megatron/core/inference/sampling/base.py @@ -10,7 +10,8 @@ class Sampling(ABC): """Abstract base for inference sampling backends. - Subclasses implement `sample_kernel`. CUDA graphs are added via `CudaGraphManager`. + Subclasses implement `sample_kernel` and `log_probs_kernel`. + CUDA graphs are added via `CudaGraphManager`. """ @abstractmethod @@ -87,3 +88,18 @@ def sample_speculative( token_to_request_index=token_to_request_index, eager=True, ) + + @abstractmethod + def log_probs_kernel( + self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor + ) -> Tensor: + """Per-row log-probs of the distribution this backend samples from. + + Args: + logits: `[num_rows, vocab_size]` raw logits. + temperature, top_k, top_p: `[num_rows]` per-row sampling params. + + Returns: + `[num_rows, vocab_size]` log-probs; filtered-out tokens are `-inf`. + """ + ... diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py index c89093daeac..f7b85a8836e 100644 --- a/megatron/core/inference/sampling/flashinfer_sampling.py +++ b/megatron/core/inference/sampling/flashinfer_sampling.py @@ -99,3 +99,20 @@ def sample_kernel( ) ) return output + + def log_probs_kernel( + self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor + ) -> Tensor: + """Per-row log-probs of the FlashInfer top-k / top-p sampling distribution.""" + temperature = temperature.clamp(min=1e-6) + probs = torch.softmax(logits / temperature.unsqueeze(1), dim=-1) + + # Sentinel values disable filtering: + # top_k=vocab_size keeps all tokens, top_p=1.0 keeps the full probability mass. + top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size) + top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0) + + # Renormalize to the kept set (top-k first, then top-p) to match + renormed = flashinfer.sampling.top_k_renorm_probs(probs, top_k_safe) + renormed = flashinfer.sampling.top_p_renorm_probs(renormed, top_p_safe) + return torch.log(renormed) diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py index 79491add5ab..f7f6f8cb662 100644 --- a/megatron/core/inference/sampling/torch_sampling.py +++ b/megatron/core/inference/sampling/torch_sampling.py @@ -19,6 +19,56 @@ def __init__(self, rng: torch.Generator, vocab_size: int) -> None: self._rng = rng self._vocab_size = vocab_size + @staticmethod + def _modify_logits_for_top_k_filtering(logits: Tensor, top_k: int) -> None: + """In-place: set logits outside the top-k set to -inf.""" + filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits.masked_fill_(filter_, float("-Inf")) + + @staticmethod + def _modify_logits_for_top_p_filtering(logits: Tensor, top_p: float) -> None: + """In-place: set logits outside the top-p (nucleus) set to -inf.""" + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + + filter_ = cumulative_probs > top_p + # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views; + # without clone, each write would corrupt the next read during the shift. + filter_[:, 1:] = filter_[:, :-1].clone() + filter_[..., 0] = 0 + + filter_ = filter_.scatter(1, sorted_indices, filter_) + logits.masked_fill_(filter_, float("-Inf")) + + @staticmethod + def filter_logits( + last_token_logits: Tensor, + temperature: float, + top_k: int, + top_p: float, + *, + vocab_size: Optional[int] = None, + ) -> Tensor: + """Temperature-scale then top-k/top-p filter logits; filtered entries become -inf. + + Returns a new tensor (input unmodified). Shared by `sample_from_logits` and + `log_probs_kernel` so sampling and processed log-probs apply the same filter. + """ + assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero" + assert top_p <= 1.0, "top-p should be in (0,1]" + # Clone needed: .div_() and the filters below modify in-place. + last_token_logits = last_token_logits.clone() + if temperature != 1.0: + last_token_logits.div_(temperature) + if top_k >= 1: + assert top_k <= last_token_logits.size(1), "top-k is larger than logit size." + if vocab_size: + assert top_k < vocab_size, "top-k is larger than vocab size." + TorchSampling._modify_logits_for_top_k_filtering(last_token_logits, top_k) + elif top_p > 0.0: + TorchSampling._modify_logits_for_top_p_filtering(last_token_logits, top_p) + return last_token_logits + @staticmethod def sample_from_logits( last_token_logits: Tensor, @@ -49,42 +99,13 @@ def sample_from_logits( assert isinstance(top_k, int) assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero" assert top_p <= 1.0, "top-p should be in (0,1]" - - def modify_logits_for_top_k_filtering(logits, top_k): - """Set the logits for none top-k values to -inf.""" - filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None] - logits.masked_fill_(filter_, float("-Inf")) - - def modify_logits_for_top_p_filtering(logits, top_p): - """Set the logits for none top-p values to -inf.""" - sorted_logits, sorted_indices = torch.sort(logits, descending=True) - cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) - - filter_ = cumulative_probs > top_p - # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views; - # without clone, each write would corrupt the next read during the shift. - filter_[:, 1:] = filter_[:, :-1].clone() - filter_[..., 0] = 0 - - filter_ = filter_.scatter(1, sorted_indices, filter_) - logits.masked_fill_(filter_, float("-Inf")) - if top_k == 1: return torch.argmax(last_token_logits, dim=-1) - # Clone needed: .div_() and masked_fill_() below modify in-place. - last_token_logits = last_token_logits.clone() - if temperature != 1.0: - last_token_logits.div_(temperature) - if top_k > 1: - assert top_k <= last_token_logits.size(1), "top-k is larger than logit size." - if vocab_size: - assert top_k < vocab_size, "top-k is larger than vocab size." - modify_logits_for_top_k_filtering(last_token_logits, top_k) - elif top_p > 0.0: - modify_logits_for_top_p_filtering(last_token_logits, top_p) - - probabilities = last_token_logits.softmax(dim=-1) + filtered = TorchSampling.filter_logits( + last_token_logits, temperature, top_k, top_p, vocab_size=vocab_size + ) + probabilities = filtered.softmax(dim=-1) sampled = torch.multinomial(probabilities, num_samples=1, generator=generator).view(-1) if vocab_size: @@ -92,6 +113,31 @@ def modify_logits_for_top_p_filtering(logits, top_p): return sampled + def log_probs_kernel( + self, logits: Tensor, temperature: Tensor, top_k: Tensor, top_p: Tensor + ) -> Tensor: + """Per-row log-probs of the temperature, top-k/top-p sampling distribution. + + Buckets rows by identical (temperature, top_k, top_p) and reuses `filter_logits` + (the same filter as `sample_from_logits`) so log-probs match how this backend + samples. `temperature`/`top_k`/`top_p` are per-row `[num_rows]` tensors. + """ + temps = temperature.tolist() + top_ks = top_k.tolist() + top_ps = top_p.tolist() + buckets: dict = defaultdict(list) + for row, key in enumerate(zip(temps, top_ks, top_ps)): + buckets[key].append(row) + + log_probs = torch.empty_like(logits) + for (t, k, p), rows in buckets.items(): + idx = torch.tensor(rows, device=logits.device, dtype=torch.long) + filtered = TorchSampling.filter_logits( + logits[idx], float(t), int(k), float(p), vocab_size=self._vocab_size + ) + log_probs[idx] = torch.log_softmax(filtered, dim=-1) + return log_probs + def sample_kernel( self, logits: Tensor, diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index b252e013250..6b75c4685ac 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -1157,6 +1157,7 @@ def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]: self._all_logits_cuda[:, :logits_seq_len, :], self._sampled_tokens_cuda[:active_request_count], only_last_token_logits=context.config.materialize_only_last_token_logits, + sampling=self._sampling, ) def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float]], Tensor]: diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index 91a9d954617..567d48ffc3b 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -382,6 +382,7 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args): use_synchronous_zmq_collectives=args.inference_use_synchronous_zmq_collectives, disable_ep_consensus=args.inference_disable_ep_consensus, sampling_backend=args.inference_dynamic_batching_sampling_backend, + logprobs_mode=args.inference_dynamic_batching_logprobs_mode, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index fdca8afe000..14105755574 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2004,6 +2004,12 @@ def _add_inference_args(parser): help='Which sampling kernels to use during inference. ' 'Falls back to "torch" with a warning if "flashinfer" ' 'is requested but the package is not installed.') + group.add_argument('--inference-dynamic-batching-logprobs-mode', + type=str, default='raw_logprobs', + choices=['raw_logprobs', 'processed_logprobs'], + help='How returned inference log-probs are computed engine-wide. ' + '"raw_logprobs" (default) uses the unmodified model logits; ' + '"processed_logprobs" uses temperature and filters by top-k/top-p.') group.add_argument('--inference-logging-step-interval', type=int, default=0, help='Step interval for logging inference metrics. ' 'Default to 0 to disable inference logging.') From e31fde1e3f47664e3d51056edd8efa48ebfd3b45 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 22 Jun 2026 18:26:40 -0500 Subject: [PATCH 2/2] Address reviewer comments Signed-off-by: Teodor-Dumitru Ene --- .../inference/contexts/dynamic_context.py | 5 +- .../contexts/test_dynamic_context.py | 85 ++++++++++++++----- 2 files changed, 64 insertions(+), 26 deletions(-) diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 24c8e56d08a..0fca69628a5 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -3829,10 +3829,7 @@ def calculate_log_probs( if only_last_token_logits or self.is_decode_only(): seq_idx = torch.arange(len(new_tokens), dtype=torch.int32, device=logits.device) log_probs = self._processed_log_probs( - logits_squeezed[seq_idx], - n_active, - None, - sampling, + logits_squeezed[seq_idx], n_active, None, sampling ) selected_log_probs = log_probs[seq_idx, new_tokens] return [[lp] for lp in selected_log_probs.tolist()], log_probs diff --git a/tests/unit_tests/inference/contexts/test_dynamic_context.py b/tests/unit_tests/inference/contexts/test_dynamic_context.py index e79df3aaebf..499b89398fe 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_context.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_context.py @@ -15,6 +15,7 @@ TokenOverflowError, ) from megatron.core.inference.inference_request import DynamicInferenceRequest +from megatron.core.inference.sampling.torch_sampling import TorchSampling from megatron.core.inference.sampling_params import SamplingParams from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed @@ -1076,7 +1077,8 @@ def test_mamba_states_cache(self, is_hybrid_model: bool): @pytest.mark.internal @rounder_override(64) - def test_calculate_and_store_log_probs(self): + @pytest.mark.parametrize("logprobs_mode", ["raw_logprobs", "processed_logprobs"]) + def test_calculate_and_store_log_probs(self, logprobs_mode): dynamic_context = self._get_dynamic_context( params_dtype=torch.float32, @@ -1088,23 +1090,27 @@ def test_calculate_and_store_log_probs(self): block_size_tokens=128, max_tokens=None, ) + dynamic_context.config.logprobs_mode = logprobs_mode - # Add a few requests to the context + # Add a few requests to the context, each with its own sampling parameters. request_data = { 1001: { "tokens": torch.randint(0, 100, (10,), device='cpu'), "prefill_len": 10, "initial_token_offset": 0, + "sampling": dict(temperature=1.0, top_k=0, top_p=0.0), # raw-equivalent }, 1002: { "tokens": torch.randint(0, 100, (5,), device='cpu'), "prefill_len": 5, "initial_token_offset": 10, + "sampling": dict(temperature=0.5, top_k=0, top_p=0.0), # temperature }, 1003: { "tokens": torch.randint(0, 100, (7,), device='cpu'), "prefill_len": 7, "initial_token_offset": 15, + "sampling": dict(temperature=1.0, top_k=8, top_p=0.0), # top-k }, } @@ -1115,7 +1121,8 @@ def test_calculate_and_store_log_probs(self): request_id=req_id, prompt_tokens=data["tokens"], sampling_params=SamplingParams( - num_tokens_to_generate=dynamic_context.max_tokens - len(data["tokens"]) + num_tokens_to_generate=dynamic_context.max_tokens - len(data["tokens"]), + **data["sampling"], ), ) ) @@ -1127,6 +1134,32 @@ def test_calculate_and_store_log_probs(self): total_active_tokens = dynamic_context.active_token_count vocab_size = 50000 + # Supplies log_probs_kernel for processed mode (unused by raw mode). + sampling = TorchSampling(rng=torch.Generator(), vocab_size=vocab_size) + + def expected_log_probs(logits, active_id_and_counts): + """Mode-aware expected log-probs over every active-token row. + + For processed mode, each active request's params are repeated across its token + count, mirroring the request->row mapping in `_processed_log_probs`. + """ + logits_2d = logits.squeeze(0).float() + if logprobs_mode == "raw_logprobs": + return torch.nn.functional.log_softmax(logits_2d, dim=-1) + temperatures, top_ks, top_ps = [], [], [] + for active_id, count in active_id_and_counts: + sp = request_data[active_id]["sampling"] + temperatures += [sp["temperature"]] * count + top_ks += [sp["top_k"]] * count + top_ps += [sp["top_p"]] * count + device = logits_2d.device + return sampling.log_probs_kernel( + logits_2d, + torch.tensor(temperatures, device=device, dtype=torch.float32), + torch.tensor(top_ks, device=device, dtype=torch.long), + torch.tensor(top_ps, device=device, dtype=torch.float32), + ) + # Populate gpu_view for calculate_log_probs (which reads from gpu_view). dynamic_context.initialize_attention_state() dynamic_context.transfer_bookkeeping_to_gpu() @@ -1143,16 +1176,15 @@ def test_calculate_and_store_log_probs(self): prefill_new_tokens = torch.randint(0, 100, (num_active_requests,), device='cuda').long() # Call the function for prefill - prefill_log_probs, _ = dynamic_context.calculate_log_probs( - prefill_logits, prefill_new_tokens + prefill_log_probs, prefill_log_probs_full = dynamic_context.calculate_log_probs( + prefill_logits, prefill_new_tokens, sampling=sampling ) # Calculate expected prefill log probs for the selected tokens - expected_prefill_log_probs = ( - torch.nn.functional.log_softmax(prefill_logits.squeeze(0), dim=-1) - .to(torch.float32) - .cpu() - ) + prefill_active = [(req_id, request_data[req_id]["prefill_len"]) for req_id in request_data] + expected_prefill_full = expected_log_probs(prefill_logits, prefill_active) + assert torch.allclose(prefill_log_probs_full, expected_prefill_full, atol=1e-6) + expected_prefill_log_probs = expected_prefill_full.to(torch.float32).cpu() for i, (req_id, data) in enumerate(request_data.items()): req_len = data["tokens"].shape[0] @@ -1187,12 +1219,15 @@ def test_calculate_and_store_log_probs(self): 1, num_active_requests, vocab_size, device='cuda', dtype=torch.float32 ) decode_new_tokens = torch.randint(0, 100, (num_active_requests,), device='cuda').long() - decode_log_probs, _ = dynamic_context.calculate_log_probs(decode_logits, decode_new_tokens) + decode_log_probs, decode_log_probs_full = dynamic_context.calculate_log_probs( + decode_logits, decode_new_tokens, sampling=sampling + ) # Verify the stored decode log probabilities - expected_decode_log_probs = torch.nn.functional.log_softmax( - decode_logits.squeeze(0), dim=-1 - ).to(torch.float32) + decode_active = [(req_id, 1) for req_id in request_data] + expected_decode_full = expected_log_probs(decode_logits, decode_active) + assert torch.allclose(decode_log_probs_full, expected_decode_full, atol=1e-6) + expected_decode_log_probs = expected_decode_full.to(torch.float32) for i, (req_id, data) in enumerate(request_data.items()): assert len(decode_log_probs[i]) == 1, len(decode_log_probs[i]) @@ -1210,12 +1245,14 @@ def test_calculate_and_store_log_probs(self): new_request_tokens = torch.randint(0, 100, (12,), device='cpu').long() new_request_prefill_len = new_request_tokens.shape[0] initial_token_offset_new_request = dynamic_context.active_token_count + new_request_sampling = dict(temperature=1.0, top_k=0, top_p=0.8) # top-p dynamic_context.add_request( DynamicInferenceRequest( request_id=new_request_id, prompt_tokens=new_request_tokens, sampling_params=SamplingParams( - num_tokens_to_generate=dynamic_context.max_tokens - len(new_request_tokens) + num_tokens_to_generate=dynamic_context.max_tokens - len(new_request_tokens), + **new_request_sampling, ), ) ) @@ -1223,6 +1260,7 @@ def test_calculate_and_store_log_probs(self): "tokens": new_request_tokens, "prefill_len": new_request_prefill_len, "initial_token_offset": initial_token_offset_new_request, + "sampling": new_request_sampling, } # Simulate the step after adding the new prefill request. @@ -1243,15 +1281,18 @@ def test_calculate_and_store_log_probs(self): 0, 100, (num_active_requests_mixed_step,), device='cuda' ).long() - mixed_step_log_probs, _ = dynamic_context.calculate_log_probs( - mixed_step_logits, mixed_step_new_tokens + mixed_step_log_probs, mixed_step_log_probs_full = dynamic_context.calculate_log_probs( + mixed_step_logits, mixed_step_new_tokens, sampling=sampling ) - expected_mixed_step_log_probs = ( - torch.nn.functional.log_softmax(mixed_step_logits.squeeze(0), dim=-1) - .to(torch.float32) - .cpu() - ) + # Existing requests are in decode (1 token each); the new request is in prefill. + mixed_active = [ + (req_id, request_data[req_id]["prefill_len"] if req_id == new_request_id else 1) + for req_id in request_data + ] + expected_mixed_full = expected_log_probs(mixed_step_logits, mixed_active) + assert torch.allclose(mixed_step_log_probs_full, expected_mixed_full, atol=1e-6) + expected_mixed_step_log_probs = expected_mixed_full.to(torch.float32).cpu() # Verify log probs for the mixed step current_global_token_offset = 0