Merge branch 'main' into main

Author: devalshahamd

tests/kernels/test_minimax_m3_amd_ops.py

@@ -424,6 +424,69 @@ def test_mxfp8_linear_emulation_bf16_at_load(
     assert _relerr(out.float(), out_ref.float()) < 2e-2
 
 
+# --------------------------------------------------------------------------- #
+# Native MXFP8 linear: the K % 128 != 0 weight is dequantized once at load
+# --------------------------------------------------------------------------- #
+@requires_gfx950
+# (5120, 576) and (5120, 288): DSv4.1 shared-expert down_proj at TP4 and TP8.
+@pytest.mark.parametrize("shape", [(5120, 576), (5120, 288), (256, 2048)])
+@torch.inference_mode()
+def test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load(shape):
+    """``dot_scaled`` tiles K by 128, so an unaligned weight must not reach it.
+
+    ``process_weights_after_loading`` converts those to BF16 once; a weight that
+    never went through it has to fall back per step rather than take the
+    invalid path, which is what keying the dispatch on the weight dtype alone
+    would do.
+    """
+    from vllm.model_executor.kernels.linear.mxfp8 import rocm_native
+    from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+        Mxfp8LinearLayerConfig,
+    )
+
+    N, K = shape
+    aligned = K % 128 == 0
+    torch.manual_seed(0)
+    w_bf16 = torch.randn(N, K, device=DEVICE, dtype=torch.bfloat16) * 0.1
+    w_fp8, w_scale = _mxfp8_e4m3_quantize_torch(w_bf16, is_sf_swizzled_layout=False)
+    x = torch.randn(7, K, device=DEVICE, dtype=torch.bfloat16) * 0.5
+    out_ref = torch.nn.functional.linear(x, dequant_mxfp8_to_bf16(w_fp8, w_scale))
+    # dot_scaled re-quantizes the activation, the dequant paths do not.
+    tol = 5e-2 if aligned else 2e-3
+
+    def make_layer():
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(w_fp8.clone(), requires_grad=False)
+        layer.weight_scale = torch.nn.Parameter(w_scale.clone(), requires_grad=False)
+        return layer
+
+    kernel = rocm_native.RocmDotScaledMxfp8LinearKernel(Mxfp8LinearLayerConfig())
+    layer = make_layer()
+    kernel.process_weights_after_loading(layer)
+    assert (layer.weight.element_size() >= 2) is not aligned
+    assert _relerr(kernel.apply_weights(layer, x).float(), out_ref.float()) < tol
+
+    # An unaligned weight that is still FP8 must not reach dot_scaled.
+    unprocessed = make_layer()
+    unprocessed.weight_scale = torch.nn.Parameter(
+        w_scale[:N, : K // 32].contiguous(), requires_grad=False
+    )
+    calls = []
+    real = rocm_native._mxfp8_dot_scaled_linear
+
+    def _spy(*a, **k):
+        calls.append(1)
+        return real(*a, **k)
+
+    rocm_native._mxfp8_dot_scaled_linear = _spy
+    try:
+        out = kernel.apply_weights(unprocessed, x)
+    finally:
+        rocm_native._mxfp8_dot_scaled_linear = real
+    assert bool(calls) is aligned
+    assert _relerr(out.float(), out_ref.float()) < tol
+
+
 # ── EP expert_mask handling for the FlyDSL (AITER_MXFP8) MoE ────────────────
 # The map/mask choice lives in ``RoutedExperts.expert_map``: it hands AITER
 # experts (``consumes_expert_mask``) the precomputed 0/1 ``expert_mask`` and

tests/test_request_input_bounds.py

@@ -12,6 +12,7 @@
 import sys
 from collections.abc import Callable
 from typing import Protocol
+from unittest.mock import Mock
 
 import pytest
 from pydantic import ValidationError
@@ -27,10 +28,39 @@
 from vllm.exceptions import VLLMValidationError
 from vllm.sampling_params import BeamSearchParams
 from vllm.tokenizers import deepseek_v4_encoding, deepseek_v32_encoding
+from vllm.v1.engine.input_processor import InputProcessor
 
 pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
 
 
+@pytest.mark.parametrize("offset", [-1, 0, 2, 3, 4, 2**127])
+def test_routed_experts_offset_validated_before_engine_submission(offset):
+    """Reject invalid offsets even when routed-expert output is disabled."""
+    processor = Mock(spec=InputProcessor)
+    processor.tokenizer = None
+    processor.generation_config_fields = {}
+    processor.renderer = Mock()
+    processor.renderer.get_eos_token_id.return_value = None
+    processor.vllm_config = Mock()
+    params = SamplingParams(routed_experts_prompt_start=offset)
+    prompt = {"type": "token", "prompt_token_ids": [1, 2, 3]}
+
+    if not 0 <= offset <= 3:
+        with pytest.raises(VLLMValidationError, match="routed_experts_prompt_start"):
+            InputProcessor.process_inputs(
+                processor, "invalid", prompt, params, ("generate",)
+            )
+        params.routed_experts_prompt_start = 0
+
+    request = InputProcessor.process_inputs(
+        processor, "valid", prompt, params, ("generate",)
+    )
+    assert (
+        request.sampling_params.routed_experts_prompt_start
+        == params.routed_experts_prompt_start
+    )
+
+
 # --- Stop strings: public requests cap the number of stop strings ---------
 
 

tests/v1/core/test_scheduler.py

@@ -51,6 +51,7 @@
     ECConnectorOutput,
     KVConnectorOutput,
     ModelRunnerOutput,
+    RoutedExpertsLists,
     SamplingMaskLists,
     make_empty_encoder_model_runner_output,
 )
@@ -131,6 +132,44 @@ def test_add_requests():
         assert len(scheduler.waiting) == i + 1
 
 
+def test_routed_experts_prompt_start_at_prompt_end():
+    """A prompt-end offset should return empty routing data, not crash."""
+    scheduler = create_scheduler()
+    (request,) = create_requests(num_requests=1)
+    request.sampling_params.routed_experts_prompt_start = request.num_prompt_tokens
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    num_tokens = scheduler_output.num_scheduled_tokens[request.request_id]
+
+    manager = Mock()
+    manager.routed_experts_by_slot = np.empty((0, 1, 1), dtype=np.uint8)
+    manager.get.return_value = np.empty((0, 1, 1), dtype=np.uint8)
+    scheduler.enable_return_routed_experts = True
+    scheduler.routed_experts_mgr = manager
+    scheduler._re_block_ids = {request.request_id: []}
+
+    outputs = scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[0]],
+            logprobs=None,
+            prompt_logprobs_dict={},
+            pooler_output=[],
+            routed_experts=RoutedExpertsLists(
+                routing_data=np.zeros((num_tokens, 1, 1), dtype=np.uint8),
+                slot_mapping=np.arange(num_tokens),
+            ),
+        ),
+    )
+
+    manager.get.assert_called_once_with(
+        [], request.num_prompt_tokens, token_start=request.num_prompt_tokens
+    )
+    assert outputs[request.client_index].outputs[0].routed_experts.shape == (0, 1, 1)
+
+
 def test_finish_request():
     scheduler = create_scheduler()
     requests = create_requests(num_requests=10)

vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

@@ -4,9 +4,9 @@
 
 Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
 activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
-matrix cores. Falls back (via the kernel selector) to the BF16
-``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
-``K % 128 != 0``.
+matrix cores. ``dot_scaled`` tiles K by 128; a weight whose K is not a multiple
+of that is dequantized to BF16 once at load and served by a plain linear
+instead, since ``can_implement`` does not filter on K.
 """
 
 import torch
@@ -194,6 +194,9 @@ def _select_cfg(M, N, K):
     return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
 
 
+_DOT_SCALED_K_ALIGN = 128
+
+
 class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
     """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
 
@@ -219,6 +222,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         N, K = weight.shape
         scale_k = K // MXFP8_BLOCK_SIZE
         weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
+        if K % _DOT_SCALED_K_ALIGN != 0:
+            weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
         layer.weight = Parameter(weight.contiguous(), requires_grad=False)
         layer.weight_scale = Parameter(weight_scale, requires_grad=False)
 
@@ -235,10 +240,13 @@ def apply_weights(
             )
         out_shape = (*x.shape[:-1], layer.weight.shape[0])
         x2d = x.reshape(-1, x.shape[-1])
-        if x2d.shape[-1] % 128 == 0:
+        if layer.weight.element_size() >= 2:
+            out = torch.nn.functional.linear(x2d, layer.weight.to(x.dtype))
+        elif x2d.shape[-1] % _DOT_SCALED_K_ALIGN == 0:
             out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
         else:
-            # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
+            # process_weights_after_loading dequantizes these weights once, so
+            # this runs only if it did not; dot_scaled would be invalid here.
             w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
             out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
         out = out.reshape(out_shape)

vllm/v1/core/sched/scheduler.py

@@ -2101,7 +2101,7 @@ def update_from_output(
                         prompt_start = (
                             request.sampling_params.routed_experts_prompt_start
                         )
-                        assert prompt_start < request.num_prompt_tokens
+                        assert prompt_start <= request.num_prompt_tokens
                     else:
                         prompt_start = 0
                     routed_experts = self.routed_experts_mgr.get(

vllm/v1/engine/input_processor.py

@@ -354,12 +354,21 @@ def process_inputs(
         if isinstance(params, SamplingParams):
             # TODO: can we avoid cloning here in multiproc case?
             sampling_params = params.clone()
+            prompt_len = length_from_prompt_token_ids_or_embeds(
+                prompt_token_ids, prompt_embeds
+            )
+            if not 0 <= sampling_params.routed_experts_prompt_start <= prompt_len:
+                raise VLLMValidationError(
+                    f"routed_experts_prompt_start must be between 0 and "
+                    f"the prompt length ({prompt_len}), inclusive.",
+                    parameter="routed_experts_prompt_start",
+                    value=sampling_params.routed_experts_prompt_start,
+                )
             # If unset max tokens, then generate up to the max_model_len.
             if sampling_params.max_tokens is None:
-                seq_len = length_from_prompt_token_ids_or_embeds(
-                    prompt_token_ids, prompt_embeds
+                sampling_params.max_tokens = (
+                    self.model_config.max_model_len - prompt_len
                 )
-                sampling_params.max_tokens = self.model_config.max_model_len - seq_len
 
             sampling_params.update_from_generation_config(
                 self.generation_config_fields,
@@ -368,12 +377,7 @@ def process_inputs(
             if self.tokenizer is not None:
                 sampling_params.update_from_tokenizer(self.tokenizer)
             if sampling_params.trace_decode_token_ids:
-                self._normalize_trace_replay_params(
-                    sampling_params,
-                    length_from_prompt_token_ids_or_embeds(
-                        prompt_token_ids, prompt_embeds
-                    ),
-                )
+                self._normalize_trace_replay_params(sampling_params, prompt_len)
         else:
             pooling_params = params.clone()