diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index b3c0cd8a3cf..c3e8b4f5e1c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2438,10 +2438,6 @@ def reset_metadata(self) -> None: self.kv_block_allocator.reset() self.request_to_kv_block_ids.fill_(-1) - # Reset step counter and LRU clock - self.step_count = 0 - self.prefix_cache_lru_clock = 0 - # Reset chunked prefill state self.chunked_prefill_request_id = -1 self.num_prefill_requests = 0 @@ -2466,6 +2462,11 @@ def reset(self) -> None: self.reset_tensors() self.reset_metadata() + # Reset lifetime counters (not reset in reset_metadata, which is also + # called during suspend/resume where these must persist). + self.step_count = 0 + self.prefix_cache_lru_clock = 0 + # Reset Mamba cache state if self.mamba_slot_allocator is not None: self.mamba_slot_allocator.reset() diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index 8a43eb0f7ae..e75000bc2b9 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -241,8 +241,8 @@ def __init__(self, controller: TextGenerationController, context: DynamicInferen if self.num_speculative_tokens > 0: assert ( model_config.mtp_use_repeated_layer - or self.num_speculative_tokens <= self.controller.num_mtp_heads - ), f"Number of speculative tokens {self.num_speculative_tokens} must be less than or equal to number of MTP heads {self.controller.num_mtp_heads}" + or self.num_speculative_tokens <= model_config.mtp_num_layers + ), f"Number of speculative tokens {self.num_speculative_tokens} must be less than or equal to number of MTP layers {model_config.mtp_num_layers}" self.track_paused_request_events = inference_config.track_paused_request_events self.track_generated_token_events = inference_config.track_generated_token_events self.enable_chunked_prefill = inference_config.enable_chunked_prefill @@ -329,9 +329,15 @@ def reset(self) -> None: self.resume_request_ids = None - # Speculative decoding acceptance tracking. - self._spec_tokens_proposed = 0 - self._spec_tokens_accepted = 0 + # Speculative decoding acceptance tracking (per-position). + # Each tensor has length num_speculative_tokens; index i tracks position i+1 + # (i.e. the i-th draft token proposed by the MTP head). + self._spec_tokens_proposed_per_pos = torch.zeros( + self.num_speculative_tokens, dtype=torch.int64 + ) + self._spec_tokens_accepted_per_pos = torch.zeros( + self.num_speculative_tokens, dtype=torch.int64 + ) self._spec_steps = 0 # Prefix caching tracking. @@ -394,7 +400,7 @@ def create_cuda_graphs(self, reset_context: bool = True): # decoder graphs within the same loop rather than in a separate pass. unwrapped = unwrap_model(controller.inference_wrapped_model.model) mtp_warmup_enabled = ( - controller.num_mtp_heads > 0 + controller.num_mtp_depths > 0 and (controller.num_speculative_tokens or 0) > 0 and hasattr(unwrapped, 'mtp') ) @@ -402,7 +408,7 @@ def create_cuda_graphs(self, reset_context: bool = True): tp_size = get_pg_size(controller.inference_wrapped_model.tp_group) sp_enabled = model_config.sequence_parallel and tp_size > 1 mtp_pass_depth = not unwrapped.mtp.mtp_use_repeated_layer - mtp_warmup_depths = range(controller._num_mtp_depths) if mtp_pass_depth else [None] + mtp_warmup_depths = range(controller.num_mtp_depths) if mtp_pass_depth else [None] mtp_seen_batch_sizes = set() tbar = enumerate(context.cuda_graph_batch_dimensions_list) @@ -1200,6 +1206,7 @@ def post_process_requests( tokens = accepted_tokens + tokens num_stop_word_trim = 0 + is_prefill = len(request.generated_tokens) == 0 if request_id != self.context.chunked_prefill_request_id: # Skip appending token for requests being finished due to stop words # (they already have their final token from the previous step) @@ -1270,13 +1277,21 @@ def post_process_requests( request ) - # Track acceptance statistics for logging. - if len(request.generated_tokens) > 0 and self.num_speculative_tokens > 0: + # Track per-position acceptance statistics for logging. + # Skip prefill requests: MTP heads only propose speculative tokens + # for decode requests, so counting prefill requests would inflate + # the denominator and artificially deflate the acceptance rate. + if ( + not is_prefill + and len(request.generated_tokens) > 0 + and self.num_speculative_tokens > 0 + ): actual_proposed = max(0, self.num_speculative_tokens - num_stop_word_trim) - actual_accepted = max(0, len(accepted_tokens) - num_stop_word_trim) - - self._spec_tokens_proposed += actual_proposed - self._spec_tokens_accepted += actual_accepted + self._spec_tokens_proposed_per_pos[:actual_proposed] += 1 + accepted_t = torch.tensor(accepted_tokens_list[:actual_proposed]) + self._spec_tokens_accepted_per_pos[:actual_proposed] += ( + accepted_t != -1 + ).long() if request_id in finished_request_ids: # Reconstruct routing from per-block storage before popping. @@ -1943,13 +1958,24 @@ async def async_bookkeep( else: metrics[f'inference/{key}'] = value - # Add speculative decoding acceptance metrics. - if self.num_speculative_tokens > 0 and self._spec_tokens_proposed > 0: - acceptance_rate = self._spec_tokens_accepted / self._spec_tokens_proposed + # Add speculative decoding acceptance metrics (aggregate + per-position). + total_proposed = sum(self._spec_tokens_proposed_per_pos) + total_accepted = sum(self._spec_tokens_accepted_per_pos) + if self.num_speculative_tokens > 0 and total_proposed > 0: + acceptance_rate = total_accepted / total_proposed metrics['inference/spec_decode_acceptance_rate'] = float(acceptance_rate * 100.0) - metrics['inference/spec_decode_tokens_proposed'] = int(self._spec_tokens_proposed) - metrics['inference/spec_decode_tokens_accepted'] = int(self._spec_tokens_accepted) + metrics['inference/spec_decode_tokens_proposed'] = int(total_proposed) + metrics['inference/spec_decode_tokens_accepted'] = int(total_accepted) metrics['inference/spec_decode_num_steps'] = int(self._spec_steps) + for pos in range(self.num_speculative_tokens): + if self._spec_tokens_proposed_per_pos[pos] > 0: + pos_rate = ( + self._spec_tokens_accepted_per_pos[pos] + / self._spec_tokens_proposed_per_pos[pos] + ) + metrics[f'inference/spec_decode_acceptance_rate_pos{pos + 1}'] = float( + pos_rate * 100.0 + ) # Add prefix caching metrics. if self.context.enable_prefix_caching and self._prefix_cache_hits > 0: @@ -2011,16 +2037,28 @@ async def async_bookkeep( mem["reserved_bytes.all.current"] / (1024**3), ) ) - if self.num_speculative_tokens > 0 and self._spec_tokens_proposed > 0: - spec_rate = self._spec_tokens_accepted / self._spec_tokens_proposed * 100.0 - output_str += " ... spec: accept %.1f%% (%d/%d in %d steps)" % ( + total_proposed = sum(self._spec_tokens_proposed_per_pos) + total_accepted = sum(self._spec_tokens_accepted_per_pos) + if self.num_speculative_tokens > 0 and total_proposed > 0: + spec_rate = total_accepted / total_proposed * 100.0 + per_pos_rates = [] + for pos in range(self.num_speculative_tokens): + if self._spec_tokens_proposed_per_pos[pos] > 0: + pos_rate = ( + self._spec_tokens_accepted_per_pos[pos] + / self._spec_tokens_proposed_per_pos[pos] + * 100.0 + ) + per_pos_rates.append("t%d=%.1f%%" % (pos + 1, pos_rate)) + output_str += " ... spec (cumul): accept %.1f%% (%d/%d in %d steps) [%s]" % ( spec_rate, - self._spec_tokens_accepted, - self._spec_tokens_proposed, + total_accepted, + total_proposed, self._spec_steps, + ", ".join(per_pos_rates), ) if self.context.enable_prefix_caching and self._prefix_cache_hits > 0: - output_str += " ... prefix cache: %d hits, %d blocks matched" % ( + output_str += " ... prefix cache (cumul): %d hits, %d blocks matched" % ( self._prefix_cache_hits, self._prefix_cache_blocks_matched, ) @@ -2028,17 +2066,6 @@ async def async_bookkeep( output_str = f"\033[94m{output_str}\033[0m" logging.info(output_str) - # Reset speculative decoding accumulators after both wandb and console logging. - if self.num_speculative_tokens > 0: - self._spec_tokens_proposed = 0 - self._spec_tokens_accepted = 0 - self._spec_steps = 0 - - # Reset prefix caching accumulators after both wandb and console logging. - if self.context.enable_prefix_caching: - self._prefix_cache_hits = 0 - self._prefix_cache_blocks_matched = 0 - nvtx_range_pop("console_logging") return { 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 3e788fec0b1..399da90202d 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -104,9 +104,21 @@ def __init__(self, inference_wrapped_model: AbstractModelInferenceWrapper, token self.vocab_size = unwrapped_model.vocab_size self.sampling_rng = torch.Generator(device=torch.cuda.current_device()) - self.num_mtp_heads = self._get_mtp_num_heads() self.sampling_rng.manual_seed(self.model_config.inference_sampling_seed) + if not self.num_speculative_tokens: + self.num_mtp_depths = 0 + else: + assert ( + self.model_config.mtp_num_layers and self.model_config.mtp_num_layers >= 1 + ), "mtp_num_layers must be >= 1 when num_speculative_tokens > 0" + if self.model_config.mtp_use_repeated_layer: + self.num_mtp_depths = self.num_speculative_tokens + else: + self.num_mtp_depths = min( + self.num_speculative_tokens, self.model_config.mtp_num_layers + ) + if ( self.model_config.cuda_graph_impl == "local" and self.model_config.expert_model_parallel_size > 1 @@ -120,13 +132,6 @@ def __init__(self, inference_wrapped_model: AbstractModelInferenceWrapper, token if self.inference_wrapped_model.inference_context.is_dynamic_batching(): self._init_dynamic_sampling_tensors() - def _get_mtp_num_heads(self) -> int: - """Get the number of MTP layers from the model config.""" - model = self.inference_wrapped_model.model - if hasattr(model, 'config') and hasattr(model.config, 'mtp_num_layers'): - return model.config.mtp_num_layers or 0 - return 0 - def set_stop_word_finished_ids_callback(self, callback): """Set a callback to get request IDs that should be marked as finished due to stop words. @@ -222,7 +227,6 @@ def _init_mtp_sampling_tensors(self): max_requests, dtype=torch.int64, device=device ) self._last_accepted_seq_indices = None - self._num_mtp_depths = min(self.num_speculative_tokens, self.num_mtp_heads) self._mtp_token_ids_buf = torch.empty([1, max_requests], dtype=torch.int64, device=device) self._mtp_position_ids_buf = torch.empty( [1, max_requests], dtype=torch.int64, device=device @@ -833,7 +837,7 @@ def _compute_serial_mtp_and_sample(self): position_ids_buf[0, active_request_count:] = 0 nvtx_range_pop("mtp-spec-decoding/serial-mtp-init") - for depth in range(self._num_mtp_depths): + for depth in range(self.num_mtp_depths): nvtx_range_push(f"mtp-spec-decoding/depth-{depth}") token_ids_buf[0, :active_request_count] = next_token_ids @@ -1508,7 +1512,7 @@ def _dummy_serial_mtp_forward(self): - When PP > 1: participate in the ``broadcast_from_last_pipeline_stage`` that the real ranks also perform. """ - if self.num_speculative_tokens == 0 or self.num_mtp_heads == 0: + if self.num_speculative_tokens == 0 or self.num_mtp_depths == 0: return if self.model_config.expert_model_parallel_size <= 1: return @@ -1549,7 +1553,7 @@ def _dummy_serial_mtp_forward(self): context = self.inference_wrapped_model.inference_context - for depth in range(self._num_mtp_depths): + for depth in range(self.num_mtp_depths): nvtx_range_push(f"mtp-spec-decoding/dummy-depth-{depth}") mtp_logits_2d = None if has_mtp: @@ -1806,6 +1810,14 @@ async def async_generate_output_tokens_dynamic_batch( ) range_pop() + # Capture before update_requests (called by _dynamic_step_context_bookkeeping) + # resets num_prefill_requests to 0, which would make num_decode_requests + # always equal to the full active count. + num_decode_requests = context.num_decode_requests + if self.num_speculative_tokens > 0: + # Prefill-only batches must not have any accepted speculative tokens. + assert num_decode_requests > 0 or (self._accepted_tokens_per_request == -1).all() + if skip_bookkeeping: # _transfer_samples_to_cpu wasn't invoked on this path, so do # a one-shot D2H here to keep "sample" as a CPU tensor for @@ -1820,9 +1832,9 @@ async def async_generate_output_tokens_dynamic_batch( ret = { "accepted_tokens": ( - # Clone needed: .fill_(-1) on line 1480 would corrupt the returned value. + # Clone needed: .fill_(-1) below would corrupt the returned value. self._accepted_tokens_per_request.clone() - if self.num_speculative_tokens > 0 + if self.num_speculative_tokens > 0 and num_decode_requests > 0 else None ), "log_probs": log_probs, diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 2f03c8cb7aa..9adae9ce607 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -3325,6 +3325,150 @@ def test_speculative_decoding_chunked_prefill_and_prefix_caching(self): assert env.engine.context.active_token_count == 0 assert env.engine.context.total_request_count == 0 + @pytest.mark.internal + @pytest.mark.skipif( + not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" + ) + @torch.inference_mode() + def test_speculative_decoding_stats_exclude_prefill(self): + """Test that MTP acceptance stats are cumulative and exclude prefill requests. + + Prefill requests don't get MTP speculative proposals (MTP heads only run for + decode requests). Verify that: + 1. Stats accumulate across the engine lifetime (no reset between logging). + 2. Prefill steps don't inflate _spec_tokens_proposed. + 3. The acceptance rate reflects only decode steps. + """ + test_config = DynamicEngineTestConfig( + num_requests=0, # Added manually below to stagger prefill vs decode + min_prompt_length=4, + max_prompt_length=4, + num_tokens_to_generate=10, + num_speculative_tokens=2, + materialize_only_last_token_logits=False, + model_provider="gpt", + ) + env = self._build_test_env(test_config) + unwrapped_model = env.engine.controller.inference_wrapped_model.model + hidden_size = unwrapped_model.config.hidden_size + + # Mock forward: all tokens get the same high-probability logit so every + # speculative token is accepted (acceptance rate should be 100%). + def mock_mtp_forward(*args, **kwargs): + tokens = kwargs.get("tokens", args[0] if args else kwargs.get("input_ids")) + base_logits = torch.zeros( + tokens.size(0), + tokens.size(1), + test_config.vocab_size, + device=tokens.device, + dtype=torch.bfloat16, + ) + base_logits[:, :, 0] = 100.0 + unwrapped_model._decoder_hidden_states_cache = torch.zeros( + tokens.size(1), 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 + ) + return base_logits + + def mock_compute_mtp_single_step( + hidden_states, next_token_ids, position_ids, depth, eager=False, cache_key=None + ): + n = hidden_states.size(0) + logits = torch.zeros( + n, 1, test_config.vocab_size, device=hidden_states.device, dtype=torch.bfloat16 + ) + logits[:, :, 0] = 100.0 + return hidden_states, logits + + unwrapped_model.forward = mock_mtp_forward + unwrapped_model.compute_mtp_single_step = mock_compute_mtp_single_step + + # Verify counters start at zero. + assert sum(env.engine._spec_tokens_proposed_per_pos) == 0 + assert sum(env.engine._spec_tokens_accepted_per_pos) == 0 + assert env.engine._spec_steps == 0 + + # Add first request and run through prefill + some decode steps. + env.engine.add_request( + request_id=0, + prompt=torch.randint( + 0, test_config.vocab_size - 1, (4,), dtype=torch.int64, device='cuda' + ), + sampling_params=SamplingParams(num_tokens_to_generate=10, termination_id=-1), + ) + + # Step 1: prefill for request 0 — should NOT count as a spec step. + # The controller returns accepted_tokens=None for prefill-only batches + # (num_decode_requests == 0), so the engine must not increment any stats. + env.engine.step_modern() + proposed_after_prefill = sum(env.engine._spec_tokens_proposed_per_pos) + accepted_after_prefill = sum(env.engine._spec_tokens_accepted_per_pos) + assert proposed_after_prefill == 0, "Prefill step should not propose any spec tokens" + assert accepted_after_prefill == 0, "Prefill step should not accept any spec tokens" + assert env.engine._spec_steps == 0, "Prefill step should not count as a spec step" + + # Step 2: decode for request 0 — should count spec tokens. + env.engine.step_modern() + assert ( + sum(env.engine._spec_tokens_proposed_per_pos) > proposed_after_prefill + ), "Decode step should have incremented _spec_tokens_proposed_per_pos" + assert ( + sum(env.engine._spec_tokens_accepted_per_pos) > accepted_after_prefill + ), "With deterministic mock, decode step should have accepted spec tokens" + + # Now add a second request while request 0 is decoding. + # The next step is a mixed prefill (req 1) + decode (req 0) step. + env.engine.add_request( + request_id=1, + prompt=torch.randint( + 0, test_config.vocab_size - 1, (4,), dtype=torch.int64, device='cuda' + ), + sampling_params=SamplingParams(num_tokens_to_generate=10, termination_id=-1), + ) + + proposed_before_mixed = sum(env.engine._spec_tokens_proposed_per_pos) + env.engine.step_modern() + proposed_after_mixed = sum(env.engine._spec_tokens_proposed_per_pos) + + # In the mixed step, only the decode request (req 0) should contribute to + # proposed count, NOT the prefilling request (req 1). With 2 spec tokens and + # 1 decode request, proposed should increase by exactly 2. + proposed_delta = proposed_after_mixed - proposed_before_mixed + assert proposed_delta == test_config.num_speculative_tokens, ( + f"Mixed prefill+decode step: expected proposed delta of " + f"{test_config.num_speculative_tokens} (1 decode request), got {proposed_delta}" + ) + + # Run to completion. + while env.engine.has_unfinished_requests(): + env.engine.step_modern() + + # Stats should be cumulative (non-zero after all requests finish). + total_proposed = sum(env.engine._spec_tokens_proposed_per_pos) + total_accepted = sum(env.engine._spec_tokens_accepted_per_pos) + assert total_proposed > 0 + assert total_accepted > 0 + assert env.engine._spec_steps > 0 + + # With deterministic mock (all tokens accepted), acceptance rate should be 100%. + acceptance_rate = total_accepted / total_proposed + assert ( + acceptance_rate == 1.0 + ), f"Expected 100% acceptance with deterministic mock, got {acceptance_rate * 100:.1f}%" + + # With deterministic mock, every position should have 100% acceptance. + for pos in range(test_config.num_speculative_tokens): + assert ( + env.engine._spec_tokens_proposed_per_pos[pos] > 0 + ), f"Position {pos} should have proposals" + pos_rate = ( + env.engine._spec_tokens_accepted_per_pos[pos] + / env.engine._spec_tokens_proposed_per_pos[pos] + ) + assert pos_rate == 1.0, ( + f"Expected 100% acceptance at position {pos} with deterministic mock, " + f"got {pos_rate * 100:.1f}%" + ) + @pytest.mark.internal @pytest.mark.skipif( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" diff --git a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py index 8fd1f4a1154..8f738ceb81c 100644 --- a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py +++ b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py @@ -399,7 +399,6 @@ def test_cuda_graph_sp_padding_end_to_end(self, mtp_use_repeated_layer): ) ctrl.num_speculative_tokens = num_spec - ctrl.num_mtp_heads = num_spec ctrl._init_mtp_sampling_tensors() ctrl._mtp_token_ids_buf.zero_() ctrl._mtp_position_ids_buf.zero_() @@ -496,7 +495,6 @@ def _run_mtp(use_cuda_graph): ) ctrl.num_speculative_tokens = num_spec - ctrl.num_mtp_heads = num_spec ctrl._init_mtp_sampling_tensors() ctrl._mtp_token_ids_buf.zero_() ctrl._mtp_position_ids_buf.zero_() 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 926104eb6c7..e150d097e8d 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 @@ -1051,6 +1051,7 @@ def test_rewind_kv_cache(self, is_hybrid_model): self.setup_model( torch.float32, static=False, + mtp_num_layers=3, num_speculative_tokens=3, block_size_tokens=4, max_requests=16, @@ -1235,11 +1236,14 @@ def test_speculative_multinomial_sampling(self): (top_k > 1, top_p > 0) by flattening 3D MTP logits for torch.multinomial.""" num_spec = 3 self.setup_model( - torch.float32, static=False, num_speculative_tokens=num_spec, max_requests=2 + torch.float32, + static=False, + num_speculative_tokens=num_spec, + mtp_num_layers=num_spec, + max_requests=2, ) # Enable speculative decoding - self.text_generation_controller.num_speculative_tokens = num_spec ctx = self.text_generation_controller.inference_wrapped_model.inference_context ctx.total_request_count = 2 ctx.paused_request_count = 0 @@ -1294,6 +1298,7 @@ def test_rewind_kv_cache_with_prefix_caching_ref_counts(self): block_size_tokens=4, enable_prefix_caching=True, max_requests=16, + mtp_num_layers=2, ) ctx = self.text_generation_controller.inference_wrapped_model.inference_context @@ -1343,6 +1348,7 @@ def test_rewind_kv_cache_does_not_release_shared_prefix_blocks(self): num_speculative_tokens=3, block_size_tokens=4, max_requests=16, + mtp_num_layers=3, ) ctx = self.text_generation_controller.inference_wrapped_model.inference_context @@ -1393,7 +1399,6 @@ def test_speculative_mtp_position_ids_with_prefill(self): ) self.text_generation_controller.num_speculative_tokens = 2 - self.text_generation_controller.num_mtp_heads = 2 ctx = self.text_generation_controller.inference_wrapped_model.inference_context ctx.total_request_count = 2 ctx.paused_request_count = 0