[BugFix][Spec Decoding] Fix negative accepted tokens metric crash (#33729)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

tests/kernels/moe/test_moe.py

@@ -715,7 +715,7 @@ def test_mixtral_moe(
 
         # need to override the forward context for unittests, otherwise it assumes
         # we're running the model forward pass (the model specified in vllm_config)
-        get_forward_context().remaining_moe_layers = None
+        get_forward_context().all_moe_layers = None
 
         # Run forward passes for both MoE blocks
         hf_states, _ = hf_moe.forward(hf_inputs)

tests/v1/core/test_scheduler.py

@@ -870,6 +870,66 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
         assert stats.num_accepted_tokens_per_pos == expected[3]
 
 
+def test_spec_decoding_stats_empty_output():
+    """Test that spec decoding stats handle empty output tokens gracefully.
+
+    This is a regression test for a bug where empty sampled_token_ids
+    would cause num_accepted = len([]) - 1 = -1, leading to a
+    ValueError when incrementing a Prometheus counter with a negative value.
+    """
+    num_spec_tokens = 3
+    scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens)
+    requests = create_requests(num_requests=1, num_tokens=1)
+    request = requests[0]
+    req_id = request.request_id
+
+    scheduler.add_request(request)
+
+    # Initial schedule (prefill)
+    output = scheduler.schedule()
+    assert len(output.scheduled_new_reqs) == 1
+
+    # Complete the prefill with a sampled token
+    model_runner_output = ModelRunnerOutput(
+        req_ids=[req_id],
+        req_id_to_index={req_id: 0},
+        sampled_token_ids=[[0]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+    scheduler.update_from_output(output, model_runner_output)
+
+    # Add draft tokens for speculation
+    draft_token_ids = DraftTokenIds([req_id], [[1, 2, 3]])
+    scheduler.update_draft_token_ids(draft_token_ids)
+
+    # Schedule the speculated tokens for validation
+    output = scheduler.schedule()
+    assert req_id in output.scheduled_spec_decode_tokens
+    assert len(output.scheduled_spec_decode_tokens[req_id]) == 3
+
+    # Simulate empty output tokens (e.g., due to request abortion or error)
+    # This would previously cause num_accepted = -1 and crash
+    model_runner_output = ModelRunnerOutput(
+        req_ids=[req_id],
+        req_id_to_index={req_id: 0},
+        sampled_token_ids=[[]],  # Empty output tokens
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+
+    # This should not raise an error
+    engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
+
+    # Spec decoding stats should be None since no tokens were generated
+    scheduler_stats = (
+        engine_core_outputs[0].scheduler_stats if engine_core_outputs else None
+    )
+    assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None
+
+
 def _assert_right_scheduler_output(
     output: SchedulerOutput,
     num_requests: int,

vllm/config/compilation.py

@@ -581,6 +581,24 @@ class CompilationConfig:
     local_cache_dir: str = field(default=None, init=False)  # type: ignore
     """local cache dir for each rank"""
 
+    fast_moe_cold_start = True
+    """Optimization for fast MOE cold start.
+
+    This is a bit of a hack that assumes that:
+    1. the only decoder forward pass being run is the current model
+    2. the decoder forward pass runs all of the MOEs in the order in which they
+       are initialized
+
+    When the above two conditions hold, this option greatly decreases cold start
+    time for MOE models.
+
+    If the above two conditions don't hold, then this option will lead to silent
+    incorrectness. The only condition in which this doesn't hold is speculative
+    decoding, where there is a draft model that may have MOEs in them.
+
+    NB: We're working on a longer-term solution that doesn't need these assumptions.
+    """
+
     # keep track of enabled and disabled custom ops
     enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
     """custom ops that are enabled"""
@@ -596,6 +614,10 @@ class CompilationConfig:
     Map from layer name to layer objects that need to be accessed outside
     model code, e.g., Attention, FusedMOE when dp_size>1."""
 
+    static_all_moe_layers: list[str] = field(default_factory=list, init=False)
+    """The names of all the MOE layers in the model
+    """
+
     # Attention ops; used for piecewise cudagraphs
     # Use PyTorch operator format: "namespace::name"
     _attention_ops: ClassVar[list[str]] = [

vllm/forward_context.py

@@ -217,9 +217,11 @@ class ForwardContext:
     # the graph.
     #
     # The workaround is to store a list of the strings that each of those
-    # custom ops needs, in reverse order, in the ForwardContext.
+    # custom ops needs in the ForwardContext (all_moe_layers)
+    # as well as a counter (moe_layer_index).
     # The ForwardContext object is alive for the duration of the forward pass.
-    # When the custom op needs the string, pop the string from this list.
+    # When the custom op needs a layer string, get the next string
+    # from all_moe_layers and increment the counter.
     #
     # This assumes that the custom operators will always be executed in
     # order and that torch.compile will not try to reorder these
@@ -233,7 +235,8 @@ class ForwardContext:
     #
     # If this value is None (like in some tests), then we end up baking the string
     # into the graph. Otherwise, the moe custom ops will pop a string from this list.
-    remaining_moe_layers: list[str] | None = None
+    all_moe_layers: list[str] | None = None
+    moe_layer_index: int = 0
 
     additional_kwargs: dict[str, Any] = field(default_factory=dict)
 
@@ -271,17 +274,22 @@ def create_forward_context(
     additional_kwargs: dict[str, Any] | None = None,
     skip_compiled: bool = False,
 ):
-    no_compile_layers = vllm_config.compilation_config.static_forward_context
-    from vllm.model_executor.layers.fused_moe.layer import FusedMoE
-
-    remaining_moe_layers = [
-        name for name, layer in no_compile_layers.items() if isinstance(layer, FusedMoE)
-    ]
-    remaining_moe_layers.reverse()
+    if vllm_config.compilation_config.fast_moe_cold_start:
+        if vllm_config.speculative_config is None:
+            all_moe_layers = vllm_config.compilation_config.static_all_moe_layers
+        else:
+            logger.warning_once(
+                "vllm_config.compilation_config.fast_moe_cold_start is not "
+                "compatible with speculative decoding so we are ignoring "
+                "fast_moe_cold_start."
+            )
+            all_moe_layers = None
+    else:
+        all_moe_layers = None
 
     return ForwardContext(
-        no_compile_layers=no_compile_layers,
-        remaining_moe_layers=remaining_moe_layers,
+        no_compile_layers=vllm_config.compilation_config.static_forward_context,
+        all_moe_layers=all_moe_layers,
         virtual_engine=virtual_engine,
         attn_metadata=attn_metadata,
         slot_mapping=slot_mapping or {},

vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py

@@ -69,9 +69,14 @@ def _supports_routing_method(
             RoutingMethodType.RenormalizeNaive,
         ]
     elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
-        # NOTE(rob): kernel requires Llama4.
-        return routing_method == RoutingMethodType.Llama4
-
+        # NOTE(dbari): as above, potentially allow others here.
+        return routing_method in [
+            RoutingMethodType.Llama4,
+            # NOTE(mgoin): Disabled to investigate accuracy issues.
+            # See https://github.com/vllm-project/vllm/issues/33532
+            # RoutingMethodType.Renormalize,
+            # RoutingMethodType.RenormalizeNaive,
+        ]
     else:
         raise ValueError("Unsupported quantization scheme.")
 
@@ -81,7 +86,23 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo
     return not moe_parallel_config.enable_eplb
 
 
-def is_supported_config_trtllm(
+def _supports_router_logits_dtype(
+    router_logits_dtype: torch.dtype | None,
+    routing_method: RoutingMethodType,
+) -> bool:
+    """
+    The FlashInfer TRTLLM FP8 kernel expects bfloat16 router_logits by default.
+    Only DeepSeekV3 routing supports float32 router_logits (which is converted
+    internally in the kernel).
+    """
+    if router_logits_dtype == torch.float32:
+        # Only DeepSeekV3 routing handles float32 logits
+        # https://github.com/flashinfer-ai/flashinfer/issues/2469
+        return routing_method == RoutingMethodType.DeepSeekV3
+    return True
+
+
+def is_supported_config_trtllm_fp8(
     moe_config: FusedMoEConfig,
     weight_key: QuantKey | None,
     activation_key: QuantKey | None,
@@ -110,13 +131,17 @@ def is_supported_config_trtllm(
         return False, _make_reason("routing method")
     elif activation_format != mk.FusedMoEActivationFormat.Standard:
         return False, _make_reason("activation format")
+    elif not _supports_router_logits_dtype(
+        moe_config.router_logits_dtype, moe_config.routing_method
+    ):
+        return False, _make_reason("float32 router_logits with non-DeepSeekV3 routing")
 
     return True, None
 
 
 def flashinfer_fused_moe_blockscale_fp8(
     routing_logits: torch.Tensor,
-    routing_bias: torch.Tensor,
+    routing_bias: torch.Tensor | None,
     x: torch.Tensor,
     w13_weight: torch.Tensor,
     w13_weight_scale_inv: torch.Tensor,
@@ -130,7 +155,7 @@ def flashinfer_fused_moe_blockscale_fp8(
     expert_offset: int,
     local_num_experts: int,
     block_shape: list[int],
-    routing_method_type: int = int(RoutingMethodType.DeepSeekV3),
+    routing_method_type: int,
     routed_scaling: float | None = 1.0,
 ) -> torch.Tensor:
     from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
@@ -143,6 +168,13 @@ def flashinfer_fused_moe_blockscale_fp8(
     # Routing kernel expects #experts <= #threads 512
     assert global_num_experts <= 512
 
+    # The DeepSeekV3 routing method requires float32 router logits.
+    if routing_method_type == RoutingMethodType.DeepSeekV3:
+        routing_logits = routing_logits.to(torch.float32)
+
+    if routing_bias is not None:
+        routing_bias = routing_bias.to(x.dtype)
+
     a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])
     # NOTE: scales of hidden states have to be transposed!
     a_sf_t = a_sf.t().contiguous()
@@ -170,7 +202,7 @@ def flashinfer_fused_moe_blockscale_fp8(
 
 def flashinfer_fused_moe_blockscale_fp8_fake(
     routing_logits: torch.Tensor,
-    routing_bias: torch.Tensor,
+    routing_bias: torch.Tensor | None,
     x: torch.Tensor,
     w13_weight: torch.Tensor,
     w13_weight_scale_inv: torch.Tensor,

vllm/model_executor/layers/fused_moe/layer.py

@@ -407,6 +407,7 @@ class FusedMoE(CustomOp):
         if prefix in compilation_config.static_forward_context:
             raise ValueError("Duplicate layer name: {}".format(prefix))
         compilation_config.static_forward_context[prefix] = self
+        compilation_config.static_all_moe_layers.append(prefix)
         self.layer_name = prefix
 
         self.enable_eplb = enable_eplb
@@ -1572,7 +1573,7 @@ class FusedMoE(CustomOp):
             # Can be unavailable or None in unittests
             if (
                 is_forward_context_available()
-                and get_forward_context().remaining_moe_layers is not None
+                and get_forward_context().all_moe_layers is not None
             ):
                 return "from_forward_context"
             return self.layer_name
@@ -1991,13 +1992,17 @@ class FusedMoE(CustomOp):
 def get_layer_from_name(layer_name: str) -> FusedMoE:
     forward_context: ForwardContext = get_forward_context()
     if layer_name == "from_forward_context":
-        if not forward_context.remaining_moe_layers:
+        all_moe_layers = forward_context.all_moe_layers
+        assert all_moe_layers is not None
+        moe_layer_index = forward_context.moe_layer_index
+        if moe_layer_index >= len(all_moe_layers):
             raise AssertionError(
-                "We expected the number of MOE layers in `remaining_moe_layers` "
+                "We expected the number of MOE layers in `all_moe_layers` "
                 "to be equal to the number of "
                 "{vllm.moe_forward, vllm.moe_forward_shared} calls."
             )
-        layer_name = forward_context.remaining_moe_layers.pop()
+        layer_name = all_moe_layers[moe_layer_index]
+        forward_context.moe_layer_index += 1
     self = cast(FusedMoE, forward_context.no_compile_layers[layer_name])
     return self
 

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -15,7 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     fp8_w8a16_moe_quant_config,
 )
 from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (
-    is_supported_config_trtllm,
+    is_supported_config_trtllm_fp8,
 )
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
     MoEPrepareAndFinalizeNoEP,
@@ -212,7 +212,7 @@ def select_fp8_moe_backend(
 
             if fi_backend == FlashinferMoeBackend.TENSORRT_LLM:
                 backend = Fp8MoeBackend.FLASHINFER_TRTLLM
-                supported, reason = is_supported_config_trtllm(
+                supported, reason = is_supported_config_trtllm_fp8(
                     config, weight_key, activation_key, activation_format
                 )
                 if supported:
@@ -239,7 +239,7 @@ def select_fp8_moe_backend(
             ]:
                 if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
                     k_cls = None
-                    supported, reason = is_supported_config_trtllm(
+                    supported, reason = is_supported_config_trtllm_fp8(
                         config,
                         weight_key,
                         activation_key,
@@ -308,7 +308,7 @@ def select_fp8_moe_backend(
     for backend in AVAILABLE_BACKENDS:
         if backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
             k_cls = None
-            supported, reason = is_supported_config_trtllm(
+            supported, reason = is_supported_config_trtllm_fp8(
                 config,
                 weight_key,
                 activation_key,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe import (
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
     FusedMoEQuantConfig,
-    RoutingMethodType,
     int4_w4a16_moe_quant_config,
     int4_w4afp8_moe_quant_config,
     int8_w8a8_moe_quant_config,
@@ -1072,17 +1071,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
         if self.block_quant:
             import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe  # noqa: E501, F401
 
-            e_score_correction_bias = (
-                layer.e_score_correction_bias.to(x.dtype)
-                if layer.e_score_correction_bias is not None
-                else None
-            )
-            routing_method_type = layer.routing_method_type
             return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
-                routing_logits=router_logits.to(torch.float32)
-                if routing_method_type == RoutingMethodType.DeepSeekV3
-                else router_logits,
-                routing_bias=e_score_correction_bias,
+                routing_logits=router_logits,
+                routing_bias=layer.e_score_correction_bias,
                 x=x,
                 w13_weight=layer.w13_weight,
                 w13_weight_scale_inv=layer.w13_weight_scale,
@@ -1096,7 +1087,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
                 expert_offset=layer.ep_rank * layer.local_num_experts,
                 local_num_experts=layer.local_num_experts,
                 block_shape=self.weight_block_size,
-                routing_method_type=routing_method_type,
+                routing_method_type=layer.routing_method_type,
                 routed_scaling=layer.routed_scaling_factor,
             )
         else:

vllm/model_executor/layers/quantization/fp8.py

@@ -26,7 +26,6 @@ from vllm.model_executor.layers.fused_moe import (
 )
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
-    RoutingMethodType,
 )
 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
 from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
@@ -990,17 +989,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
         if self.block_quant:
             import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe  # noqa: E501, F401
 
-            e_score_correction_bias = (
-                layer.e_score_correction_bias.to(x.dtype)
-                if layer.e_score_correction_bias is not None
-                else None
-            )
-            routing_method_type = layer.routing_method_type
             return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
-                routing_logits=router_logits.to(torch.float32)
-                if routing_method_type == RoutingMethodType.DeepSeekV3
-                else router_logits,
-                routing_bias=e_score_correction_bias,
+                routing_logits=router_logits,
+                routing_bias=layer.e_score_correction_bias,
                 x=x,
                 w13_weight=layer.w13_weight,
                 w13_weight_scale_inv=layer.w13_weight_scale_inv,
@@ -1014,7 +1005,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 expert_offset=layer.ep_rank * layer.local_num_experts,
                 local_num_experts=layer.local_num_experts,
                 block_shape=self.weight_block_size,
-                routing_method_type=routing_method_type,
+                routing_method_type=layer.routing_method_type,
                 routed_scaling=layer.routed_scaling_factor,
             )
         else:

vllm/model_executor/models/minimax_m2.py

@@ -107,6 +107,7 @@ class MiniMaxM2MoE(nn.Module):
             renormalize=True,
             quant_config=quant_config,
             prefix=f"{prefix}.experts",
+            router_logits_dtype=torch.float32,
         )
 
         self.gate = ReplicatedLinear(

vllm/v1/core/sched/scheduler.py

@@ -1284,7 +1284,7 @@ class Scheduler(SchedulerInterface):
             scheduled_spec_token_ids = (
                 scheduler_output.scheduled_spec_decode_tokens.get(req_id)
             )
-            if scheduled_spec_token_ids:
+            if scheduled_spec_token_ids and generated_token_ids:
                 num_draft_tokens = len(scheduled_spec_token_ids)
                 num_accepted = len(generated_token_ids) - 1
                 num_rejected = num_draft_tokens - num_accepted