diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index e1a36ff1563..c4b092309c2 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -297,10 +297,10 @@ class InferenceConfig: Defaults to 0, which means no logging. """ - request_metadata_types: Optional[List[Tuple[str, torch.dtype, bool]]] = None + 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 - consisting of the string label, the target dtype, and whether to store the data on GPU. + consisting of the string label and the target dtype. """ use_synchronous_zmq_collectives: bool = False diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 5b264b36302..2a6e4fb5342 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -844,7 +844,7 @@ def initialize_all_tensors(self) -> None: label: torch.empty( (self.max_requests,), dtype=dtype, device=torch.cuda.current_device() ) - for label, dtype, _ in self.request_metadata_types + for label, dtype in self.request_metadata_types } # Per-token state. @@ -860,6 +860,11 @@ def initialize_all_tensors(self) -> None: self.token_to_position_in_request = torch.empty_like(self.token_to_input_ids) self.token_to_local_position_within_kv_block = torch.empty_like(self.token_to_input_ids) + # Static tensor addresses of active slices to enable fast inference kernels. + self.active_request_metadata = { + label: torch.empty_like(tensor) for label, tensor in self.request_metadata.items() + } + # NOTE: Need to build this outside the UVM / TMS context to avoid IMA. if self.is_hybrid_model: self.mamba_metadata = MambaMetadata( @@ -1062,6 +1067,23 @@ def get_active_request_count(self): """Returns the current number of active requests.""" return self.total_request_count - self.paused_request_count + def build_active_slices(self, batch_size: int): + """Build the active slices of specific tensors. This is run on every forward step. + + If the context is reordered to active -> paused -> finished, this can be graphed. + """ + padded_slice = slice(self.paused_request_count, self.paused_request_count + batch_size) + + # Request metadata all needs to be sliced. + for label in self.request_metadata: + self.active_request_metadata[label][:batch_size].copy_( + self.request_metadata[label][padded_slice], non_blocking=True + ) + + def pad_active_slices(self): + """Pad the active slices of specific tensors.""" + pass + def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None: """Append to KV cache. @@ -1391,7 +1413,7 @@ def add_dummy_requests_parallel( self.request_output_lengths[request_slice] = lengths_tensor + tokens_to_generate_tensor self.request_kv_length_offsets[request_slice] = 0 self.request_kv_block_counts[request_slice] = block_counts - for i, (label, dtype, _) in enumerate(self.request_metadata_types): + for i, (label, dtype) in enumerate(self.request_metadata_types): self.request_metadata[label][request_slice] = torch.tensor( metadata_cols[i], dtype=dtype, device=torch.cuda.current_device() ) @@ -1702,6 +1724,9 @@ def initialize_attention_state( self.padded_active_request_count = self.padded_batch_dimensions.req_count self.padding_slice = slice(self.active_token_count, self.padded_active_token_count) + self.build_active_slices(self.padded_active_request_count) + self.pad_active_slices() + # Update token position indexes. self.token_to_block_idx[self.active_token_count : self.padded_active_token_count] = ( self.kv_block_allocator.dummy_block_idx @@ -1912,6 +1937,20 @@ def speculative_required_logit_indices(self, device: torch.device) -> Tensor: return torch.cat([decode_indices, prefill_last_indices]) + @property + def num_last_token_logits(self) -> int: + """Number of rows produced by `last_token_logits` for the current step. + + Single source of truth for the bound: one row per request, with + `(num_speculative_tokens + 1)` rows per decode request when MTP is active. + """ + if self.num_speculative_tokens > 0: + return ( + self.num_decode_requests * (self.num_speculative_tokens + 1) + + self.num_prefill_requests + ) + return self.total_request_count - self.paused_request_count + def last_token_logits(self, logits: Tensor) -> Tensor: """Select the logit positions needed for token generation. @@ -1925,7 +1964,7 @@ def last_token_logits(self, logits: Tensor) -> Tensor: logits (Tensor): Output logits of forward pass, shape [1, S, H]. Return: - (Tensor) Selected logits, shape [N, H]. + (Tensor) Selected logits, shape [N, H], where N == num_last_token_logits. """ # todo: @lmcafee, remove these asserts? assert logits.size(0) == 1, f"logits.size(0) ({tuple(logits.shape)}) != 1" @@ -1937,12 +1976,14 @@ def last_token_logits(self, logits: Tensor) -> Tensor: if self.num_speculative_tokens > 0: selected = self.speculative_required_logit_indices(logits.device) + assert selected.numel() == self.num_last_token_logits return logits_2d[selected, :] paused = self.paused_request_count total = self.total_request_count query_lengths = self.request_query_lengths[paused:total] last_token_idxs = torch.cumsum(query_lengths, dim=0) - 1 + assert last_token_idxs.numel() == self.num_last_token_logits return logits_2d[last_token_idxs, :] def _compute_prefix_match( @@ -2183,7 +2224,7 @@ def add_request( metadata = req.tracked_metadata metadata_types = req.get_metadata_types() for m, m_type in zip(metadata, metadata_types): - label, _, _ = m_type + label, _ = m_type if not isinstance(m, torch.Tensor): m = torch.as_tensor( m, diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 2a40937ae6b..1c5a4c4cfb1 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -400,39 +400,39 @@ def create_cuda_graphs(self, reset_context: bool = True): # Enable routing recording during warmup if routing replay is enabled. # This ensures the record_indices copy operation is captured in the CUDA graph. - model_config = controller.inference_wrapped_model.model.config if model_config.moe_enable_routing_replay: RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD) # Forward pass -> logits. - controller._dynamic_step_forward_logits(input_ids, position_ids) - - # MTP CUDA graph warmup for this batch dimension. - if mtp_warmup_enabled: - n = cuda_graph_batch_dimension.req_count - # pylint: disable-next=possibly-used-before-assignment - if sp_enabled: - n = round_up_to_nearest_multiple(n, tp_size) - # pylint: disable-next=possibly-used-before-assignment - if n > 0 and n not in mtp_seen_batch_sizes: - mtp_seen_batch_sizes.add(n) - device = torch.cuda.current_device() - batch_dim = n // tp_size if sp_enabled else n - # Use zeros (not empty) — garbage token IDs cause OOB embedding lookups during graph capture/replay. - for depth in mtp_warmup_depths: - unwrapped.compute_mtp_single_step( - hidden_states=torch.zeros( - (batch_dim, 1, model_config.hidden_size), - device=device, - dtype=model_config.params_dtype, - ), - next_token_ids=torch.zeros((1, n), device=device, dtype=torch.long), - position_ids=torch.zeros((1, n), device=device, dtype=torch.int64), - depth=depth, - cache_key=("mtp", n, depth), - ) + with torch.inference_mode(): + controller._dynamic_step_forward_logits(input_ids, position_ids) + + # MTP CUDA graph warmup for this batch dimension. + if mtp_warmup_enabled: + n = cuda_graph_batch_dimension.req_count + # pylint: disable-next=possibly-used-before-assignment + if sp_enabled: + n = round_up_to_nearest_multiple(n, tp_size) + # pylint: disable-next=possibly-used-before-assignment + if n > 0 and n not in mtp_seen_batch_sizes: + mtp_seen_batch_sizes.add(n) + device = torch.cuda.current_device() + batch_dim = n // tp_size if sp_enabled else n + # Use zeros (not empty) — garbage token IDs cause OOB embedding lookups during graph capture/replay. + for depth in mtp_warmup_depths: + unwrapped.compute_mtp_single_step( + hidden_states=torch.zeros( + (batch_dim, 1, model_config.hidden_size), + device=device, + dtype=model_config.params_dtype, + ), + next_token_ids=torch.zeros((1, n), device=device, dtype=torch.long), + position_ids=torch.zeros((1, n), device=device, dtype=torch.int64), + depth=depth, + cache_key=("mtp", n, depth), + ) - context.reset() + context.reset() # Disable inference dispatcher after graph capture if is_inference_optimized_ep: diff --git a/megatron/core/inference/inference_request.py b/megatron/core/inference/inference_request.py index 3917bbba720..33fbcdf6518 100644 --- a/megatron/core/inference/inference_request.py +++ b/megatron/core/inference/inference_request.py @@ -476,26 +476,25 @@ def tracked_metadata(self) -> List[Any]: "in its sampling_params. Defaulting to -1." ) sp.termination_id = -1 - return [getattr(sp, field) for field, _, _ in self.get_metadata_types()] + return [getattr(sp, field) for field, _ in self.get_metadata_types()] @staticmethod - def get_metadata_types() -> List[Tuple[str, torch.dtype, bool]]: - """Keeps track of all request metadata names, dtypes, and target device. + def get_metadata_types() -> List[Tuple[str, torch.dtype]]: + """Keeps track of all request metadata names and dtypes. Returns: - List[Tuple[str, torch.dtype, bool]]: Mapping from metadata name to: + List[Tuple[str, torch.dtype]]: Mapping from metadata name to: name (str) - The name of the metadata field. dtype (torch.dtype) - The datatype of the metadata. - on_device (bool) - Whether the metadata lives on GPU (True) or CPU (False). """ return [ - ("temperature", torch.float32, False), # CPU for torch sampling - ("top_k", torch.int32, False), # CPU for torch sampling - ("top_p", torch.float32, False), # CPU for torch sampling - ("termination_id", torch.int64, True), - ("return_log_probs", torch.bool, False), # CPU for non-selective logprobs - ("skip_prompt_log_probs", torch.bool, False), # CPU for non-selective logprobs - ("top_n_logprobs", torch.int32, False), # CPU for torch sampling + ("temperature", torch.float32), + ("top_k", torch.int32), + ("top_p", torch.float32), + ("termination_id", torch.int64), + ("return_log_probs", torch.bool), + ("skip_prompt_log_probs", torch.bool), + ("top_n_logprobs", torch.int32), ] def add_event( 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 e66591edad0..061d0083ea2 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -137,6 +137,11 @@ def _init_dynamic_sampling_tensors(self): """Initialize tensors needed for dynamic sampling.""" context = self.inference_wrapped_model.inference_context max_requests = context.max_requests + if context.config.materialize_only_last_token_logits: + # Under MTP, each decode request emits (num_speculative_tokens + 1) logit rows + max_logits = max_requests * (self.num_speculative_tokens + 1) + else: + max_logits = context.max_tokens # Callback to get request IDs that should be marked as finished due to stop words self._get_stop_word_finished_ids_callback = None @@ -145,17 +150,16 @@ def _init_dynamic_sampling_tensors(self): logits_dtype = self.inference_wrapped_model.config.params_dtype self._sampling_backend = "torch" - self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device) + self._enable_cuda_graph = self.model_config.cuda_graph_impl == "local" - # Keep track of request metadata. - self._request_metadata: Dict[str, Tensor] = {} - for label, dtype, on_gpu in context.request_metadata_types: - tensor = context.request_metadata[label] - if not on_gpu: - # Create pinned tensors for request metadata that lives on CPU. - # This is metadata which requires D2H copies, such as top_k for torch sampling. - tensor = torch.empty_like(tensor, device="cpu", pin_memory=True) - self._request_metadata[label] = tensor + # Initialize bookkeeping tensors. + if self._enable_cuda_graph: + self._all_logits_cuda = torch.zeros( + (1, max_logits, self.vocab_size), dtype=logits_dtype, device=device + ) + else: + self._all_logits_cuda = None + self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device) # Used for inefficient torch sampling. if self._sampling_backend == "torch": @@ -599,7 +603,6 @@ def _dynamic_step_context_init( position_ids (Tensor): The active position IDs. """ context = self.inference_wrapped_model.inference_context - active_request_slice = slice(context.paused_request_count, context.total_request_count) # Remove Float16Module wrapper if it exists unwrapped_model = unwrap_model(self.inference_wrapped_model.model) @@ -652,14 +655,6 @@ def _dynamic_step_context_init( # Turn off symmetric all reduces for prefill unwrapped_model.set_symmetric_ar(None) - # Get request metadata for this step. - for label, dtype, on_gpu in context.request_metadata_types: - if not on_gpu: - # We need a D2H copy from the context to the pinned memory buffer. - self._request_metadata[label].copy_( - context.request_metadata[label], non_blocking=True - ) - # Get flat tokens, position ids. # If we are running a dummy forward step we want to use the token count agreed upon # by all EP ranks rather than the minimum number of tokens. @@ -670,7 +665,7 @@ def _dynamic_step_context_init( else: return context.current_input_and_position_ids() - def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) -> Tensor: + def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor): """Forward step the model to get logits for dynamic batching. This also handles logits-broadcasting for pipeline parallelism. @@ -680,7 +675,10 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) position_ids (Tensor): The position IDs. """ context = self.inference_wrapped_model.inference_context - active_request_count = context.total_request_count - context.paused_request_count + if context.config.materialize_only_last_token_logits: + logits_seq_len = context.num_last_token_logits + else: + logits_seq_len = context.padded_active_token_count with torch.inference_mode(): logits = self.inference_wrapped_model.run_one_forward_step( @@ -688,6 +686,9 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) ) # logits shape: [1, seq_len, vocab_size] + if not context.config.materialize_only_last_token_logits: + assert logits_seq_len == input_ids.shape[1] + # Note: When speculative decoding is active (num_speculative_tokens > 0), # the model skips MTP computation during the forward pass. MTP logits # will be computed serially after verification to ensure they are @@ -695,13 +696,7 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) if self.model_is_pipeline_parallel: if context.config.materialize_only_last_token_logits: - if self.num_speculative_tokens > 0: - logits_seq_len = ( - context.num_decode_requests * (self.num_speculative_tokens + 1) - + context.num_prefill_requests - ) - else: - logits_seq_len = active_request_count + logits_seq_len = context.num_last_token_logits else: logits_seq_len = input_ids.shape[1] logits_shape = [1, logits_seq_len, self.vocab_size] @@ -716,12 +711,16 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor) pp_group=self.pp_group, ) - return logits + # Copy logits to contiguous buffer. + if self._enable_cuda_graph: + self._all_logits_cuda[:, :logits_seq_len, :].copy_(logits[:, :logits_seq_len, :]) + else: + self._all_logits_cuda = logits def _dynamic_step_sample_bookkeeping(self): """Perform bookkeeping necessary to sample logits for dynamic batching.""" context = self.inference_wrapped_model.inference_context - active_request_slice = slice(context.paused_request_count, context.total_request_count) + active_request_count = context.total_request_count - context.paused_request_count if self._sampling_backend == "torch": # Bucketize the core sampling parameters. @@ -729,9 +728,9 @@ def _dynamic_step_sample_bookkeeping(self): bucket_map = defaultdict(list) # Shorthands for the dictionary comprehension. - temp = self._request_metadata["temperature"][active_request_slice].tolist() - top_k = self._request_metadata["top_k"][active_request_slice].tolist() - top_p = self._request_metadata["top_p"][active_request_slice].tolist() + temp = context.active_request_metadata["temperature"][:active_request_count].tolist() + top_k = context.active_request_metadata["top_k"][:active_request_count].tolist() + top_p = context.active_request_metadata["top_p"][:active_request_count].tolist() for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)): sampling_params = (t, k, p) @@ -1042,7 +1041,7 @@ def _verify_speculative_tokens( num_speculative_tokens=self.num_speculative_tokens, ) - def _dynamic_step_sample_logits_and_verify_tokens(self, logits: Tensor, input_ids: Tensor): + def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor): """ Sample tokens from logits for dynamic batching with speculative tokens and verify the tokens. """ @@ -1057,6 +1056,8 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, logits: Tensor, input_id # These indices are always needed for input_ids slicing and tracking # accepted sequence positions, even when logits are pre-sliced. nvtx_range_push("mtp-spec-decoding/verify/logit-indices") + # Use pre-allocated buffer for CUDA graph compatibility. + logits = self._all_logits_cuda required_logit_indices = context.speculative_required_logit_indices(logits.device) if context.config.materialize_only_last_token_logits: @@ -1137,24 +1138,23 @@ def _prepare_speculative_tokens_for_next_forward_pass( # Expose the active slice so downstream code sees the right length. self._last_accepted_seq_indices = self._last_accepted_seq_indices_buf[:active_request_count] - def _dynamic_step_sample_logits(self, logits: Tensor): - """Sample tokens from logits for dynamic batching. - - Args: - logits (Tensor): The logits from the forward pass. - """ + def _dynamic_step_sample_logits(self): + """Sample tokens from logits for dynamic batching.""" # TODO(ksanthanam): Evaluate whether it makes more sense to sample on 1 rank # and then broadcast the sampled tokens rather than broadcasting the raw logits. # Last token logits. context = self.inference_wrapped_model.inference_context + active_request_count = context.total_request_count - context.paused_request_count + if context.config.materialize_only_last_token_logits: # When materialize_only_last_token_logits is true, last_token_logits is # already called in the forward pass of GPT. - required_token_logits = logits.squeeze(0) + required_token_logits = self._all_logits_cuda.squeeze(0)[:active_request_count, :] else: - # todo : Should do verification here and get approrpiate las token logits - required_token_logits = context.last_token_logits(logits) + required_token_logits = context.last_token_logits( + self._all_logits_cuda[:, : context.padded_active_token_count, :] + ) if self._sampling_backend == "torch": # Concatenate the outputs once to prevent repeated small writes. @@ -1185,12 +1185,12 @@ def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]: return_log_probs (bool): Whether to return the sampled log_probs. """ context = self.inference_wrapped_model.inference_context - active_request_slice = slice(context.paused_request_count, context.total_request_count) - - return_log_probs = self._request_metadata["return_log_probs"][active_request_slice] - top_n_log_probs = self._request_metadata["top_n_logprobs"][active_request_slice] > 0 + active_request_count = context.total_request_count - context.paused_request_count - return return_log_probs.any(), top_n_log_probs.any() + return ( + (context.active_request_metadata["return_log_probs"][:active_request_count]).any(), + (context.active_request_metadata["top_n_logprobs"][:active_request_count] > 0).any(), + ) def _router_record_bookkeeping(self) -> Optional[np.ndarray]: """Collect flat routing indices for MoE router recording. @@ -1251,20 +1251,25 @@ def _router_record_bookkeeping(self) -> Optional[np.ndarray]: _ri_dtype = np.int16 if (config.num_moe_experts or 0) <= 32768 else np.int32 return stacked_routing[:active_token_count].cpu().numpy().astype(_ri_dtype) - def _dynamic_step_calculate_log_probs(self, logits: Tensor) -> Optional[Tensor]: + def _dynamic_step_calculate_log_probs(self) -> Optional[Tensor]: """Calculate log probs from logits.""" context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count + # This code cannot be reached when we are using speculative decode. + assert self.num_speculative_tokens == 0 + logits_seq_len = ( + active_request_count + if context.config.materialize_only_last_token_logits + else context.padded_active_token_count + ) return context.calculate_log_probs( - logits, + 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, ) - def _dynamic_step_calculate_log_probs_speculative( - self, logits: Tensor - ) -> Tuple[List[List[float]], Tensor]: + def _dynamic_step_calculate_log_probs_speculative(self) -> Tuple[List[List[float]], Tensor]: """Calculate log probs from logits for speculative decoding. For decode requests, computes log probs for each accepted speculative token @@ -1275,9 +1280,6 @@ def _dynamic_step_calculate_log_probs_speculative( - log_prob(accepted_token[j]) comes from logits at position j - log_prob(newly_sampled_token) comes from logits at position accepted_count - Args: - logits (Tensor): The main model logits [1, seq_len, vocab_size]. - Returns: Tuple of (log_probs_list, log_probs_tensor): log_probs_list: List of lists, one per active request, containing @@ -1298,6 +1300,8 @@ def _dynamic_step_calculate_log_probs_speculative( num_decode_requests = active_request_count - num_prefill_requests only_last = context.config.materialize_only_last_token_logits + # Use pre-allocated buffer for CUDA graph compatibility. + logits = self._all_logits_cuda logits_squeezed = logits.squeeze(0).float() if only_last: log_probs_tensor = F.log_softmax(logits_squeezed, dim=-1) @@ -1395,7 +1399,6 @@ def _dynamic_step_calculate_top_n_logprobs_speculative( """ context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - active_request_slice = slice(context.paused_request_count, context.total_request_count) request_in_prefill_status_tensor = context.request_in_prefill_status_tensor[ context.paused_request_count : context.total_request_count @@ -1415,7 +1418,7 @@ def _dynamic_step_calculate_top_n_logprobs_speculative( num_decode_requests, self.num_speculative_tokens + 1, -1 ) accepted_counts = self._accepted_token_counts_per_request[:num_decode_requests] - top_n_per_request = self._request_metadata["top_n_logprobs"][active_request_slice][ + top_n_per_request = context.active_request_metadata["top_n_logprobs"][ :num_decode_requests ] max_top_n = int(top_n_per_request.max().item()) @@ -1444,8 +1447,8 @@ def _dynamic_step_calculate_top_n_logprobs_speculative( prefill_log_probs = log_probs_tensor[decode_len:] # Batch metadata reads: single CPU transfer for all prefill requests. - prefill_top_n = self._request_metadata["top_n_logprobs"][active_request_slice][ - num_decode_requests: + prefill_top_n = context.active_request_metadata["top_n_logprobs"][ + num_decode_requests:active_request_count ].tolist() max_top_n_prefill = int(max(prefill_top_n)) if prefill_top_n else 0 @@ -1472,7 +1475,7 @@ def _dynamic_step_calculate_top_n_logprobs_speculative( prefill_log_probs_per_request = prefill_log_probs.split( prefill_query_lengths.tolist(), dim=0 ) - prefill_skip_prompt = self._request_metadata["skip_prompt_log_probs"][ + prefill_skip_prompt = context.active_request_metadata["skip_prompt_log_probs"][ num_decode_requests:active_request_count ].tolist() @@ -1500,12 +1503,11 @@ def _dynamic_step_calculate_top_n_logprobs_speculative( return top_n_results if top_n_results else None def _dynamic_step_calculate_top_n_logprobs( - self, logits: Tensor, log_probs_tensor: Optional[Tensor] = None + self, log_probs_tensor: Optional[Tensor] = None ) -> Optional[Dict[int, List[Tuple[Tensor, Tensor]]]]: """Calculate top-n log probs from logits for dynamic batching. Args: - logits (Tensor): The logits to compute top-n log probs from. log_probs_tensor (Optional[Tensor]): Pre-computed log probabilities tensor. If provided, avoids recomputing log_softmax. Should be the tensor returned by calculate_log_probs. @@ -1531,9 +1533,7 @@ def _dynamic_step_calculate_top_n_logprobs( top_n_results = {} for req_idx in range(active_request_count): - top_n = int( - self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item() - ) + top_n = int(context.active_request_metadata["top_n_logprobs"][req_idx].item()) if top_n > 0: # Get top-n logprobs and indices for this request (single token) top_n_logits = torch.topk(log_probs[req_idx], k=top_n) @@ -1555,14 +1555,14 @@ def _dynamic_step_calculate_top_n_logprobs( top_n_results = {} for req_idx in range(active_request_count): - top_n = int( - self._request_metadata["top_n_logprobs"][active_request_slice][req_idx].item() - ) + top_n = int(context.active_request_metadata["top_n_logprobs"][req_idx].item()) if top_n > 0: request_log_probs = log_probs_per_request[ req_idx ] # [num_tokens_for_request, vocab_size] - skip_prompt = bool(self._request_metadata["skip_prompt_log_probs"][req_idx].item()) + skip_prompt = bool( + context.active_request_metadata["skip_prompt_log_probs"][req_idx].item() + ) # If skip_prompt_log_probs is True, only compute for last token if skip_prompt and request_log_probs.size(0) > 1: @@ -1752,10 +1752,11 @@ def _dynamic_step_context_bookkeeping(self) -> Dict[str, Tensor]: # Note: termination_id tensor has per-request termination IDs from mixed sampling active_request_mask = ( self._sampled_tokens_cuda[:active_request_count] - != self._request_metadata["termination_id"][active_request_slice] + != context.active_request_metadata["termination_id"][:active_request_count] ).byte() & torch.less(active_sequence_lengths, max_sequence_lengths).byte() - # Mark requests as finished if they hit stop words (detected in previous step's post_process_requests) + # Mark requests as finished if they hit stop words + # (detected in previous step's post_process_requests) if self._get_stop_word_finished_ids_callback is not None: request_ids_list = active_request_ids.tolist() stop_word_finished_ids = self._get_stop_word_finished_ids_callback(request_ids_list) @@ -1841,7 +1842,7 @@ async def async_generate_output_tokens_dynamic_batch( # Forward pass produces only base logits. When speculative decoding is # active, MTP logits are computed serially after verification. - logits = self._dynamic_step_forward_logits(input_ids, position_ids) + self._dynamic_step_forward_logits(input_ids, position_ids) # Commit Mamba intermediate states before update_requests, which # may swap request indices. The Python lists tracking EOS block IDs @@ -1872,7 +1873,7 @@ async def async_generate_output_tokens_dynamic_batch( if self.num_speculative_tokens > 0: # Phase 1: Verify speculative tokens using base logits only. nvtx_range_push("mtp-spec-decoding/verify") - self._dynamic_step_sample_logits_and_verify_tokens(logits, input_ids) + self._dynamic_step_sample_logits_and_verify_tokens(input_ids) nvtx_range_pop("mtp-spec-decoding/verify") # Phase 2: Rewind KV cache for rejected tokens. nvtx_range_push("mtp-spec-decoding/rewind-kv-cache") @@ -1894,24 +1895,24 @@ async def async_generate_output_tokens_dynamic_batch( # data-dependent boolean-mask sync overlaps with MTP GPU work. context.kv_block_allocator.release_memory_blocks(blocks_to_release[remove_mask]) else: - self._dynamic_step_sample_logits(logits) + self._dynamic_step_sample_logits() log_probs = None top_n_logprobs = None if return_log_probs or return_top_n_logprobs: if self.num_speculative_tokens > 0: log_probs, log_probs_tensor = ( - self._dynamic_step_calculate_log_probs_speculative(logits) + self._dynamic_step_calculate_log_probs_speculative() ) if return_top_n_logprobs: top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs_speculative( log_probs_tensor ) else: - log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs(logits) + log_probs, log_probs_tensor = self._dynamic_step_calculate_log_probs() if return_top_n_logprobs: top_n_logprobs = self._dynamic_step_calculate_top_n_logprobs( - logits, log_probs_tensor + log_probs_tensor ) if skip_bookkeeping: diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index f0c41ca7d83..8ce21761f66 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py @@ -67,6 +67,7 @@ def setup_model( expert_model_parallel_size: int = 1, num_moe_experts: int = None, hybrid_layer_pattern: str = None, + cuda_graph_impl: str = 'none', ): if use_training_random_init: # This is necessary to induce the training behavior which permutes the random seed @@ -102,6 +103,7 @@ def setup_model( if hybrid_layer_pattern else {} ), + cuda_graph_impl=cuda_graph_impl, ) if dtype == torch.bfloat16: transformer_config.bf16 = True @@ -315,12 +317,9 @@ def test_sample_from_dynamic_logits( temp_values = torch.Tensor([s.temperature for s in rev_sampling_dict]) top_k_values = torch.Tensor([s.top_k for s in rev_sampling_dict]).to(torch.int32) top_p_values = torch.Tensor([s.top_p for s in rev_sampling_dict]) - request_metadata = { - "temperature": temp_values, - "top_k": top_k_values, - "top_p": top_p_values, - } - self.text_generation_controller._request_metadata = request_metadata + context.active_request_metadata["temperature"][:batch_size].copy_(temp_values) + context.active_request_metadata["top_k"][:batch_size].copy_(top_k_values) + context.active_request_metadata["top_p"][:batch_size].copy_(top_p_values) self.text_generation_controller._sampling_backend = backend context.padded_active_token_count = batch_size @@ -333,7 +332,8 @@ def test_sample_from_dynamic_logits( # Sampling. logits = torch.arange(0, self.vocab_size).repeat(batch_size, 1).unsqueeze(0).float().cuda() - self.text_generation_controller._dynamic_step_sample_logits(logits) + self.text_generation_controller._all_logits_cuda = logits + self.text_generation_controller._dynamic_step_sample_logits() sampled_logits = self.text_generation_controller._sampled_tokens_cuda[:batch_size] vocab_indices = torch.arange(self.vocab_size).cuda() @@ -857,15 +857,10 @@ def test_dynamic_top_n_logprobs_calculation( # Prepare sampling params top_n = 5 - request_metadata = { - "top_n_logprobs": torch.full((batch_size,), top_n, dtype=torch.int32).cuda(), - "skip_prompt_log_probs": torch.full( - (batch_size,), float(skip_prompt_log_probs), dtype=torch.float32 - ).cuda(), - } - self.text_generation_controller._request_metadata = request_metadata - self.text_generation_controller._active_request_count = batch_size - self.text_generation_controller._active_request_slice = slice(0, batch_size) + context.active_request_metadata["top_n_logprobs"][:batch_size].fill_(top_n) + context.active_request_metadata["skip_prompt_log_probs"][:batch_size].fill_( + skip_prompt_log_probs + ) if materialize_only_last_token_logits: # Decode mode: logits for last tokens only @@ -881,7 +876,7 @@ def test_dynamic_top_n_logprobs_calculation( # Calculate top-n logprobs top_n_results = self.text_generation_controller._dynamic_step_calculate_top_n_logprobs( - logits, log_probs_tensor + log_probs_tensor ) # Validate results @@ -926,7 +921,7 @@ def test_dynamic_top_n_logprobs_calculation( # Calculate top-n logprobs top_n_results = self.text_generation_controller._dynamic_step_calculate_top_n_logprobs( - logits, log_probs_tensor + log_probs_tensor ) # Validate results @@ -1010,10 +1005,9 @@ def mock_sampling_func(logits, *args, **kwargs): # Mock logits matching input shape logits = torch.randn(1, 6, self.vocab_size, device='cuda') + self.text_generation_controller._all_logits_cuda = logits - self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens( - logits, input_ids - ) + self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens(input_ids) # Verify acceptance counts accepted_counts = self.text_generation_controller._accepted_token_counts_per_request[:2] @@ -1242,10 +1236,9 @@ def test_speculative_multinomial_sampling(self): # Since we are actually testing the internal math of `_torch_sampling_func` handling the shapes, # we DO NOT mock `_torch_sampling_func` here. We want it to run natively to prove it doesn't crash. + self.text_generation_controller._all_logits_cuda = logits try: - self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens( - logits, input_ids - ) + self.text_generation_controller._dynamic_step_sample_logits_and_verify_tokens(input_ids) except RuntimeError as e: if "prob_dist must be 1 or 2 dim" in str(e): pytest.fail("MTP logits were not flattened before calling multinomial sampling.")