Merge branch 'main' into wentao-epd-support-for-MRv2

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/lm-eval-harness/configs/Meta-Llama-4-Maverick-17B-128E-Instruct-FP8.yaml

@@ -1,6 +1,9 @@
 # For hf script, without -t option (tensor parallel size).
 # bash .buildkite/lm-eval-harness/run-lm-eval-mmlupro-vllm-baseline.sh -m meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8 -l 250 -t 8 -f 5
 model_name: "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
+required_gpu_arch:
+  - gfx942
+  - gfx950
 tasks:
 - name: "mmlu_pro"
   metrics:

.buildkite/lm-eval-harness/configs/Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml

@@ -1,6 +1,9 @@
 # For vllm script, with -t option (tensor parallel size)
 # bash .buildkite/lm-eval-harness/run-lm-eval-gsm-vllm-baseline.sh -m RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic -l 1319 -t 1
 model_name: "RedHatAI/Qwen2.5-VL-3B-Instruct-FP8-Dynamic"
+required_gpu_arch:
+  - gfx942
+  - gfx950
 tasks:
 - name: "gsm8k"
   metrics:

.buildkite/lm-eval-harness/configs/Qwen3-235B-A22B-Instruct-2507-FP8.yaml

@@ -1,4 +1,7 @@
 model_name: "Qwen/Qwen3-235B-A22B-Instruct-2507-FP8"
+required_gpu_arch:
+  - gfx942
+  - gfx950
 tasks:
   - name: "mmlu_pro"
     metrics:

.buildkite/lm-eval-harness/configs/models-small-rocm.txt

@@ -1,5 +1,6 @@
 Qwen2.5-1.5B-Instruct.yaml
 Meta-Llama-3.2-1B-Instruct-INT8-compressed-tensors.yaml
+Meta-Llama-3-8B-Instruct-INT8-compressed-tensors-asym.yaml
 Meta-Llama-3-8B-Instruct-nonuniform-compressed-tensors.yaml
 Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
 Qwen1.5-MoE-W4A16-compressed-tensors.yaml

.buildkite/lm-eval-harness/test_lm_eval_correctness.py

@@ -13,6 +13,7 @@ import os
 from contextlib import contextmanager
 
 import lm_eval
+import pytest
 import yaml
 
 from vllm.platforms import current_platform
@@ -89,9 +90,40 @@ def launch_lm_eval(eval_config, tp_size):
     return results
 
 
+def _check_rocm_gpu_arch_requirement(eval_config):
+    """Skip the test if the model requires a ROCm GPU arch not present.
+
+    Model YAML configs can specify::
+
+        required_gpu_arch:
+          - gfx942
+          - gfx950
+
+    The check only applies on ROCm.  On other platforms (e.g. CUDA) the
+    field is ignored so that shared config files work for both NVIDIA and
+    AMD CI pipelines.
+    """
+    required_archs = eval_config.get("required_gpu_arch")
+    if not required_archs:
+        return
+
+    if not current_platform.is_rocm():
+        return
+
+    from vllm.platforms.rocm import _GCN_ARCH  # noqa: E402
+
+    if not any(arch in _GCN_ARCH for arch in required_archs):
+        pytest.skip(
+            f"Model requires GPU arch {required_archs}, "
+            f"but detected arch is '{_GCN_ARCH}'"
+        )
+
+
 def test_lm_eval_correctness_param(config_filename, tp_size):
     eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
 
+    _check_rocm_gpu_arch_requirement(eval_config)
+
     results = launch_lm_eval(eval_config, tp_size)
 
     rtol = eval_config.get("rtol", DEFAULT_RTOL)

.buildkite/test-amd.yaml

@@ -751,6 +751,7 @@ steps:
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
   agent_pool: mi250_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/
@@ -2035,7 +2036,6 @@ steps:
   timeout_in_minutes: 38
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
   agent_pool: mi325_1
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/
@@ -2690,6 +2690,24 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
 
 
+- label: LM Eval Small Models (MI325) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
+  agent_pool: mi325_1
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt
+
+
 - label: LM Eval Small Models (B200-MI325) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]

.buildkite/test_areas/basic_correctness.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Basic Correctness
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/basic_correctness/test_basic_correctness

.buildkite/test_areas/benchmarks.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Benchmarks CLI Test
   timeout_in_minutes: 20
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/benchmarks/

.buildkite/test_areas/cuda.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Platform Tests (CUDA)
   timeout_in_minutes: 15
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/cuda

.buildkite/test_areas/engine.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Engine
   timeout_in_minutes: 15
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/engine
@@ -25,6 +26,7 @@ steps:
 
 - label: e2e Scheduling (1 GPU)
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
     - vllm/v1/
     - tests/v1/e2e/general/

.buildkite/test_areas/entrypoints.yaml

@@ -61,6 +61,7 @@ steps:
 
 - label: Entrypoints Integration (API Server openai - Part 3)
   timeout_in_minutes: 50
+  device: h200_18gb
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/
@@ -105,6 +106,7 @@ steps:
 
 - label: OpenAI API Correctness
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
   - csrc/
   - vllm/entrypoints/openai/

.buildkite/test_areas/expert_parallelism.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: EPLB Algorithm
   timeout_in_minutes: 15
+  device: h200_18gb
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/distributed/eplb

.buildkite/test_areas/kernels.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: vLLM IR Tests
   timeout_in_minutes: 10
+  device: h200_18gb
   working_dir: "/vllm-workspace/"
   source_file_dependencies:
     - vllm/ir
@@ -179,3 +180,21 @@ steps:
     - pytest -v -s kernels/moe/test_flashinfer_moe.py
     - pytest -v -s kernels/moe/test_nvfp4_moe.py
     - pytest -v -s kernels/moe/test_ocp_mx_moe.py
+
+
+- label: Kernels FusedMoE Layer Test (2 H100s)
+  timeout_in_minutes: 90
+  device: h100
+  num_devices: 2
+  optional: true
+  commands:
+    - pytest -v -s kernels/moe/test_moe_layer.py
+
+
+- label: Kernels FusedMoE Layer Test (2 B200s)
+  timeout_in_minutes: 90
+  device: b200
+  num_devices: 2
+  optional: true
+  commands:
+    - pytest -v -s kernels/moe/test_moe_layer.py

.buildkite/test_areas/misc.yaml

@@ -19,6 +19,7 @@ steps:
 
 - label: V1 Sample + Logits
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
     - vllm/
     - tests/v1/sample
@@ -86,6 +87,7 @@ steps:
 
 - label: Regression
   timeout_in_minutes: 20
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/test_regression

.buildkite/test_areas/models_basic.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Basic Models Tests (Initialization)
   timeout_in_minutes: 45
+  device: h200_18gb
   torch_nightly: true
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/models_language.yaml

@@ -67,6 +67,7 @@ steps:
 
 - label: Language Models Test (PPL)
   timeout_in_minutes: 110
+  device: h200_18gb
   optional: true
   source_file_dependencies:
   - vllm/
@@ -90,6 +91,7 @@ steps:
 
 - label: Language Models Test (MTEB)
   timeout_in_minutes: 110
+  device: h200_18gb
   optional: true
   source_file_dependencies:
   - vllm/

.buildkite/test_areas/models_multimodal.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
   timeout_in_minutes: 45
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
@@ -19,6 +20,7 @@ steps:
 
 - label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
   timeout_in_minutes: 45
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
@@ -77,6 +79,7 @@ steps:
 
 - label: Multi-Modal Processor # 44min
   timeout_in_minutes: 60
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal
@@ -131,6 +134,7 @@ steps:
 
 - label: Multi-Modal Models (Extended Pooling)
   optional: true
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/models/multimodal/pooling

.buildkite/test_areas/pytorch.yaml

@@ -49,6 +49,7 @@ steps:
 
 - label: PyTorch Fullgraph
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
   - vllm/
   - tests/compile
@@ -60,6 +61,7 @@ steps:
   # if this test fails, it means the nightly torch version is not compatible with some
   # of the dependencies. Please check the error message and add the package to whitelist
   # in /vllm/tools/pre_commit/generate_nightly_torch_test.py
+  device: h200_18gb
   soft_fail: true
   source_file_dependencies:
   - requirements/nightly_torch_test.txt

.buildkite/test_areas/ray_compat.yaml

@@ -7,6 +7,7 @@ steps:
   # If this fails, it means the PR introduces a dependency that
   # conflicts with Ray's dependency constraints.
   # See https://github.com/vllm-project/vllm/issues/33599
+  device: h200_18gb
   soft_fail: true
   timeout_in_minutes: 10
   source_file_dependencies:

.buildkite/test_areas/spec_decode.yaml

@@ -4,6 +4,7 @@ depends_on:
 steps:
 - label: Spec Decode Eagle
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
@@ -13,6 +14,7 @@ steps:
 
 - label: Spec Decode Speculators + MTP
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
@@ -23,6 +25,7 @@ steps:
 
 - label: Spec Decode Ngram + Suffix
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/
@@ -32,6 +35,7 @@ steps:
 
 - label: Spec Decode Draft Model
   timeout_in_minutes: 30
+  device: h200_18gb
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/v1/worker/gpu/spec_decode/

csrc/cpu/cpu_attn_vec.hpp

@@ -53,7 +53,7 @@ class TileGemm82 {
                          const int64_t ldb, const int64_t ldc,
                          const int32_t block_size, const int32_t dynamic_k_size,
                          const bool accum_c) {
-    static_assert(0 < M <= 8);
+    static_assert(0 < M && M <= 8);
     using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
 
     kv_cache_t* __restrict__ curr_b_0 = b_tile;

csrc/cpu/cpu_attn_vec16.hpp

@@ -68,7 +68,7 @@ class TileGemm161 {
                          const int64_t ldb, const int64_t ldc,
                          const int32_t block_size, const int32_t dynamic_k_size,
                          const bool accum_c) {
-    static_assert(0 < M <= 16);
+    static_assert(0 < M && M <= 16);
     using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
 
     kv_cache_t* __restrict__ curr_b_0 = b_tile;

csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp

@@ -39,7 +39,7 @@ class TileGemm82 {
 
   template <int32_t M>
   static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
-    static_assert(0 < M <= 8);
+    static_assert(0 < M && M <= 8);
     using load_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
 
     scalar_t* __restrict__ curr_b_0 = b_ptr;

docs/design/moe_kernel_features.md

@@ -57,8 +57,8 @@ Modular kernels are supported by the following `FusedMoEMethodBase` classes.
 
 - [`ModelOptFp8MoEMethod`][vllm.model_executor.layers.quantization.modelopt.ModelOptFp8MoEMethod]
 - [`Fp8MoEMethod`][vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod]
-- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoEMethod]
-- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsW8A8Fp8MoEMethod]
+- [`CompressedTensorsW4A4Nvfp4MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4.CompressedTensorsW4A4Nvfp4MoEMethod]
+- [`CompressedTensorsW8A8Fp8MoEMethod`][vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_fp8.CompressedTensorsW8A8Fp8MoEMethod]
 - [`Mxfp4MoEMethod`][vllm.model_executor.layers.quantization.mxfp4.Mxfp4MoEMethod]
 - [`UnquantizedFusedMoEMethod`][vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod]
 

docs/mkdocs/hooks/autoref_code.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+MkDocs hook to automatically convert inline code references to API doc links.
+
+For example, `WeightTransferConfig` becomes
+[`WeightTransferConfig`][vllm.config.WeightTransferConfig]
+
+This works with the `autorefs` plugin to create clickable cross-references
+to API documentation pages generated by `mkdocstrings`.
+
+The hook builds an index of all documented public Python names (classes and
+functions with docstrings) from the vllm package at startup using AST parsing,
+then substitutes matching inline code spans on each page. Names without
+docstrings are excluded because mkdocstrings will not generate a page for them.
+"""
+
+import ast
+import logging
+from pathlib import Path
+
+import regex as re
+from mkdocs.config.defaults import MkDocsConfig
+from mkdocs.structure.files import Files
+from mkdocs.structure.pages import Page
+
+logger = logging.getLogger("mkdocs")
+
+ROOT_DIR = Path(__file__).parent.parent.parent.parent.resolve()
+VLLM_DIR = ROOT_DIR / "vllm"
+
+# Maps short name -> qualified name (e.g. "ModelConfig" -> "vllm.config.ModelConfig")
+_name_index: dict[str, str] = {}
+
+# Fenced code block pattern (``` or ~~~, with optional language specifier).
+_FENCED_BLOCK = re.compile(
+    r"(?:^|\n)(?P<fence>`{3,}|~{3,})[^\n]*\n.*?(?:\n(?P=fence))", re.DOTALL
+)
+
+# Inline code that is NOT already part of a markdown link.
+# Matches `Name` but not [`Name`] and not [`Name`][...] or [`Name`](...).
+_INLINE_CODE = re.compile(
+    r"(?<!\[)"  # not preceded by [
+    r"`(?P<name>[A-Za-z0-9_]*)`"  # `UpperCamelCase` or `UPPER_SNAKE`
+    r"(?!\])"  # not followed by ]
+)
+
+
+def _has_docstring(node: ast.AST) -> bool:
+    """Check if a class or function node has a docstring."""
+    if not isinstance(node, ast.ClassDef | ast.FunctionDef | ast.AsyncFunctionDef):
+        return False
+    return ast.get_docstring(node, clean=False) is not None
+
+
+def _module_path(filepath: Path) -> str:
+    """Convert a filesystem path to a dotted module path."""
+    rel = filepath.relative_to(ROOT_DIR)
+    parts = list(rel.with_suffix("").parts)
+    if parts[-1] == "__init__":
+        parts = parts[:-1]
+    return ".".join(parts)
+
+
+def _index_file(filepath: Path) -> dict[str, str]:
+    """Extract documented public names from a Python file using AST parsing.
+
+    Only classes and functions with docstrings are included, since
+    mkdocstrings won't generate a page for undocumented symbols.
+    """
+    names: dict[str, str] = {}
+    try:
+        source = filepath.read_text(encoding="utf-8")
+        tree = ast.parse(source, filename=str(filepath))
+    except (SyntaxError, UnicodeDecodeError):
+        return names
+
+    module = _module_path(filepath)
+
+    for node in ast.iter_child_nodes(tree):
+        if (
+            # Class definitions (with docstring)
+            isinstance(node, ast.ClassDef)
+            and not node.name.startswith("_")
+            and _has_docstring(node)
+        ) or (
+            # Function definitions (with docstring, only uppercase/CamelCase)
+            isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef)
+            and not node.name.startswith("_")
+            and node.name[0].isupper()
+            and _has_docstring(node)
+        ):
+            names[node.name] = f"{module}.{node.name}"
+
+    return names
+
+
+def _build_index() -> dict[str, str]:
+    """Walk the vllm package and build a name -> qualified path index."""
+    index: dict[str, str] = {}
+    # Track conflicts: if multiple modules define the same name,
+    # prefer shallower modules (more likely to be the public API).
+    depth: dict[str, int] = {}
+
+    for filepath in sorted(VLLM_DIR.rglob("*.py")):
+        # Skip internal/private modules
+        if any(part.startswith("_") and part != "__init__" for part in filepath.parts):
+            continue
+        # Skip third-party vendored code
+        rel = filepath.relative_to(VLLM_DIR)
+        if rel.parts and rel.parts[0] in ("third_party", "vllm_flash_attn"):
+            continue
+
+        module_depth = len(filepath.relative_to(ROOT_DIR).parts)
+        file_names = _index_file(filepath)
+
+        for name, qualified in file_names.items():
+            if name not in index or module_depth < depth[name]:
+                index[name] = qualified
+                depth[name] = module_depth
+
+    return index
+
+
+def on_startup(*, command: str, dirty: bool) -> None:
+    """Build the name index once at startup."""
+    global _name_index
+    _name_index = _build_index()
+    logger.info("autoref_code: indexed %d names from vllm/", len(_name_index))
+
+
+def on_page_markdown(
+    markdown: str, *, page: Page, config: MkDocsConfig, files: Files
+) -> str:
+    """Replace inline code references with autoref links."""
+    if not _name_index:
+        return markdown
+
+    # Skip API reference pages to avoid circular/redundant links.
+    if page.file.src_path.startswith("api/"):
+        return markdown
+
+    # Step 1: Mask fenced code blocks so we don't touch code inside them.
+    masks: list[str] = []
+
+    def _mask_block(match: re.Match) -> str:
+        masks.append(match.group(0))
+        return f"\ue000CODEBLOCK{len(masks) - 1}\ue000"
+
+    masked = _FENCED_BLOCK.sub(_mask_block, markdown)
+
+    # Step 2: Replace inline code references.
+    def _replace(match: re.Match) -> str:
+        name = match.group("name")
+        qualified = _name_index.get(name)
+        if qualified is None:
+            return match.group(0)
+        logger.debug("autoref_code: linking `%s` to [%s]", name, qualified)
+        return f"[`{name}`][{qualified}]"
+
+    result = _INLINE_CODE.sub(_replace, masked)
+
+    # Step 3: Restore masked code blocks.
+    result = re.sub(
+        r"\ue000CODEBLOCK(\d+)\ue000", lambda m: masks[int(m.group(1))], result
+    )
+    return result

docs/models/supported_models.md

@@ -457,6 +457,7 @@ th {
 | `PanguEmbeddedForCausalLM` | openPangu-Embedded-7B | `FreedomIntelligence/openPangu-Embedded-7B-V1.1` | ✅︎ | ✅︎ |
 | `PanguProMoEV2ForCausalLM` | openpangu-pro-moe-v2 | | ✅︎ | ✅︎ |
 | `PanguUltraMoEForCausalLM` | openpangu-ultra-moe-718b-model | `FreedomIntelligence/openPangu-Ultra-MoE-718B-V1.1` | ✅︎ | ✅︎ |
+| `Param2MoEForCausalLM` | param2moe | `bharatgenai/Param2-17B-A2.4B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `PhiForCausalLM` | Phi | `microsoft/phi-1_5`, `microsoft/phi-2`, etc. | ✅︎ | ✅︎ |
 | `Phi3ForCausalLM` | Phi-4, Phi-3 | `microsoft/Phi-4-mini-instruct`, `microsoft/Phi-4`, `microsoft/Phi-3-mini-4k-instruct`, `microsoft/Phi-3-mini-128k-instruct`, `microsoft/Phi-3-medium-128k-instruct`, etc. | ✅︎ | ✅︎ |
 | `PhiMoEForCausalLM` | Phi-3.5-MoE | `microsoft/Phi-3.5-MoE-instruct`, etc. | ✅︎ | ✅︎ |

mkdocs.yaml

@@ -54,6 +54,7 @@ hooks:
   - docs/mkdocs/hooks/generate_argparse.py
   - docs/mkdocs/hooks/generate_metrics.py
   - docs/mkdocs/hooks/url_schemes.py
+  - docs/mkdocs/hooks/autoref_code.py
 
 plugins:
   - meta

requirements/common.txt

@@ -31,7 +31,7 @@ partial-json-parser # used for parsing partial JSON outputs
 pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0
-mistral_common[image] >= 1.10.0
+mistral_common[image] >= 1.11.0
 opencv-python-headless >= 4.13.0    # required for video IO
 pyyaml
 six>=1.16.0; python_version > '3.11' # transitive dependency of pandas that needs to be the latest version for python 3.12

requirements/nightly_torch_test.txt

@@ -23,7 +23,7 @@ jiwer # required for audio tests
 timm # required for internvl test
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio] >= 1.9.1 # required for voxtral test
+mistral_common[image,audio] >= 1.11.0 # required for voxtral test
 num2words # required for smolvlm test
 opencv-python-headless >= 4.13.0 # required for video test
 datamodel_code_generator # required for minicpm3 test

requirements/rocm-test.in

@@ -31,7 +31,7 @@ tblib # for pickling test exceptions
 timm>=1.0.17 # required for internvl and gemma3n-mm test
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio]>=1.10.0 # required for voxtral test
+mistral_common[image,audio]>=1.11.0 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 opencv-python-headless>=4.13.0 # required for video test

requirements/rocm-test.txt

@@ -604,7 +604,7 @@ mcp==1.27.0
     # via -r requirements/common.txt
 mdurl==0.1.2
     # via markdown-it-py
-mistral-common==1.10.0
+mistral-common==1.11.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/common.txt

requirements/test.in

@@ -32,7 +32,7 @@ torchaudio==2.10.0
 torchvision==0.25.0
 transformers_stream_generator # required for qwen-vl test
 matplotlib # required for qwen-vl test
-mistral_common[image,audio] >= 1.9.1 # required for voxtral test
+mistral_common[image,audio] >= 1.11.0 # required for voxtral test
 num2words # required for smolvlm test
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 opencv-python-headless >= 4.13.0 # required for video test

requirements/test.txt

@@ -508,7 +508,7 @@ mbstrdecoder==1.1.3
     #   typepy
 mdurl==0.1.2
     # via markdown-it-py
-mistral-common==1.10.0
+mistral-common==1.11.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test.in

tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-baseline.yaml

@@ -3,4 +3,4 @@
 model_name: openai/gpt-oss-20b
 metric_threshold: 0.568
 reasoning_effort: low
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN"
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --tensor-parallel-size 2"

tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml

@@ -3,6 +3,6 @@
 model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
 metric_threshold: 0.568
 reasoning_effort: low
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter"
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter --tokenizer openai/gpt-oss-20b --tensor-parallel-size 2"
 env:
-  VLLM_ROCM_USE_AITER: "1"
+  VLLM_ROCM_USE_AITER: "1"

tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml

@@ -3,4 +3,4 @@
 model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
 metric_threshold: 0.568
 reasoning_effort: low
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton"
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton --tokenizer openai/gpt-oss-20b --tensor-parallel-size 2"

tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml

@@ -3,6 +3,6 @@
 model_name: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8
 metric_threshold: 0.568
 reasoning_effort: low
-server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN"
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --tensor-parallel-size 2"
 env:
   VLLM_ROCM_USE_AITER: "1"

tests/kernels/core/test_layernorm.py

@@ -6,13 +6,21 @@ import torch
 
 from tests.kernels.quant_utils import FP8_DTYPE
 from tests.kernels.utils import opcheck
-from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
+from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
+if current_platform.is_rocm():
+    from vllm.platforms.rocm import on_gfx90a
+
+    on_mi250 = on_gfx90a()
+else:
+    on_mi250 = False
+
 DTYPES = [torch.half, torch.bfloat16, torch.float]
 NUM_TOKENS = [7, 83, 4096]  # Arbitrary values for testing
 HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192]  # Arbitrary values for testing
-ADD_RESIDUAL = [False, True]
+ADD_RESIDUAL = [False, True] if not on_mi250 else [True]
 SEEDS = [0]
 CUDA_DEVICES = [
     f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
@@ -154,3 +162,31 @@ def test_fused_rms_norm_quant(
         atol=1e-3,
         rtol=1e-3,
     )
+
+
+@torch.inference_mode()
+def test_gemma_rms_norm_mixed_input_weight_dtype(default_vllm_config) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+
+    device = CUDA_DEVICES[0]
+    torch.set_default_device(device)
+
+    num_tokens, hidden_size = 32, 1024
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
+    layer = GemmaRMSNorm(hidden_size, eps=1e-6).to(device=device)
+    layer.weight.data.normal_(mean=0.0, std=0.1)
+
+    # Gemma uses fp32 weight parameter while activations can be bf16.
+    assert layer.weight.dtype == torch.float32
+    out = layer(x)
+
+    x_fp32 = x.float()
+    weight_fp32 = layer.weight.data.float() + 1.0
+    variance = x_fp32.pow(2).mean(dim=-1, keepdim=True)
+    ref = (x_fp32 * torch.rsqrt(variance + layer.variance_epsilon) * weight_fp32).to(
+        x.dtype
+    )
+
+    assert out.dtype == x.dtype
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)

tests/kernels/moe/conftest.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--subtests", action="store", type=str, default=None, help="subtest ids"
+    )
+
+
+@pytest.fixture
+def subtests(request):
+    return request.config.getoption("--subtests")

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

@@ -11,7 +11,11 @@ from torch.multiprocessing import spawn  # pyright: ignore[reportPrivateImportUs
 from typing_extensions import ParamSpec
 
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed import init_distributed_environment, initialize_model_parallel
+from vllm.distributed import (
+    cleanup_dist_env_and_memory,
+    init_distributed_environment,
+    initialize_model_parallel,
+)
 from vllm.utils.network_utils import get_open_port
 
 ## Parallel Processes Utils
@@ -36,10 +40,17 @@ def _set_vllm_config(
 
     temp_file = tempfile.mkstemp()[1]
 
+    # When DP is enabled, processes are organized as:
+    #  rank = dp_rank * tp_pp_world_size + tp_pp_rank
+    tp_pp_world_size = vllm_config.parallel_config.world_size
+    vllm_config.parallel_config.data_parallel_rank = rank // tp_pp_world_size
+    tp_pp_rank = rank % tp_pp_world_size
+    vllm_config.parallel_config.rank = tp_pp_rank
+
     with set_current_vllm_config(vllm_config):
         init_distributed_environment(
-            world_size=world_size,
-            rank=rank,
+            world_size=tp_pp_world_size,
+            rank=tp_pp_rank,
             distributed_init_method=f"file://{temp_file}",
             local_rank=local_rank,
             backend="nccl",
@@ -59,11 +70,11 @@ def _worker_parallel_launch(
     world_local_size: int,
     node_rank: int,
     init_method: str,
-    worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig | None, Any, P], None],
+    worker: Callable[..., None],
     vllm_config: VllmConfig | None,
     env_dict: dict | None,
-    *args: P.args,
-    **kwargs: P.kwargs,
+    worker_kwargs: dict[str, Any],
+    *args: Any,
 ) -> None:
     rank = node_rank * world_local_size + local_rank
     torch.accelerator.set_device_index(local_rank)
@@ -98,14 +109,17 @@ def _worker_parallel_launch(
             vllm_config,
             cpu_group,
             *args,
-            **kwargs,
+            **worker_kwargs,
         )
     except Exception as ex:
         print(ex)
         traceback.print_exc()
         raise
     finally:
-        torch.distributed.destroy_process_group()
+        if vllm_config is not None:
+            cleanup_dist_env_and_memory()
+        else:
+            torch.distributed.destroy_process_group()
 
 
 def parallel_launch_with_config(
@@ -116,7 +130,6 @@ def parallel_launch_with_config(
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> None:
-    assert not kwargs
     spawn(
         _worker_parallel_launch,
         args=(
@@ -127,6 +140,7 @@ def parallel_launch_with_config(
             worker,
             vllm_config,
             env_dict,
+            kwargs,
         )
         + args,
         nprocs=world_size,

tests/kernels/moe/test_cutedsl_moe.py

@@ -17,7 +17,7 @@ from flashinfer import fp4_quantize
 from torch.nn import functional as F
 
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import (
+from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import (  # noqa: E501
     flashinfer_cutedsl_moe_masked,
 )
 from vllm.utils.flashinfer import (

tests/kernels/moe/test_gpt_oss_triton_kernels.py

@@ -23,16 +23,12 @@ from triton_kernels.numerics_details.mxfp import downcast_to_mxfp, upcast_from_m
 from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
 from triton_kernels.tensor_details import layout
 from triton_kernels.testing import assert_close
-from triton_kernels.topk import topk as topk_fn
 
 from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
-    legacy_routing,
-    make_routing_data,
     triton_kernel_moe_forward,
 )
 from vllm.utils.math_utils import round_up
-from vllm.utils.torch_utils import set_random_seed
 
 from .utils import shuffle_weight
 
@@ -97,10 +93,18 @@ def init_compute_data(M, K, N, E, a_dtype: str, w_dtype: str, num_warps: int):
     if w_dtype != "mx4":
         pytest.skip("NYI")
     else:  # quantize to mx4
-        # careful on the padding here, the activation padding need to be
-        # multiple of 64, the actual engine is not implemented
-        w1_bottom_pad = round_up(w1_tri.shape[1], 64) - w1_tri.shape[1]
-        w1_right_pad = round_up(w1_tri.shape[2], 128) - w1_tri.shape[2]
+        # Padding alignment depends on the platform.  On CDNA4 the scale
+        # swizzle requires SCALE_K % 8 == 0 (K % 256) and
+        # SCALE_N % 32 == 0 (2*N % 512), matching the production
+        # alignment in mxfp4_round_up_hidden_size_and_intermediate_size.
+        # On CUDA (Hopper) the scale layout pads internally, so the
+        # original 64/128 alignment is sufficient.
+        if current_platform.is_rocm():
+            k_align, n2_align = 256, 512
+        else:
+            k_align, n2_align = 64, 128
+        w1_bottom_pad = round_up(w1_tri.shape[1], k_align) - w1_tri.shape[1]
+        w1_right_pad = round_up(w1_tri.shape[2], n2_align) - w1_tri.shape[2]
 
         w2_bottom_pad = w1_right_pad // 2
         w2_right_pad = w1_bottom_pad
@@ -367,52 +371,3 @@ def test_unit_shuffle():
     )
 
     assert_close(ref=out_ref, tri=out)
-
-
-@pytest.mark.parametrize("num_tokens", [2, 8, 64])
-@pytest.mark.parametrize("num_experts", [32, 128])
-@pytest.mark.parametrize("topk", [1, 4])
-@pytest.mark.parametrize("renormalize", [True, False])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_legacy_routing(
-    num_tokens: int, num_experts: int, topk: int, renormalize: bool, dtype: torch.dtype
-):
-    set_random_seed(0)
-    gating_output = torch.randn(num_tokens, num_experts, device="cuda", dtype=dtype)
-
-    sm_first = not renormalize
-    logits = gating_output
-    if sm_first:
-        logits = torch.softmax(logits, dim=-1)
-    topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
-    # topk_fn returns SparseMatrix on NVIDIA, plain tuple on ROCm.
-    if isinstance(topk_result, tuple):
-        topk_weights, topk_ids_raw, bitmatrix = topk_result
-        from triton_kernels.routing import routing_from_bitmatrix
-
-        routing_data_ref, gather_indx_ref, scatter_indx_ref = routing_from_bitmatrix(
-            bitmatrix, topk_weights, topk_ids_raw, num_experts, topk
-        )
-    else:
-        topk_ids = topk_result.indx.to(torch.long)
-        topk_weights = topk_result.vals
-        routing_data_ref, gather_indx_ref, scatter_indx_ref = make_routing_data(
-            topk_ids, topk_weights, num_experts
-        )
-
-    routing_data, gather_indx, scatter_indx = legacy_routing(
-        gating_output, topk, sm_first=sm_first
-    )
-
-    assert_close(
-        ref=gather_indx_ref.src_indx, tri=gather_indx.src_indx, maxtol=0, rmstol=0
-    )
-    assert_close(
-        ref=gather_indx_ref.dst_indx, tri=gather_indx.dst_indx, maxtol=0, rmstol=0
-    )
-    assert_close(
-        ref=scatter_indx_ref.src_indx, tri=scatter_indx.src_indx, maxtol=0, rmstol=0
-    )
-    assert_close(
-        ref=scatter_indx_ref.dst_indx, tri=scatter_indx.dst_indx, maxtol=0, rmstol=0
-    )

tests/kernels/moe/utils.py

@@ -248,7 +248,7 @@ def make_quantized_test_activations(
     return a, a_q, a_scale
 
 
-def moe_quantize_weights(
+def moe_quantize_weights_2d(
     w: torch.Tensor,
     w_s: torch.Tensor | None,
     quant_dtype: torch.dtype | str | None,
@@ -293,6 +293,40 @@ def moe_quantize_weights(
     return w, w_s, w_gs
 
 
+def moe_quantize_weights(
+    w: torch.Tensor,
+    w_s: torch.Tensor | None,
+    quant_dtype: torch.dtype | str | None,
+    per_token_quant: bool,
+    block_shape: list[int] | None,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    assert w.dim() == 3
+    e, rows, cols = w.shape
+    w_l = [None] * e
+    w_s_l = [None] * e
+    w_gs_l = [None] * e
+    for idx in range(e):
+        w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights_2d(
+            w[idx], None, quant_dtype, per_token_quant, block_shape
+        )
+
+    w = torch.stack(w_l)
+    w_s = torch.stack(w_s_l)
+    w_gs = torch.stack(w_gs_l) if e > 0 and w_gs_l[0] is not None else None
+
+    if w_s.ndim == 2:
+        assert w_s.shape[-1] == 1
+        w_s = w_s.view(-1, 1, 1)
+
+    if block_shape is not None:
+        block_n, block_k = block_shape
+        n_tiles = (rows + block_n - 1) // block_n
+        k_tiles = (cols + block_k - 1) // block_k
+        assert w_s.shape == (e, n_tiles, k_tiles)
+
+    return w, w_s, w_gs
+
+
 def make_test_weight(
     e: int,
     rows: int,
@@ -303,30 +337,11 @@ def make_test_weight(
     per_out_ch_quant: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
     w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15
-    w_gs = None
 
     if quant_dtype is not None:
-        w_l = [None] * e
-        w_s_l = [None] * e
-        w_gs_l = [None] * e
-        for idx in range(e):
-            w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights(
-                w_16[idx], None, quant_dtype, per_out_ch_quant, block_shape
-            )
-
-        w = torch.stack(w_l)
-        w_s = torch.stack(w_s_l)
-        if e > 0 and w_gs_l[0] is not None:
-            w_gs = torch.stack(w_gs_l)
-        if w_s.ndim == 2:
-            assert w_s.shape[-1] == 1
-            w_s = w_s.view(-1, 1, 1)
-
-        if block_shape is not None:
-            block_n, block_k = block_shape
-            n_tiles = (rows + block_n - 1) // block_n
-            k_tiles = (cols + block_k - 1) // block_k
-            assert w_s.shape == (e, n_tiles, k_tiles)
+        w, w_s, w_gs = moe_quantize_weights(
+            w_16, None, quant_dtype, per_out_ch_quant, block_shape
+        )
     else:
         w = w_16
         w_s = None
@@ -454,7 +469,6 @@ def fused_moe(
     )
 
 
-# CustomOp?
 class BaselineMM(torch.nn.Module):
     def __init__(
         self,
@@ -462,13 +476,22 @@ class BaselineMM(torch.nn.Module):
         out_dtype: torch.dtype,
     ):
         super().__init__()
-        self.b = b.to(dtype=torch.float32)
+        self.b = torch.nn.Parameter(b.to(dtype=torch.float32))
         self.out_dtype = out_dtype
 
     def forward(self, a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
         return torch.mm(a.to(dtype=torch.float32), self.b).to(self.out_dtype), None
 
 
+class BaselineSiluAndMul(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        return torch.nn.functional.silu(x[..., :d]) * x[..., d:]
+
+
 class TestMLP(torch.nn.Module):
     def __init__(
         self,
@@ -479,7 +502,7 @@ class TestMLP(torch.nn.Module):
         super().__init__()
         self.gate_up_proj = BaselineMM(w1, out_dtype)
         self.down_proj = BaselineMM(w2, out_dtype)
-        self.act_fn = SiluAndMul()
+        self.act_fn = BaselineSiluAndMul()
 
     def forward(self, x):
         x, _ = self.gate_up_proj(x)
@@ -564,35 +587,24 @@ class RealMLP(torch.nn.Module):
         return x
 
 
-def make_shared_experts(
+def make_shared_experts_with_weights(
     N: int,
     K: int,
-    in_dtype: torch.dtype = torch.bfloat16,
+    in_dtype: torch.dtype,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_s: torch.Tensor | None = None,
+    w2_s: torch.Tensor | None = None,
     quant_dtype: torch.dtype | str | None = None,
 ) -> torch.nn.Module:
-    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
-
-    (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
-        1,
-        N,
-        K,
-        in_dtype=in_dtype,
-        quant_dtype=quant_dtype,
-    )
     old_dtype = torch.get_default_dtype()
     try:
         torch.set_default_dtype(in_dtype)
         if quant_dtype == torch.float8_e4m3fn:
-            w1 = w1[0].transpose(0, 1)
-            w2 = w2[0].transpose(0, 1)
-            w1_s = w1_s[0].transpose(0, 1) if w1_s is not None else None
-            w2_s = w2_s[0].transpose(0, 1) if w2_s is not None else None
+            from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
             quant_config = Fp8Config(True)
         else:
-            w1 = w1[0]
-            w2 = w2[0]
-            w1_s = None
-            w2_s = None
             quant_config = None
 
         return RealMLP(K, N, w1, w2, "silu", quant_config, w1_s=w1_s, w2_s=w2_s)
@@ -614,3 +626,22 @@ def modular_triton_fused_moe(
         TritonExperts(moe_config, quant_config),
         inplace=False,
     )
+
+
+def make_shared_experts(
+    N: int,
+    K: int,
+    in_dtype: torch.dtype = torch.bfloat16,
+    quant_dtype: torch.dtype | str | None = None,
+) -> torch.nn.Module:
+    (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
+        1,
+        N,
+        K,
+        in_dtype=in_dtype,
+        quant_dtype=quant_dtype,
+    )
+
+    return make_shared_experts_with_weights(
+        N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype
+    )

tests/kernels/quantization/test_mxfp4_triton_ep.py

@@ -4,12 +4,9 @@
 Tests that triton_kernel_moe_forward correctly applies expert_map
 remapping when expert parallelism (EP) is enabled.
 
-Previously, legacy_routing was always used and it produced routing data
-with global expert IDs that didn't correspond to local weight indices,
-causing illegal memory access with EP.  The fix splits routing: when
-expert_map is provided, topk selection is performed first, expert_map is
-applied to remap global→local IDs, and make_routing_data builds routing
-structures from the local IDs.
+Both EP and non-EP paths use topk + make_routing_data. When expert_map
+is provided, global expert IDs are remapped to local IDs before building
+routing structures.
 """
 
 from unittest.mock import MagicMock, patch
@@ -24,21 +21,15 @@ class TestTritonMoeForwardExpertMap:
 
     @pytest.mark.parametrize("expert_map_present", [False, True])
     def test_routing_path_selection(self, expert_map_present):
-        """Verify that the EP-aware routing path is taken when expert_map
-        is present, and the legacy_routing path is taken otherwise."""
+        """Verify that both EP and non-EP paths use topk + make_routing_data,
+        and that expert_map remapping is applied when present."""
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        # This is a structural test: we mock the routing functions to
-        # verify the correct path is exercised.
         mock_expert_map = (
             torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
         )
 
         with (
-            patch(
-                "vllm.model_executor.layers.fused_moe."
-                "gpt_oss_triton_kernels_moe.legacy_routing"
-            ) as mock_legacy,
             patch("triton_kernels.topk.topk") as mock_topk,
             patch(
                 "vllm.model_executor.layers.fused_moe."
@@ -53,27 +44,19 @@ class TestTritonMoeForwardExpertMap:
                 triton_kernel_moe_forward,
             )
 
-            # Set up return values
             mock_routing_data = MagicMock()
             mock_gather = MagicMock()
             mock_scatter = MagicMock()
 
-            if expert_map_present:
-                sparse_result = MagicMock()
-                sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32)
-                sparse_result.vals = torch.tensor([[0.6, 0.4]])
-                mock_topk.return_value = sparse_result
-                mock_make_routing.return_value = (
-                    mock_routing_data,
-                    mock_gather,
-                    mock_scatter,
-                )
-            else:
-                mock_legacy.return_value = (
-                    mock_routing_data,
-                    mock_gather,
-                    mock_scatter,
-                )
+            sparse_result = MagicMock()
+            sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32)
+            sparse_result.vals = torch.tensor([[0.6, 0.4]])
+            mock_topk.return_value = sparse_result
+            mock_make_routing.return_value = (
+                mock_routing_data,
+                mock_gather,
+                mock_scatter,
+            )
 
             mock_fused_experts.return_value = torch.zeros((1, 8), device=device)
 
@@ -92,20 +75,14 @@ class TestTritonMoeForwardExpertMap:
                 expert_map=mock_expert_map,
             )
 
+            # Both paths use topk + make_routing_data
+            mock_topk.assert_called_once()
+            mock_make_routing.assert_called_once()
+
             if expert_map_present:
-                # EP path: should use topk + make_routing_data, NOT
-                # legacy_routing
-                mock_topk.assert_called_once()
-                mock_make_routing.assert_called_once()
-                mock_legacy.assert_not_called()
                 # expert_map should be None in the fused_experts call
                 # (already applied)
                 call_kwargs = mock_fused_experts.call_args
                 assert call_kwargs[1].get("expert_map") is None or (
                     len(call_kwargs[0]) > 0
                 )
-            else:
-                # Non-EP path: should use legacy_routing
-                mock_legacy.assert_called_once()
-                mock_topk.assert_not_called()
-                mock_make_routing.assert_not_called()

tests/models/language/pooling/test_colbert.py

@@ -109,6 +109,14 @@ def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device):
         **extra,
     ).to(device)
     model.eval()
+
+    # Transformers 5.0 weight materialization can clear non-persistent
+    # buffers (e.g. rotary inv_freq) that were registered with
+    # persistent=False.  Re-compute them so the model produces valid output.
+    for mod in model.modules():
+        if hasattr(mod, "_compute_inv_freq") and hasattr(mod, "inv_freq"):
+            mod.inv_freq = mod._compute_inv_freq(device=device)
+
     return model
 
 

tests/models/quantization/test_nvfp4.py

@@ -89,22 +89,33 @@ def test_models(example_prompts, model_name) -> None:
 
 EAGER = [True, False]
 
+SM_100_NVFP4_BACKENDS = [
+    "flashinfer-cudnn",
+    "flashinfer-trtllm",
+    "flashinfer-cutlass",
+]
+
 
-@pytest.mark.skipif(
-    not current_platform.has_device_capability(100),
-    reason="modelopt_fp4 is not supported on this GPU type.",
-)
 @pytest.mark.parametrize("model", ["nvidia/Llama-3.1-8B-Instruct-NVFP4"])
 @pytest.mark.parametrize("eager", EAGER)
 @pytest.mark.parametrize(
     "backend",
     [
+        "emulation",
         "flashinfer-cudnn",
         "flashinfer-trtllm",  # the small seq_len ensures trtllm_8x4_layout backend is used
         "flashinfer-cutlass",
     ],
 )
 def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):
+    if (
+        not current_platform.has_device_capability(100)
+        and backend in SM_100_NVFP4_BACKENDS
+    ):
+        pytest.skip(
+            f"The backend {backend} is not supported with current_platform.has_device_capability(100) == False"
+        )
+
     monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend)
     with vllm_runner(model, enforce_eager=eager) as llm:
         output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)

tests/models/registry.py

@@ -461,6 +461,10 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         trust_remote_code=True,
         is_available_online=False,
     ),
+    "Param2MoEForCausalLM": _HfExamplesInfo(
+        "bharatgenai/Param2-17B-A2.4B-Thinking",
+        trust_remote_code=True,
+    ),
     "PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),
     "PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
     "Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),
@@ -1246,6 +1250,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         use_original_num_layers=True,
         max_model_len=10240,
     ),
+    "Eagle3MiniMaxM2ForCausalLM": _HfExamplesInfo(
+        "MiniMaxAI/MiniMax-M2",
+        trust_remote_code=True,
+        speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
+        tokenizer="MiniMaxAI/MiniMax-M2",
+    ),
     "EagleMistralLarge3ForCausalLM": _HfExamplesInfo(
         "mistralai/Mistral-Large-3-675B-Instruct-2512",
         speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle",

tests/quantization/test_compressed_tensors.py

@@ -366,9 +366,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
         assert output
 
 
-@pytest.mark.skipif(
-    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
-)
 @pytest.mark.parametrize(
     "args",
     [
@@ -398,7 +395,7 @@ def test_compressed_tensors_nvfp4(vllm_runner, args):
             assert qkv_proj.scheme.group_size == 16
 
         llm.apply_model(check_model)
-        output = llm.generate_greedy("Hello my name is", max_tokens=4)
+        output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
         print(output)
         assert output
 

tests/tokenizers_/test_mistral.py

@@ -3,8 +3,10 @@
 
 from typing import Any
 
+import llguidance
 import pytest
 from mistral_common.exceptions import InvalidMessageStructureException
+from mistral_common.guidance.grammar_factory import GrammarFactory
 from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy
 
 from vllm.tokenizers.mistral import (
@@ -2407,3 +2409,29 @@ class TestMistralTokenizer:
         assert actual_tokens == expected_tokens
 
         assert mistral_tokenizer.convert_ids_to_tokens([]) == []
+
+    def test_grammar_factory(self, mistral_tokenizer: MistralTokenizer) -> None:
+        # works in this case cause Mistral 7B is < v11 and SPM
+        if not mistral_tokenizer.is_tekken:
+            with pytest.raises(AttributeError):
+                mistral_tokenizer.grammar_factory  # noqa: B018
+            return
+        factory = mistral_tokenizer.grammar_factory
+        assert isinstance(factory, GrammarFactory)
+
+        # Test caching
+        factory_2 = mistral_tokenizer.grammar_factory
+        assert factory is factory_2
+
+    def test_llg_tokenizer(self, mistral_tokenizer: MistralTokenizer) -> None:
+        if not mistral_tokenizer.is_tekken:
+            with pytest.raises(ValueError):
+                mistral_tokenizer.llg_tokenizer  # noqa: B018
+            return
+
+        llg_tokenizer = mistral_tokenizer.llg_tokenizer
+        assert isinstance(llg_tokenizer, llguidance.LLTokenizer)
+
+        # Test caching
+        llg_tokenizer_2 = mistral_tokenizer.llg_tokenizer
+        assert llg_tokenizer is llg_tokenizer_2

tests/tool_parsers/test_gemma4_tool_parser.py

@@ -502,3 +502,32 @@ class TestStreamingExtraction:
         results = self._simulate_streaming(parser, mock_request, chunks)
         name = self._collect_function_name(results)
         assert name == "get_status"
+
+    def test_streaming_split_delimiter_no_invalid_json(self, parser, mock_request):
+        """Partial <|"|> delimiter chars must not leak into streamed JSON.
+
+        Reproduces the bug from https://github.com/vllm-project/vllm/issues/38946
+        where a token boundary splits the string delimiter, leaving fragments
+        like '<|' at the end of a parsed value which then corrupt the JSON.
+        """
+        chunks = [
+            "<|tool_call>",
+            "call:todowrite{",
+            'content:<|"|>Buy milk<|',
+            '"|>}',
+            "<tool_call|>",
+        ]
+
+        results = self._simulate_streaming(parser, mock_request, chunks)
+
+        args_text = self._collect_arguments(results)
+        assert args_text, "No arguments were streamed"
+
+        # Must be valid JSON — the original bug caused a JSON parse error
+        parsed_args = json.loads(args_text)
+        assert parsed_args["content"] == "Buy milk"
+
+        # Ensure no raw delimiter fragments leaked into the JSON
+        assert "<|" not in args_text, (
+            f"Partial delimiter leaked into JSON: {args_text!r}"
+        )

tests/tool_parsers/test_mistral_tool_parser.py

@@ -3,19 +3,43 @@
 
 import json
 from collections.abc import Generator
+from unittest.mock import MagicMock, patch
 
 import partial_json_parser
 import pytest
 from mistral_common.protocol.instruct.messages import AssistantMessage
 from mistral_common.protocol.instruct.request import InstructRequest
-from mistral_common.protocol.instruct.tool_calls import FunctionCall, ToolCall
+from mistral_common.protocol.instruct.tool_calls import (
+    FunctionCall,
+    ToolCall,
+)
+from mistral_common.protocol.instruct.tool_calls import (
+    NamedToolChoice as MistralNamedToolChoice,
+)
+from mistral_common.protocol.instruct.tool_calls import (
+    ToolChoice as MistralToolChoice,
+)
+from mistral_common.protocol.instruct.tool_calls import (
+    ToolChoiceEnum as MistralToolChoiceEnum,
+)
 from partial_json_parser.core.options import Allow
 
-from vllm.entrypoints.openai.engine.protocol import DeltaMessage, DeltaToolCall
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+    DeltaMessage,
+    DeltaToolCall,
+    StructuralTagResponseFormat,
+)
+from vllm.sampling_params import StructuredOutputsParams
 from vllm.tokenizers import TokenizerLike, get_tokenizer
 from vllm.tokenizers.detokenizer_utils import detokenize_incrementally
 from vllm.tokenizers.mistral import MistralTokenizer
-from vllm.tool_parsers.mistral_tool_parser import MistralToolParser
+from vllm.tool_parsers.mistral_tool_parser import (
+    _DEFAULT_JSON_SCHEMA,
+    MistralToolParser,
+)
 
 
 @pytest.fixture(scope="module")
@@ -40,6 +64,13 @@ def mistral_tool_parser(mistral_tokenizer):
     return MistralToolParser(mistral_tokenizer)
 
 
+@pytest.fixture
+def non_mistral_parser() -> MistralToolParser:
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.get_vocab.return_value = {"[TOOL_CALLS]": 1}
+    return MistralToolParser(mock_tokenizer)
+
+
 def assert_tool_calls(
     actual_tool_calls: list[ToolCall] | list[DeltaToolCall],
     expected_tool_calls: list[ToolCall],
@@ -951,3 +982,313 @@ def test_fast_detokenization_text_detection_pre_v11(
     assert len(delta_message.tool_calls) > 0
     assert delta_message.tool_calls[0].function is not None
     assert delta_message.tool_calls[0].function.name == "add"
+
+
+SAMPLE_TOOLS_DICTS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get the weather",
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string"}},
+                "required": ["city"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "add",
+            "description": "Add two numbers",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "a": {"type": "number"},
+                    "b": {"type": "number"},
+                },
+                "required": ["a", "b"],
+            },
+        },
+    },
+]
+
+
+def _make_request(**kwargs) -> ChatCompletionRequest:
+    defaults: dict = {
+        "messages": [],
+        "model": "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
+        "tools": SAMPLE_TOOLS_DICTS,
+        "tool_choice": "auto",
+    }
+    defaults.update(kwargs)
+    return ChatCompletionRequest(**defaults)
+
+
+@pytest.mark.parametrize(
+    "request_kwargs,expected_mode,expected_parallel",
+    [
+        ({"tool_choice": "auto"}, MistralToolChoiceEnum.auto, True),
+        ({"tool_choice": "none"}, MistralToolChoiceEnum.none, True),
+        ({"tool_choice": "required"}, MistralToolChoiceEnum.required, True),
+        ({"tool_choice": None, "tools": None}, MistralToolChoiceEnum.auto, True),
+        (
+            {
+                "tool_choice": {
+                    "type": "function",
+                    "function": {"name": "get_weather"},
+                }
+            },
+            MistralNamedToolChoice.model_validate(
+                {"type": "function", "function": {"name": "get_weather"}}
+            ),
+            True,
+        ),
+        (
+            {"tool_choice": "auto", "parallel_tool_calls": False},
+            MistralToolChoiceEnum.auto,
+            False,
+        ),
+        (
+            {"tool_choice": "auto", "response_format": {"type": "text"}},
+            MistralToolChoiceEnum.auto,
+            True,
+        ),
+    ],
+    ids=[
+        "auto",
+        "none",
+        "required",
+        "null_tool_choice",
+        "named_tool_choice",
+        "parallel_false",
+        "response_format_text",
+    ],
+)
+def test_adjust_request_grammar_factory(
+    mistral_tool_parser: MistralToolParser,
+    request_kwargs: dict,
+    expected_mode: MistralToolChoice,
+    expected_parallel: bool,
+) -> None:
+    request = _make_request(**request_kwargs)
+    factory = mistral_tool_parser.model_tokenizer.grammar_factory
+
+    with patch.object(
+        factory,
+        "get_lark_from_jinja",
+        wraps=factory.get_lark_from_jinja,
+    ) as mock_get_lark:
+        result = mistral_tool_parser.adjust_request(request)
+
+        mock_get_lark.assert_called_once()
+        call_kwargs = mock_get_lark.call_args
+
+        assert call_kwargs.kwargs["mode"] == expected_mode
+        assert call_kwargs.kwargs["json_schema"] is None
+        assert call_kwargs.kwargs["parallel_tool_calls"] == expected_parallel
+
+    assert result.structured_outputs is not None
+    assert isinstance(result.structured_outputs.grammar, str)
+    assert len(result.structured_outputs.grammar) > 0
+
+
+def test_adjust_request_unsupported_grammar_for_tokenizer(mistral_tokenizer) -> None:
+    with patch.object(
+        type(mistral_tokenizer),
+        "supports_grammar",
+        new_callable=lambda: property(lambda self: False),
+    ):
+        parser = MistralToolParser(mistral_tokenizer)
+        request = _make_request()
+        result = parser.adjust_request(request)
+
+        assert result.structured_outputs is None
+
+
+@pytest.mark.parametrize(
+    "tool_choice,expected_skip",
+    [("auto", False), ("none", True)],
+    ids=["auto_skip_false", "none_skip_true"],
+)
+def test_adjust_request_non_mistral_tokenizer(
+    non_mistral_parser: MistralToolParser,
+    tool_choice: str,
+    expected_skip: bool,
+) -> None:
+    request = _make_request(tool_choice=tool_choice)
+    result = non_mistral_parser.adjust_request(request)
+
+    assert result.skip_special_tokens is expected_skip
+
+
+@pytest.mark.parametrize(
+    "so_kwargs",
+    [
+        {"regex": r"\d+"},
+        {"choice": ["a", "b"]},
+        {"structural_tag": '{"key": "value"}'},
+        {"grammar": "start: 'hello'"},
+    ],
+    ids=["regex", "choice", "structural_tag", "grammar"],
+)
+def test_adjust_request_unsupported_structured_outputs(
+    mistral_tool_parser: MistralToolParser,
+    so_kwargs: dict,
+) -> None:
+    request = _make_request(
+        structured_outputs=StructuredOutputsParams(**so_kwargs),
+    )
+    result = mistral_tool_parser.adjust_request(request)
+
+    assert result.structured_outputs == request.structured_outputs
+
+
+def test_adjust_request_unsupported_response_format(
+    mistral_tool_parser: MistralToolParser,
+) -> None:
+    request = _make_request(
+        response_format=StructuralTagResponseFormat(
+            type="structural_tag", format={"some": "config"}
+        ),
+    )
+    result = mistral_tool_parser.adjust_request(request)
+    assert result.structured_outputs is None
+    assert result.response_format == request.response_format
+
+
+@pytest.mark.parametrize(
+    "so_kwargs,expected_json_schema",
+    [
+        ({"json_object": True}, _DEFAULT_JSON_SCHEMA),
+        ({"json": '{"type": "object"}'}, {"type": "object"}),
+        (
+            {"json": {"type": "object", "properties": {"x": {"type": "integer"}}}},
+            {"type": "object", "properties": {"x": {"type": "integer"}}},
+        ),
+    ],
+    ids=["json_object", "json_str", "json_dict"],
+)
+def test_adjust_request_structured_outputs_generates_grammar(
+    mistral_tool_parser: MistralToolParser,
+    so_kwargs: dict,
+    expected_json_schema: str,
+) -> None:
+    request = _make_request(
+        structured_outputs=StructuredOutputsParams(**so_kwargs),
+    )
+    factory = mistral_tool_parser.model_tokenizer.grammar_factory
+
+    with patch.object(
+        factory,
+        "get_lark_from_jinja",
+        wraps=factory.get_lark_from_jinja,
+    ) as mock_get_lark:
+        result = mistral_tool_parser.adjust_request(request)
+
+        mock_get_lark.assert_called_once()
+        assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema
+
+    assert result.structured_outputs is not None
+    assert isinstance(result.structured_outputs.grammar, str)
+    assert len(result.structured_outputs.grammar) > 0
+
+
+@pytest.mark.parametrize(
+    "response_format_kwargs,expected_json_schema",
+    [
+        ({"type": "json_object"}, _DEFAULT_JSON_SCHEMA),
+        (
+            {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "my_schema",
+                    "schema": {
+                        "type": "object",
+                        "properties": {"x": {"type": "integer"}},
+                    },
+                },
+            },
+            {"type": "object", "properties": {"x": {"type": "integer"}}},
+        ),
+    ],
+    ids=["json_object", "json_schema_with_schema"],
+)
+def test_adjust_request_response_format_generates_grammar(
+    mistral_tool_parser: MistralToolParser,
+    response_format_kwargs: dict,
+    expected_json_schema: str,
+) -> None:
+    request = _make_request(response_format=response_format_kwargs)
+    factory = mistral_tool_parser.model_tokenizer.grammar_factory
+
+    with patch.object(
+        factory,
+        "get_lark_from_jinja",
+        wraps=factory.get_lark_from_jinja,
+    ) as mock_get_lark:
+        result = mistral_tool_parser.adjust_request(request)
+
+        mock_get_lark.assert_called_once()
+        assert mock_get_lark.call_args.kwargs["json_schema"] == expected_json_schema
+
+    assert result.structured_outputs is not None
+    assert isinstance(result.structured_outputs.grammar, str)
+    assert len(result.structured_outputs.grammar) > 0
+
+
+def test_adjust_request_tool_choice_none_with_json_schema_uses_json_schema_factory(
+    mistral_tool_parser: MistralToolParser,
+) -> None:
+    request = _make_request(
+        tool_choice="none",
+        structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
+    )
+    factory = mistral_tool_parser.model_tokenizer.grammar_factory
+
+    with patch.object(
+        factory,
+        "get_lark_for_json_schema",
+        wraps=factory.get_lark_for_json_schema,
+    ) as mock_json_schema:
+        result = mistral_tool_parser.adjust_request(request)
+
+        mock_json_schema.assert_called_once()
+        assert mock_json_schema.call_args.kwargs["json_schema"] == {"type": "object"}
+
+    assert result.structured_outputs is not None
+    assert isinstance(result.structured_outputs.grammar, str)
+    assert len(result.structured_outputs.grammar) > 0
+
+
+def test_adjust_request_tool_choice_auto_with_json_schema_uses_jinja_factory(
+    mistral_tool_parser: MistralToolParser,
+) -> None:
+    request = _make_request(
+        tool_choice="auto",
+        structured_outputs=StructuredOutputsParams(json='{"type": "object"}'),
+    )
+    factory = mistral_tool_parser.model_tokenizer.grammar_factory
+
+    with (
+        patch.object(
+            factory,
+            "get_lark_for_json_schema",
+            wraps=factory.get_lark_for_json_schema,
+        ) as mock_json_schema,
+        patch.object(
+            factory,
+            "get_lark_from_jinja",
+            wraps=factory.get_lark_from_jinja,
+        ) as mock_jinja,
+    ):
+        result = mistral_tool_parser.adjust_request(request)
+
+        mock_jinja.assert_called_once()
+        assert mock_jinja.call_args.kwargs["json_schema"] == {"type": "object"}
+        mock_json_schema.assert_not_called()
+
+    assert result.structured_outputs is not None
+    assert isinstance(result.structured_outputs.grammar, str)
+    assert len(result.structured_outputs.grammar) > 0

tests/v1/core/test_async_scheduler.py

@@ -1,10 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections import deque
+from unittest.mock import Mock
 
 import pytest
 
-from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.async_scheduler import AsyncScheduler
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import RequestStatus
 from vllm.v1.utils import ConstantList
@@ -247,3 +249,66 @@ def test_prefix_caching_for_multi_turn():
     # requests.
     for req in next_turn_requests:
         assert req.num_cached_tokens == req.num_prompt_tokens // BLOCK_SIZE * BLOCK_SIZE
+
+
+def test_abort_request_when_structured_output_fsm_cannot_advance():
+    scheduler = object.__new__(AsyncScheduler)
+    request = create_requests(num_requests=1, num_tokens=1)[0]
+    request.structured_output_request = Mock()
+    request.structured_output_request.grammar = Mock()
+    request.structured_output_request.grammar.accept_tokens.return_value = False
+    request.status = RequestStatus.RUNNING
+    request.num_computed_tokens = request.num_tokens
+    request.num_output_placeholders = 1
+
+    scheduler.perf_metrics = None
+    scheduler.connector = None
+    scheduler.structured_output_manager = Mock()
+    scheduler.structured_output_manager.should_advance.return_value = True
+    scheduler.requests = {request.request_id: request}
+    scheduler.running = [request]
+    scheduler.waiting = Mock()
+    scheduler.kv_cache_manager = Mock()
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = Mock()
+    scheduler.finished_req_ids = set()
+    scheduler.finished_req_ids_dict = None
+    scheduler.vllm_config = Mock()
+    scheduler.vllm_config.model_config.enable_return_routed_experts = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.make_stats = Mock(return_value=None)
+    scheduler.max_model_len = 128
+
+    def free_request(req, delay_free_blocks=False):
+        scheduler.finished_req_ids.add(req.request_id)
+        scheduler.requests.pop(req.request_id, None)
+        return None
+
+    scheduler._free_request = Mock(side_effect=free_request)
+
+    output = SchedulerOutput(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        num_scheduled_tokens={request.request_id: 1},
+        total_num_scheduled_tokens=1,
+        scheduled_encoder_inputs={},
+        scheduled_spec_decode_tokens={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+    model_runner_output = ModelRunnerOutput(
+        req_ids=[request.request_id],
+        req_id_to_index={request.request_id: 0},
+        sampled_token_ids=[[123]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+
+    scheduler.update_from_output(output, model_runner_output)
+
+    assert request.resumable is False
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert request.request_id not in scheduler.requests
+    assert not scheduler.running

tests/v1/core/test_scheduler.py

@@ -26,6 +26,7 @@ from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.engine import FinishReason
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
@@ -2463,6 +2464,86 @@ def test_schedule_skip_tokenizer_init_structured_output_request():
     assert len(scheduler.skipped_waiting) == 1
 
 
+def test_abort_request_when_structured_output_fsm_cannot_advance():
+    scheduler = object.__new__(Scheduler)
+    sampling_params = SamplingParams(ignore_eos=True, max_tokens=4)
+    sampling_params.update_from_generation_config({}, EOS_TOKEN_ID)
+
+    request = Request(
+        request_id="0",
+        prompt_token_ids=[0, 1],
+        mm_features=None,
+        sampling_params=sampling_params,
+        pooling_params=None,
+    )
+    request.structured_output_request = Mock()
+    request.structured_output_request.grammar = Mock()
+    request.structured_output_request.grammar.accept_tokens.return_value = False
+    request.status = RequestStatus.RUNNING
+    request.num_computed_tokens = request.num_tokens
+
+    scheduler.perf_metrics = None
+    scheduler.connector = None
+    scheduler.structured_output_manager = Mock()
+    scheduler.structured_output_manager.should_advance.return_value = True
+    scheduler.requests = {request.request_id: request}
+    scheduler.running = [request]
+    scheduler.waiting = Mock()
+    scheduler.kv_cache_manager = Mock()
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = Mock()
+    scheduler.finished_req_ids = set()
+    scheduler.finished_req_ids_dict = None
+    scheduler.vllm_config = Mock()
+    scheduler.vllm_config.model_config.enable_return_routed_experts = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.make_stats = Mock(return_value=None)
+    scheduler.max_model_len = 128
+
+    def free_request(req: Request, delay_free_blocks: bool = False):
+        scheduler.finished_req_ids.add(req.request_id)
+        scheduler.requests.pop(req.request_id, None)
+        return None
+
+    scheduler._free_request = Mock(side_effect=free_request)
+
+    output = SchedulerOutput(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=CachedRequestData.make_empty(),
+        num_scheduled_tokens={request.request_id: 1},
+        total_num_scheduled_tokens=1,
+        scheduled_encoder_inputs={},
+        scheduled_spec_decode_tokens={},
+        num_common_prefix_blocks=[],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+
+    model_runner_output = ModelRunnerOutput(
+        req_ids=[request.request_id],
+        req_id_to_index={request.request_id: 0},
+        sampled_token_ids=[[123]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[],
+    )
+    engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
+
+    request.structured_output_request.grammar.accept_tokens.assert_called_once_with(
+        request.request_id, [123]
+    )
+    assert request.resumable is False
+    assert request.status == RequestStatus.FINISHED_ERROR
+    assert request.request_id not in scheduler.requests
+    assert not scheduler.running
+    scheduler._free_request.assert_called_once_with(request)
+    assert len(engine_core_outputs[0].outputs) == 1
+    engine_core_output = engine_core_outputs[0].outputs[0]
+    assert engine_core_output.request_id == request.request_id
+    assert engine_core_output.new_token_ids == [123]
+    assert engine_core_output.finish_reason == FinishReason.ERROR
+
+
 @pytest.mark.parametrize(
     "use_ec_connector, ec_role", [(False, None), (True, "ec_consumer")]
 )

tests/v1/cudagraph/test_encoder_cudagraph.py

@@ -14,17 +14,17 @@ from typing import Any
 
 import pytest
 import torch
-from vllm.v1.worker.gpu.mm.encoder_cudagraph import (
+
+from vllm.platforms import current_platform
+from vllm.v1.worker.encoder_cudagraph import (
     EncoderCudaGraphManager,
 )
-from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
+from vllm.v1.worker.encoder_cudagraph_defs import (
     EncoderCudaGraphCaptureInputs,
     EncoderCudaGraphConfig,
     EncoderCudaGraphReplayBuffers,
 )
 
-from vllm.platforms import current_platform
-
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------

tests/v1/ec_connector/integration/run_epd_correctness_test.sh

@@ -15,8 +15,9 @@
 
 # set -xe
 
-# Find the git repository root directory
-GIT_ROOT=$(git rev-parse --show-toplevel)
+# Resolve the repository root from the script location instead of `.git`.
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 
 # Model to test
 MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -19,9 +19,9 @@ dp_ep_configs=(
 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
 )
 hybrid_ssm_configs=(
-  "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
+  "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
   # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
-  "ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
+  "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
 )
 sw_attn_configs=(
   "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -85,8 +85,11 @@ DECODE_BLOCK_SIZE=${DECODE_BLOCK_SIZE:-128}
 # Comma-separated extra args for vllm serve (e.g. --max-model-len,2048)
 VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
 
-# Find the git repository root directory
-GIT_ROOT=$(git rev-parse --show-toplevel)
+# Resolve the repository root from the script location instead of `.git`.
+# The ROCm CI image copies `/vllm-workspace` without the Git metadata, so
+# `git rev-parse --show-toplevel` is not reliable at runtime.
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 
 SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
 

tests/v1/kv_connector/nixl_integration/run_edge_case_test.sh

@@ -33,8 +33,9 @@ MODELS=(
     "Qwen/Qwen3-0.6B"
 )
 
-# Find the git repository root directory
-GIT_ROOT=$(git rev-parse --show-toplevel)
+# Resolve the repository root from the script location instead of `.git`.
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 
 # Trap the SIGINT signal (triggered by Ctrl+C)
 trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT

tests/v1/kv_connector/nixl_integration/run_tpu_disagg_accuracy_test.sh

@@ -20,7 +20,8 @@ BLOCK_SIZE=${BLOCK_SIZE:-32}
 
 
 # execution env
-GIT_ROOT=$(git rev-parse --show-toplevel)
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
 CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"}
 CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"}
@@ -153,4 +154,4 @@ echo "-----P/D success----"
 rm "${OUTPUT_FILE}"
 cleanup
 
-exit 0
+exit 0

tests/v1/kv_connector/nixl_integration/run_tpu_edge_case_test.sh

@@ -20,7 +20,8 @@ BLOCK_SIZE=${BLOCK_SIZE:-32}
 
 
 # execution env
-GIT_ROOT=$(git rev-parse --show-toplevel)
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
 CONDA_PATH=${CONDA_PATH:-"/home/${USER}/anaconda3"}
 CONDA_ENV_NAME=${CONDA_ENV_NAME:-"nixl"}

tests/v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh

@@ -44,7 +44,8 @@ DECODER_ZE_AFFINITY_MASK=${DECODER_ZE_AFFINITY_MASK:-$(generate_affinity_mask "$
 
 
 # execution env
-GIT_ROOT=$(git rev-parse --show-toplevel)
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 EXP_ROOT="${GIT_ROOT}/tests/v1/kv_connector/nixl_integration"
 
 OUTPUT_FILE=${OUTPUT_FILE:-"${EXP_ROOT}/.xpu_accuracy_test_outputs.txt"}

tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

@@ -52,7 +52,9 @@ DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
 GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7}
 BLOCK_SIZE=${BLOCK_SIZE:-16}
 
-GIT_ROOT=$(git rev-parse --show-toplevel)
+# Resolve the repository root from the script location instead of `.git`.
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
 
 SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
 

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -224,6 +224,8 @@ def test_get_block_descs_ids_hybrid_ssm():
     worker._has_mamba = True
     worker._is_mamba_group = [False, True]
     worker._physical_blocks_per_logical_kv_block = 1
+    worker._mamba_phys_ratio = {engine_id: 1}
+    worker.block_len_per_layer = [100]
     # num_descs = num_regions * num_blocks (no blocks_first doubling)
     worker.num_descs = 2 * num_blocks
 
@@ -234,9 +236,10 @@ def test_get_block_descs_ids_hybrid_ssm():
     # FA group: stride=num_blocks=100, offset=0
     #   region0: [3, 5],  region1: [103, 105]
     # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),
-    #   offset=num_descs=200
-    #   region0: [201, 202],  region1: [301, 302]
-    expected = [3, 5, 103, 105, 201, 202, 301, 302]
+    #   offset=num_fa_descs=200, 4 regions per Mamba layer (x, B, C, ssm)
+    #   region0: [201, 202], region1: [301, 302],
+    #   region2: [401, 402], region3: [501, 502]
+    expected = [3, 5, 103, 105, 201, 202, 301, 302, 401, 402, 501, 502]
     assert list(result) == expected, f"Expected {expected}, got {list(result)}"
 
 
@@ -259,6 +262,8 @@ def test_get_block_descs_ids_kernel_block_mismatch():
     worker._has_mamba = True
     worker._is_mamba_group = [False, True]
     worker._physical_blocks_per_logical_kv_block = ratio
+    worker._mamba_phys_ratio = {engine_id: ratio}
+    worker.block_len_per_layer = [100]
     worker.num_descs = 2 * num_blocks  # 800
 
     fa_blocks = [3, 7]  # kernel-level block IDs
@@ -267,9 +272,11 @@ def test_get_block_descs_ids_kernel_block_mismatch():
 
     # FA group: stride=num_blocks=400, offset=0
     #   region0: [3, 7],  region1: [403, 407]
-    # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800
-    #   region0: [801, 802],  region1: [901, 902]
-    expected = [3, 7, 403, 407, 801, 802, 901, 902]
+    # SSM group: stride=logical_blocks=400//4=100, offset=num_fa_descs=800,
+    #   4 regions per Mamba layer (x, B, C, ssm)
+    #   region0: [801, 802], region1: [901, 902],
+    #   region2: [1001, 1002], region3: [1101, 1102]
+    expected = [3, 7, 403, 407, 801, 802, 901, 902, 1001, 1002, 1101, 1102]
     assert list(result) == expected, f"Expected {expected}, got {list(result)}"
 
 
@@ -418,3 +425,29 @@ def test_has_mamba_init(
     )
     assert scheduler._has_mamba is expected_has_mamba
     assert scheduler._is_hma_required is expected_is_hma
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "ssm_sizes,block_len,expected_ratio",
+    [
+        # Nemotron 30B TP=1: ceil((36864 + 2097152) / 8192) = 261
+        ((36864, 2097152), 8192, 261),
+        # Nemotron 30B TP=2: ceil((18432 + 1048576) / 4096) = 261
+        ((18432, 1048576), 4096, 261),
+        # Nemotron 30B TP=4: ceil((9216 + 524288) / 4096) = 131
+        ((9216, 524288), 4096, 131),
+    ],
+)
+def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio):
+    """Verify that compute_mamba_phys_ratio is TP-dependent.
+
+    With dimension-sharded Mamba state, the ratio differs across TP sizes
+    (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why
+    _mamba_phys_ratio must be stored per-engine.
+    """
+    from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
+        compute_mamba_phys_ratio,
+    )
+
+    assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio

tests/v1/kv_offload/test_cpu_manager.py

@@ -15,6 +15,7 @@ from vllm.v1.kv_offload.abstract import (
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
 from vllm.v1.kv_offload.mediums import CPULoadStoreSpec
+from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
 
 
 @dataclass
@@ -243,335 +244,300 @@ def test_cpu_manager():
     )
 
 
-def test_arc_manager_basic():
-    """
-    Tests CPUOffloadingManager with arc policy.
-    Verifies that ARC handles store, load, and lookup operations correctly.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+class TestARCPolicy:
+    """Unit tests for CPUOffloadingManager with ARC eviction policy."""
 
-    # prepare store [1, 2]
-    prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2]))
-    verify_store_output(
-        prepare_store_output,
-        ExpectedPrepareStoreOutput(
-            block_hashes_to_store=[1, 2],
-            store_block_ids=[0, 1],
-            block_hashes_evicted=[],
-        ),
-    )
+    def _make_manager(
+        self, num_blocks: int = 4, enable_events: bool = True
+    ) -> tuple[CPUOffloadingManager, ARCCachePolicy]:
+        manager = CPUOffloadingManager(
+            block_size=256,
+            num_blocks=num_blocks,
+            cache_policy="arc",
+            enable_events=enable_events,
+        )
+        policy = manager._policy
+        assert isinstance(policy, ARCCachePolicy)
+        return manager, policy
 
-    # lookup [1, 2] -> not ready
-    assert arc_manager.lookup(to_hashes([1, 2])) == 0
+    def test_basic(self):
+        """
+        Tests CPUOffloadingManager with arc policy.
+        Verifies that ARC handles store, load, and lookup operations correctly.
+        """
+        cpu_manager, arc_policy = self._make_manager()
 
-    # no events so far
-    assert list(arc_manager.take_events()) == []
+        # prepare store [1, 2]
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2]))
+        verify_store_output(
+            prepare_store_output,
+            ExpectedPrepareStoreOutput(
+                block_hashes_to_store=[1, 2],
+                store_block_ids=[0, 1],
+                block_hashes_evicted=[],
+            ),
+        )
 
-    # complete store [1, 2]
-    arc_manager.complete_store(to_hashes([1, 2]))
-    verify_events(
-        arc_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
-    )
+        # lookup [1, 2] -> not ready
+        assert cpu_manager.lookup(to_hashes([1, 2])) == 0
 
-    # lookup [1, 2]
-    assert arc_manager.lookup(to_hashes([1])) == 1
-    assert arc_manager.lookup(to_hashes([1, 2])) == 2
-    assert arc_manager.lookup(to_hashes([1, 2, 3])) == 2
+        # no events so far
+        assert list(cpu_manager.take_events()) == []
 
-    # blocks should be in T1 (recent)
-    assert len(arc_policy.t1) == 2
-    assert len(arc_policy.t2) == 0
+        # complete store [1, 2]
+        cpu_manager.complete_store(to_hashes([1, 2]))
+        verify_events(
+            cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},)
+        )
 
+        # lookup [1, 2]
+        assert cpu_manager.lookup(to_hashes([1])) == 1
+        assert cpu_manager.lookup(to_hashes([1, 2])) == 2
+        assert cpu_manager.lookup(to_hashes([1, 2, 3])) == 2
 
-def test_arc_manager_t1_to_t2_promotion():
-    """
-    Tests that accessing a block in T1 promotes it to T2 (frequent).
-    This is a key feature of ARC's adaptive behavior.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # blocks should be in T1 (recent)
+        assert len(arc_policy.t1) == 2
+        assert len(arc_policy.t2) == 0
 
-    # store and complete block 1
-    arc_manager.prepare_store(to_hashes([1]))
-    arc_manager.complete_store(to_hashes([1]))
+    def test_t1_to_t2_promotion(self):
+        """
+        Tests that accessing a block in T1 promotes it to T2 (frequent).
+        This is a key feature of ARC's adaptive behavior.
+        """
+        cpu_manager, arc_policy = self._make_manager(enable_events=False)
 
-    # block 1 starts in T1 (recent)
-    assert to_hashes([1])[0] in arc_policy.t1
-    assert to_hashes([1])[0] not in arc_policy.t2
+        # store and complete block 1
+        cpu_manager.prepare_store(to_hashes([1]))
+        cpu_manager.complete_store(to_hashes([1]))
 
-    # touch block 1 (simulate second access)
-    arc_manager.touch(to_hashes([1]))
+        # block 1 starts in T1 (recent)
+        assert to_hashes([1])[0] in arc_policy.t1
+        assert to_hashes([1])[0] not in arc_policy.t2
 
-    # block 1 should now be in T2 (frequent)
-    assert to_hashes([1])[0] not in arc_policy.t1
-    assert to_hashes([1])[0] in arc_policy.t2
+        # touch block 1 (simulate second access)
+        cpu_manager.touch(to_hashes([1]))
 
+        # block 1 should now be in T2 (frequent)
+        assert to_hashes([1])[0] not in arc_policy.t1
+        assert to_hashes([1])[0] in arc_policy.t2
 
-def test_arc_manager_eviction_with_load():
-    """
-    Tests ARC eviction behavior similar to LRU test.
-    Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
-    )
+    def test_eviction_with_load(self):
+        """
+        Tests ARC eviction behavior similar to LRU test.
+        Verifies that blocks being loaded (ref_cnt > 0) cannot be evicted.
+        """
+        cpu_manager, _ = self._make_manager()
 
-    # prepare and complete store [1, 2, 3, 4]
-    prepare_store_output = arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
-    verify_store_output(
-        prepare_store_output,
-        ExpectedPrepareStoreOutput(
-            block_hashes_to_store=[1, 2, 3, 4],
-            store_block_ids=[0, 1, 2, 3],
-            block_hashes_evicted=[],
-        ),
-    )
-    arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
+        # prepare and complete store [1, 2, 3, 4]
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
+        verify_store_output(
+            prepare_store_output,
+            ExpectedPrepareStoreOutput(
+                block_hashes_to_store=[1, 2, 3, 4],
+                store_block_ids=[0, 1, 2, 3],
+                block_hashes_evicted=[],
+            ),
+        )
+        cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
 
-    # prepare load [2, 3] (increases ref_cnt)
-    prepare_load_output = arc_manager.prepare_load(to_hashes([2, 3]))
-    verify_load_output(prepare_load_output, [1, 2])
+        # prepare load [2, 3] (increases ref_cnt)
+        prepare_load_output = cpu_manager.prepare_load(to_hashes([2, 3]))
+        verify_load_output(prepare_load_output, [1, 2])
 
-    # prepare store [5, 6, 7] with [2, 3] being loaded
-    # should fail because [2, 3] have ref_cnt > 0
-    assert arc_manager.prepare_store(to_hashes([5, 6, 7])) is None
+        # prepare store [5, 6, 7] with [2, 3] being loaded
+        # should fail because [2, 3] have ref_cnt > 0
+        assert cpu_manager.prepare_store(to_hashes([5, 6, 7])) is None
 
-    # complete load [2, 3]
-    arc_manager.complete_load(to_hashes([2, 3]))
+        # complete load [2, 3]
+        cpu_manager.complete_load(to_hashes([2, 3]))
 
-    # now prepare store [5, 6, 7] should succeed
-    # ARC will evict blocks one at a time from T1 as needed
-    prepare_store_output = arc_manager.prepare_store(to_hashes([5, 6, 7]))
-    assert prepare_store_output is not None
-    # Should successfully evict enough blocks to make room (at least 1)
-    assert len(prepare_store_output.block_hashes_evicted) >= 1
+        # now prepare store [5, 6, 7] should succeed
+        # ARC will evict blocks one at a time from T1 as needed
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([5, 6, 7]))
+        assert prepare_store_output is not None
+        # Should successfully evict enough blocks to make room (at least 1)
+        assert len(prepare_store_output.block_hashes_evicted) >= 1
 
+    def test_adaptive_target(self):
+        """
+        Tests ARC's adaptive target adjustment via ghost lists.
+        When a block in B1 (ghost list) is accessed, target_t1_size increases.
+        When a block in B2 is accessed, target_t1_size decreases.
+        """
+        cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
 
-def test_arc_manager_adaptive_target():
-    """
-    Tests ARC's adaptive target adjustment via ghost lists.
-    When a block in B1 (ghost list) is accessed, target_t1_size increases.
-    When a block in B2 is accessed, target_t1_size decreases.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # store blocks 1, 2 (fills cache)
+        cpu_manager.prepare_store(to_hashes([1, 2]))
+        cpu_manager.complete_store(to_hashes([1, 2]))
 
-    # store blocks 1, 2 (fills cache)
-    arc_manager.prepare_store(to_hashes([1, 2]))
-    arc_manager.complete_store(to_hashes([1, 2]))
+        initial_target = arc_policy.target_t1_size
 
-    initial_target = arc_policy.target_t1_size
+        # store block 3, evicting block 1 (moves to B1 ghost list)
+        cpu_manager.prepare_store(to_hashes([3]))
+        cpu_manager.complete_store(to_hashes([3]))
 
-    # store block 3, evicting block 1 (moves to B1 ghost list)
-    arc_manager.prepare_store(to_hashes([3]))
-    arc_manager.complete_store(to_hashes([3]))
+        # block 1 should be in B1 (ghost list)
+        assert to_hashes([1])[0] in arc_policy.b1
 
-    # block 1 should be in B1 (ghost list)
-    assert to_hashes([1])[0] in arc_policy.b1
+        # touch block 1 (cache miss, but in B1)
+        # this should increase target_t1_size (favor recency)
+        cpu_manager.touch(to_hashes([1]))
 
-    # touch block 1 (cache miss, but in B1)
-    # this should increase target_t1_size (favor recency)
-    arc_manager.touch(to_hashes([1]))
+        # target should have increased
+        assert arc_policy.target_t1_size > initial_target
 
-    # target should have increased
-    assert arc_policy.target_t1_size > initial_target
+    def test_t1_t2_eviction_policy(self):
+        """
+        Tests that ARC evicts from T1 or T2 based on target_t1_size.
+        If |T1| >= target_t1_size, evict from T1, otherwise from T2.
+        """
+        cpu_manager, arc_policy = self._make_manager(enable_events=False)
 
+        # store blocks 1, 2, 3, 4
+        cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
+        cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
 
-def test_arc_manager_t1_t2_eviction_policy():
-    """
-    Tests that ARC evicts from T1 or T2 based on target_t1_size.
-    If |T1| >= target_t1_size, evict from T1, otherwise from T2.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=False
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # promote blocks 3, 4 to T2 by touching them
+        cpu_manager.touch(to_hashes([3, 4]))
 
-    # store blocks 1, 2, 3, 4
-    arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
-    arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
+        # now: T1 = {1, 2}, T2 = {3, 4}
+        assert len(arc_policy.t1) == 2
+        assert len(arc_policy.t2) == 2
 
-    # promote blocks 3, 4 to T2 by touching them
-    arc_manager.touch(to_hashes([3, 4]))
+        # set target_t1_size to prefer evicting from T1
+        # (when |T1| >= target, evict from T1)
+        arc_policy.target_t1_size = 1
 
-    # now: T1 = {1, 2}, T2 = {3, 4}
-    assert len(arc_policy.t1) == 2
-    assert len(arc_policy.t2) == 2
+        # store block 5, should evict from T1 (block 1, LRU in T1)
+        output = cpu_manager.prepare_store(to_hashes([5]))
+        assert output is not None
+        assert to_hashes([1]) == output.block_hashes_evicted
 
-    # set target_t1_size to prefer evicting from T1
-    # (when |T1| >= target, evict from T1)
-    arc_policy.target_t1_size = 1
+        cpu_manager.complete_store(to_hashes([5]))
 
-    # store block 5, should evict from T1 (block 1, LRU in T1)
-    output = arc_manager.prepare_store(to_hashes([5]))
-    assert output is not None
-    assert to_hashes([1]) == output.block_hashes_evicted
+        # block 1 should be in B1 (ghost list)
+        assert to_hashes([1])[0] in arc_policy.b1
+        # block 5 should be in T1
+        assert to_hashes([5])[0] in arc_policy.t1
 
-    arc_manager.complete_store(to_hashes([5]))
+    def test_ghost_list_bounds(self):
+        """
+        Tests that ghost lists (B1, B2) don't grow unbounded.
+        They should be capped at cache_capacity.
+        """
+        cpu_manager, arc_policy = self._make_manager(num_blocks=2, enable_events=False)
 
-    # block 1 should be in B1 (ghost list)
-    assert to_hashes([1])[0] in arc_policy.b1
-    # block 5 should be in T1
-    assert to_hashes([5])[0] in arc_policy.t1
+        # fill cache with blocks 1, 2
+        cpu_manager.prepare_store(to_hashes([1, 2]))
+        cpu_manager.complete_store(to_hashes([1, 2]))
 
+        # store many blocks to fill ghost lists
+        for i in range(3, 20):
+            cpu_manager.prepare_store(to_hashes([i]))
+            cpu_manager.complete_store(to_hashes([i]))
 
-def test_arc_manager_ghost_list_bounds():
-    """
-    Tests that ghost lists (B1, B2) don't grow unbounded.
-    They should be capped at cache_capacity.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=2, cache_policy="arc", enable_events=False
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # ghost lists should not exceed cache_capacity
+        assert len(arc_policy.b1) <= arc_policy.cache_capacity
+        assert len(arc_policy.b2) <= arc_policy.cache_capacity
 
-    # fill cache with blocks 1, 2
-    arc_manager.prepare_store(to_hashes([1, 2]))
-    arc_manager.complete_store(to_hashes([1, 2]))
+    def test_touch_ordering(self):
+        """
+        Tests that touch() correctly updates access patterns.
+        Similar to LRU test but verifies T1/T2 ordering.
+        """
+        cpu_manager, arc_policy = self._make_manager()
 
-    # store many blocks to fill ghost lists
-    for i in range(3, 20):
-        arc_manager.prepare_store(to_hashes([i]))
-        arc_manager.complete_store(to_hashes([i]))
+        # store blocks 1, 2, 3, 4
+        cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
+        cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
 
-    # ghost lists should not exceed cache_capacity
-    assert len(arc_policy.b1) <= arc_policy.cache_capacity
-    assert len(arc_policy.b2) <= arc_policy.cache_capacity
+        # promote 3, 4 to T2
+        cpu_manager.touch(to_hashes([3, 4]))
 
+        # T1 = {1, 2}, T2 = {3, 4}
+        # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
+        cpu_manager.touch(to_hashes([1, 3, 4]))
 
-def test_arc_manager_touch_ordering():
-    """
-    Tests that touch() correctly updates access patterns.
-    Similar to LRU test but verifies T1/T2 ordering.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
+        assert len(arc_policy.t1) == 1
+        assert len(arc_policy.t2) == 3
 
-    # store blocks 1, 2, 3, 4
-    arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
-    arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
+        # store block 5, should evict from T1 (block 2, only one in T1)
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([5]))
+        verify_store_output(
+            prepare_store_output,
+            ExpectedPrepareStoreOutput(
+                block_hashes_to_store=[5],
+                store_block_ids=[1],  # reuses block 2's storage
+                block_hashes_evicted=[2],
+            ),
+        )
 
-    # promote 3, 4 to T2
-    arc_manager.touch(to_hashes([3, 4]))
+    def test_failed_store(self):
+        """
+        Tests that failed store operations clean up correctly.
+        Similar to LRU test but for ARC.
+        """
+        cpu_manager, arc_policy = self._make_manager()
 
-    # T1 = {1, 2}, T2 = {3, 4}
-    # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2
-    arc_manager.touch(to_hashes([1, 3, 4]))
+        # store blocks 1, 2, 3, 4
+        cpu_manager.prepare_store(to_hashes([1, 2, 3, 4]))
+        cpu_manager.complete_store(to_hashes([1, 2, 3, 4]))
 
-    # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent)
-    assert len(arc_policy.t1) == 1
-    assert len(arc_policy.t2) == 3
+        # prepare store block 5 (will evict block 1)
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([5]))
+        assert prepare_store_output is not None
+        assert len(prepare_store_output.block_hashes_evicted) == 1
 
-    # store block 5, should evict from T1 (block 2, only one in T1)
-    prepare_store_output = arc_manager.prepare_store(to_hashes([5]))
-    verify_store_output(
-        prepare_store_output,
-        ExpectedPrepareStoreOutput(
-            block_hashes_to_store=[5],
-            store_block_ids=[1],  # reuses block 2's storage
-            block_hashes_evicted=[2],
-        ),
-    )
+        # complete store with failure
+        cpu_manager.complete_store(to_hashes([5]), success=False)
 
+        # block 5 should not be in cache
+        assert cpu_manager.lookup(to_hashes([5])) == 0
+        # block 5 should not be in T1 or T2
+        assert to_hashes([5])[0] not in arc_policy.t1
+        assert to_hashes([5])[0] not in arc_policy.t2
 
-def test_arc_manager_failed_store():
-    """
-    Tests that failed store operations clean up correctly.
-    Similar to LRU test but for ARC.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # evicted block should still be gone (in B1 ghost list)
+        evicted_hash = prepare_store_output.block_hashes_evicted[0]
+        assert evicted_hash in arc_policy.b1
 
-    # store blocks 1, 2, 3, 4
-    arc_manager.prepare_store(to_hashes([1, 2, 3, 4]))
-    arc_manager.complete_store(to_hashes([1, 2, 3, 4]))
+    def test_full_scenario(self):
+        """
+        Comprehensive test covering multiple ARC operations in sequence.
+        Similar to the full LRU test but adapted for ARC behavior.
+        """
+        cpu_manager, arc_policy = self._make_manager()
 
-    # prepare store block 5 (will evict block 1)
-    prepare_store_output = arc_manager.prepare_store(to_hashes([5]))
-    assert prepare_store_output is not None
-    assert len(prepare_store_output.block_hashes_evicted) == 1
+        # store [1, 2]
+        cpu_manager.prepare_store(to_hashes([1, 2]))
+        cpu_manager.complete_store(to_hashes([1, 2]))
 
-    # complete store with failure
-    arc_manager.complete_store(to_hashes([5]), success=False)
+        # store [3, 4, 5] -> evicts [1]
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([3, 4, 5]))
+        assert prepare_store_output is not None
+        assert len(prepare_store_output.block_hashes_evicted) == 1
+        cpu_manager.complete_store(to_hashes([3, 4, 5]))
 
-    # block 5 should not be in cache
-    assert arc_manager.lookup(to_hashes([5])) == 0
-    # block 5 should not be in T1 or T2
-    assert to_hashes([5])[0] not in arc_policy.t1
-    assert to_hashes([5])[0] not in arc_policy.t2
+        # promote some blocks to T2
+        cpu_manager.touch(to_hashes([2, 3]))
 
-    # evicted block should still be gone (in B1 ghost list)
-    evicted_hash = prepare_store_output.block_hashes_evicted[0]
-    assert evicted_hash in arc_policy.b1
+        # T1 has {4, 5}, T2 has {2, 3}
+        assert len(arc_policy.t1) == 2
+        assert len(arc_policy.t2) == 2
 
+        # store [6] -> should evict from T1 (4 is oldest in T1)
+        prepare_store_output = cpu_manager.prepare_store(to_hashes([6]))
+        assert prepare_store_output is not None
+        cpu_manager.complete_store(to_hashes([6]))
 
-def test_arc_manager_full_scenario():
-    """
-    Comprehensive test covering multiple ARC operations in sequence.
-    Similar to the full LRU test but adapted for ARC behavior.
-    """
-    block_size = 256
-    arc_manager = CPUOffloadingManager(
-        block_size=block_size, num_blocks=4, cache_policy="arc", enable_events=True
-    )
-    arc_policy = arc_manager._policy
-    assert isinstance(arc_policy, ARCCachePolicy)
+        # verify blocks 2, 3 (in T2) are still present
+        assert cpu_manager.lookup(to_hashes([2])) == 1
+        assert cpu_manager.lookup(to_hashes([3])) == 1
 
-    # store [1, 2]
-    arc_manager.prepare_store(to_hashes([1, 2]))
-    arc_manager.complete_store(to_hashes([1, 2]))
-
-    # store [3, 4, 5] -> evicts [1]
-    prepare_store_output = arc_manager.prepare_store(to_hashes([3, 4, 5]))
-    assert prepare_store_output is not None
-    assert len(prepare_store_output.block_hashes_evicted) == 1
-    arc_manager.complete_store(to_hashes([3, 4, 5]))
-
-    # promote some blocks to T2
-    arc_manager.touch(to_hashes([2, 3]))
-
-    # T1 has {4, 5}, T2 has {2, 3}
-    assert len(arc_policy.t1) == 2
-    assert len(arc_policy.t2) == 2
-
-    # store [6] -> should evict from T1 (4 is oldest in T1)
-    prepare_store_output = arc_manager.prepare_store(to_hashes([6]))
-    assert prepare_store_output is not None
-    arc_manager.complete_store(to_hashes([6]))
-
-    # verify blocks 2, 3 (in T2) are still present
-    assert arc_manager.lookup(to_hashes([2])) == 1
-    assert arc_manager.lookup(to_hashes([3])) == 1
-
-    # verify events
-    events = list(arc_manager.take_events())
-    assert len(events) > 0  # should have store and eviction events
+        # verify events
+        events = list(cpu_manager.take_events())
+        assert len(events) > 0  # should have store and eviction events
 
 
 def test_filter_reused_manager():
@@ -583,8 +549,6 @@ def test_filter_reused_manager():
         block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
     )
 
-    from vllm.v1.kv_offload.reuse_manager import FilterReusedOffloadingManager
-
     manager = FilterReusedOffloadingManager(
         backing=lru_manager, store_threshold=2, max_tracker_size=3
     )

tests/v1/structured_output/test_backend_guidance.py

@@ -11,6 +11,7 @@ from vllm.config.model import ModelConfig
 from vllm.config.parallel import ParallelConfig
 from vllm.config.speculative import SpeculativeConfig
 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
+from vllm.tokenizers import get_tokenizer
 from vllm.v1.request import Request
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.v1.structured_output.backend_guidance import GuidanceBackend
@@ -19,6 +20,14 @@ from vllm.v1.structured_output.backend_types import StructuredOutputOptions
 TOKENIZER = "gpt2"
 
 
+@pytest.fixture(scope="module")
+def mistral_tokenizer():
+    return get_tokenizer(
+        tokenizer_name="mistralai/Mistral-Small-3.2-24B-Instruct-2506",
+        tokenizer_mode="mistral",
+    )
+
+
 def test_backend_guidance_rollback_terminated():
     # Test that the backend guidance successfully rollbacks from a
     # terminated state. This can happen with speculative decoding,
@@ -187,3 +196,38 @@ def test_grammar_init_async_and_sync(async_grammar):
 
     # Verify the grammar can accept valid tokens
     assert grammar.accept_tokens(request.request_id, prompt)
+
+
+@pytest.mark.parametrize(
+    "request_type,grammar_spec",
+    [
+        pytest.param(
+            StructuredOutputOptions.JSON,
+            '{"type": "object"}',
+            id="json",
+        ),
+        pytest.param(
+            StructuredOutputOptions.GRAMMAR,
+            'start: "hello" | "world"',
+            id="lark",
+        ),
+    ],
+)
+def test_mistral_tokenizer_compile_grammar(
+    mistral_tokenizer,
+    request_type: StructuredOutputOptions,
+    grammar_spec: str,
+) -> None:
+    vllm_config = VllmConfig(
+        structured_outputs_config=StructuredOutputsConfig(backend="guidance"),
+    )
+    backend = GuidanceBackend(
+        vllm_config,
+        tokenizer=mistral_tokenizer,
+        vocab_size=mistral_tokenizer.vocab_size,
+    )
+    assert backend.ll_tokenizer is mistral_tokenizer.llg_tokenizer
+
+    grammar = backend.compile_grammar(request_type, grammar_spec)
+    assert grammar is not None
+    assert not grammar.is_terminated()

vllm/compilation/compiler_interface.py

@@ -345,9 +345,9 @@ class InductorStandaloneAdaptor(CompilerInterface):
         # Inductor's pre-grad passes don't do anything for vLLM.
         # The pre-grad passes get run even on cache-hit and negatively impact
         # vllm cold compile times by O(1s)
-        # Can remove this after the following issue gets fixed
+        # Fixed upstream in PyTorch 2.12:
         # https://github.com/pytorch/pytorch/issues/174502
-        if envs.VLLM_ENABLE_PREGRAD_PASSES:
+        if is_torch_equal_or_newer("2.12.0.dev") or envs.VLLM_ENABLE_PREGRAD_PASSES:
             pregrad_ctx: Any = contextlib.nullcontext()
         else:
             pregrad_ctx = patch(

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -12,6 +12,9 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
 from torch._inductor.pattern_matcher import PatternMatcherPass
 
 import vllm.ir.ops
+from vllm.compilation.passes.fusion.rms_quant_fusion import (
+    _rms_input_weight_dtype_match,
+)
 from vllm.config import VllmConfig
 from vllm.config.utils import Range
 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
@@ -320,7 +323,12 @@ class AllReduceRMSNormPattern(BasePattern):
             return allreduce[3], allreduce[1]
 
         pm.register_replacement(
-            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+            pattern,
+            replacement,
+            self.get_inputs(),
+            pm.fwd_only,
+            pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 
@@ -459,7 +467,12 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
             return allreduce[4], allreduce[1]
 
         pm.register_replacement(
-            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+            pattern,
+            replacement,
+            self.get_inputs(),
+            pm.fwd_only,
+            pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 
@@ -621,7 +634,12 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
             return allreduce[4], allreduce[1], allreduce[5]
 
         pm.register_replacement(
-            pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
+            pattern,
+            replacement,
+            self.get_inputs(),
+            pm.fwd_only,
+            pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 

vllm/compilation/passes/fusion/rms_quant_fusion.py

@@ -38,6 +38,22 @@ FP8_DTYPE = current_platform.fp8_dtype()
 FP4_DTYPE = torch.uint8
 
 
+_RMS_NORM_OP = torch.ops.vllm_ir.rms_norm.default
+
+
+# TODO: extend rmsnorm quant kernels to support mixed input/weight dtypes,
+# and remove this check.
+def _rms_input_weight_dtype_match(match: pm.Match) -> bool:
+    """Prevent fusion when rms_norm input and weight dtypes differ."""
+    for node in match.nodes:
+        if node.target == _RMS_NORM_OP:
+            # rms_norm(x, weight, epsilon, variance_size)
+            x, weight = node.args[0], node.args[1]
+            if isinstance(x, fx.Node) and isinstance(weight, fx.Node):
+                return x.meta["val"].dtype == weight.meta["val"].dtype
+    return True
+
+
 def empty_bf16(*args: Any, **kwargs: Any) -> torch.Tensor:
     return torch.empty(*args, **kwargs, dtype=torch.bfloat16, device="cuda")
 
@@ -186,7 +202,14 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
         ]
         pattern(*inputs)
 
-        pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
+        pm.register_replacement(
+            pattern,
+            replacement,
+            inputs,
+            pm.fwd_only,
+            pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
+        )
 
 
 class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
@@ -249,6 +272,7 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
             inputs,
             pm.fwd_only,
             pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 
@@ -350,6 +374,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
             self.rmsnorm_matcher.inputs() + [scale],
             pm.fwd_only,
             pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 
@@ -445,6 +470,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
             ],
             pm.fwd_only,
             pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 
@@ -503,6 +529,7 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
             ],
             pm.fwd_only,
             pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 
@@ -559,6 +586,7 @@ class FusedAddRMSNormDynamicQuantPattern(RMSNormQuantPattern):
             self.rmsnorm_matcher.inputs(),
             pm.fwd_only,
             pm_pass,
+            extra_check=_rms_input_weight_dtype_match,
         )
 
 

vllm/config/attention.py

@@ -30,7 +30,7 @@ class AttentionConfig:
     use_cudnn_prefill: bool = False
     """Whether to use cudnn prefill."""
 
-    use_trtllm_ragged_deepseek_prefill: bool = True
+    use_trtllm_ragged_deepseek_prefill: bool = False
     """Whether to use TRTLLM ragged deepseek prefill."""
 
     use_trtllm_attention: bool | None = None

vllm/config/speculative.py

@@ -817,6 +817,7 @@ class SpeculativeConfig:
             "deepseek_v3",
             "kimi_k2",
             "kimi_k25",
+            "minimax_m2",
         ]
         if (
             self.method in ("eagle3", "extract_hidden_states", "dflash")

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -5,7 +5,7 @@ KV cache helper for store.
 """
 
 from collections.abc import Iterator
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 import torch
@@ -516,6 +516,338 @@ class TpKVTopology:
         return cache if self.split_k_and_v else [cache]
 
 
+# ---- Mamba-HMA hetero-TP transfer config ----
+#
+# Key insight: with hetero-TP (P_TP > D_TP), FA KV cache may be
+# replicated across P ranks (when P_TP > num_kv_heads), but Mamba
+# conv/SSM state is almost always uniquely sharded per P rank.  So the
+# number of P ranks D must read from can differ between FA and Mamba,
+# and they must be handled separately.
+
+
+def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range:
+    """Physical KV head range stored in a rank's KV cache tensor.
+
+    When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank.
+    When ``tp_size > num_heads``: 1 physical head per rank.  Heads are
+    distributed **contiguously** (matching vLLM's GQA weight partitioning):
+    consecutive ranks share a head before moving to the next one.
+    """
+    if tp_size <= num_heads:
+        assert num_heads % tp_size == 0
+        per_rank = num_heads // tp_size
+        return range(rank * per_rank, (rank + 1) * per_rank)
+    else:
+        h = rank * num_heads // tp_size
+        return range(h, h + 1)
+
+
+def _range_overlap(a: range, b: range) -> range:
+    start = max(a.start, b.start)
+    stop = min(a.stop, b.stop)
+    return range(start, max(start, stop))
+
+
+@dataclass
+class HeteroTPTransferConfig:
+    """Precomputed transfer plan for one (D rank, P engine) pair.
+
+    Currently only instantiated for Mamba-HMA (hybrid SSM+Attention) models
+    where FA and mamba require different splitting factors.  Could be extended
+    to other model types that need non-uniform hetero-TP transfer sizing.
+
+    All descriptor sizes are computed here.  The guarantee is:
+        local_entry_size == remote_entry_size   (for NIXL)
+
+    Attributes that start with ``fa_`` concern FlashAttention KV cache.
+    Attributes that start with ``mamba_`` concern Mamba conv/SSM state.
+    """
+
+    # ---- Input parameters (from handshake) ----
+    tp_ratio: int
+    K: int  # total_num_kv_heads (before TP sharding)
+    d_tp: int  # D engine's tensor_parallel_size
+    p_tp: int  # P engine's tensor_parallel_size
+    d_rank: int  # this D worker's TP rank
+    use_mla: bool
+
+    # Per-layer block lengths (bytes, K+V combined for blocks_first).
+    # Uniform across layers for current models.
+    d_block_len: int  # D's block_len_per_layer (representative)
+    p_block_len: int  # P's block_len_per_layer (from handshake)
+    is_blocks_first: bool  # kv_topo.is_kv_layout_blocks_first
+
+    # ---- Derived: computed in __post_init__ ----
+    #
+    # Physical heads per rank (what the KV tensor actually stores)
+    d_physical_heads: int = field(init=False)
+    p_physical_heads: int = field(init=False)
+
+    # How many distinct P ranks D needs for FA data
+    physical_fa_num_reads: int = field(init=False)
+
+    # Which P ranks contribute unique FA heads (ordered by head index)
+    fa_read_targets: list[int] = field(init=False)
+
+    # All P ranks needed for mamba (always abs_tp for tp_ratio < 0)
+    mamba_num_reads: int = field(init=False)
+
+    # All P ranks this D rank communicates with (FA ∪ mamba)
+    transfer_targets: list[int] = field(init=False)
+
+    # FA descriptor entry size (K or V side, for blocks_first layout)
+    # Guaranteed: fa_entry_size is the SAME for local handle AND remote desc.
+    fa_entry_size: int = field(init=False)
+
+    # Replication flags
+    is_d_replicated: bool = field(init=False)
+    is_p_replicated: bool = field(init=False)
+
+    # Pre-built set for fast lookup
+    _fa_target_set: frozenset[int] = field(init=False, repr=False)
+    # Map: P rank → index in fa_read_targets (for head slot offset)
+    _fa_target_index: dict[int, int] = field(init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        K = self.K
+        self.is_d_replicated = self.d_tp > K
+        self.is_p_replicated = self.p_tp > K
+
+        self.d_physical_heads = max(1, K // self.d_tp)
+        self.p_physical_heads = max(1, K // self.p_tp)
+
+        abs_tp = -self.tp_ratio if self.tp_ratio < 0 else 1
+
+        # ---- Mamba range (computed first so FA can prefer ranks in it) ----
+        mamba_range: range | None = None
+        if self.tp_ratio < 0:
+            mamba_range = range(self.d_rank * abs_tp, (self.d_rank + 1) * abs_tp)
+
+        # ---- FA read targets ----
+        if self.use_mla or self.tp_ratio >= 0:
+            self.physical_fa_num_reads = 1
+            self.fa_read_targets = (
+                [0]
+                if self.use_mla
+                # Must match kv_topo.get_target_remote_ranks (d_rank // tp_ratio).
+                else [
+                    self.d_rank // self.tp_ratio if self.tp_ratio > 0 else self.d_rank
+                ]
+            )
+        else:
+            d_needs = _physical_head_range(self.d_tp, K, self.d_rank)
+            # When mamba range exists, prefer P ranks within it so that
+            # FA targets are a subset of mamba transfer_targets (avoids
+            # orphaned FA targets outside the transfer loop).
+            search_range = mamba_range if mamba_range is not None else range(self.p_tp)
+            seen: set[tuple[int, int]] = set()
+            targets: list[int] = []
+            for p in search_range:
+                p_has = _physical_head_range(self.p_tp, K, p)
+                ov = _range_overlap(d_needs, p_has)
+                if len(ov) > 0:
+                    key = (ov.start, ov.stop)
+                    if key not in seen:
+                        seen.add(key)
+                        targets.append(p)
+            if not targets:
+                # Fallback: search globally (should not happen in practice)
+                for p in range(self.p_tp):
+                    p_has = _physical_head_range(self.p_tp, K, p)
+                    ov = _range_overlap(d_needs, p_has)
+                    if len(ov) > 0:
+                        key = (ov.start, ov.stop)
+                        if key not in seen:
+                            seen.add(key)
+                            targets.append(p)
+            self.fa_read_targets = targets
+            self.physical_fa_num_reads = len(targets)
+
+        self._fa_target_set = frozenset(self.fa_read_targets)
+        self._fa_target_index = {r: i for i, r in enumerate(self.fa_read_targets)}
+
+        # ---- Mamba targets ----
+        if mamba_range is not None and abs_tp > self.physical_fa_num_reads:
+            self.mamba_num_reads = abs_tp
+            self.transfer_targets = list(mamba_range)
+        else:
+            self.mamba_num_reads = self.physical_fa_num_reads
+            self.transfer_targets = list(self.fa_read_targets)
+
+        # ---- FA entry size ----
+        # For blocks_first: block_len_per_layer includes K+V; // 2 gives K (or V).
+        # Use min(D, P) because D indexes into P when tp_ratio > 0,
+        # and P is the natural unit when tp_ratio < 0.
+        effective_block_len = min(self.d_block_len, self.p_block_len)
+        if self.is_blocks_first:
+            self.fa_entry_size = effective_block_len // 2
+        else:
+            self.fa_entry_size = effective_block_len
+
+        self._validate()
+
+    def _validate(self) -> None:
+        """Cross-check internal consistency."""
+        if self.is_d_replicated and self.is_p_replicated and self.tp_ratio > 0:
+            logger.info(
+                "Both-replicated hetero-TP: D_TP=%d > P_TP=%d > K=%d. "
+                "Using d_rank // tp_ratio routing with relative head offset.",
+                self.d_tp,
+                self.p_tp,
+                self.K,
+            )
+
+        # FA targets must be a subset of transfer_targets
+        tt_set = set(self.transfer_targets)
+        for t in self.fa_read_targets:
+            if t not in tt_set:
+                logger.error(
+                    "FA target P rank %d is NOT in transfer_targets %s. "
+                    "This will cause missed FA reads!",
+                    t,
+                    self.transfer_targets,
+                )
+
+        # For tp_ratio < 0 with blocks_first: D_K_half / reads should == P_K_half
+        if (
+            self.is_blocks_first
+            and self.tp_ratio < 0
+            and self.physical_fa_num_reads > 0
+        ):
+            d_k_half = self.d_block_len // 2
+            p_k_half = self.p_block_len // 2
+            expected_local = d_k_half // self.physical_fa_num_reads
+            if expected_local != p_k_half:
+                logger.warning(
+                    "FA size mismatch: D_K_half=%d / reads=%d = %d, "
+                    "but P_K_half=%d.  This may indicate a head count or "
+                    "Mamba-HMA inflation inconsistency.",
+                    d_k_half,
+                    self.physical_fa_num_reads,
+                    expected_local,
+                    p_k_half,
+                )
+
+    # ---- Query methods ----
+
+    def should_skip_fa(self, p_rank: int) -> bool:
+        """Whether to skip FA groups for this P rank (mamba-only transfer)."""
+        return p_rank not in self._fa_target_set
+
+    def fa_head_slot(self, p_rank: int) -> int:
+        """Index into D's FA block for this P rank's head data.
+
+        For P ranks in fa_read_targets, returns 0, 1, ..., reads-1.
+        For P ranks NOT in fa_read_targets (replicated duplicates),
+        returns the slot of the matching FA target with the same head.
+        """
+        if p_rank in self._fa_target_index:
+            return self._fa_target_index[p_rank]
+        # Duplicate head: find which fa_target has the same physical head
+        p_head = _physical_head_range(self.p_tp, self.K, p_rank)
+        for target in self.fa_read_targets:
+            t_head = _physical_head_range(self.p_tp, self.K, target)
+            if _range_overlap(p_head, t_head):
+                return self._fa_target_index[target]
+        return 0  # fallback
+
+    def fa_rank_offset(self, remote_kv_block_len: int) -> int:
+        """Byte offset into P's FA block for this D rank.
+
+        When D is replicated (D_TP > K), multiple D ranks share a head.
+        Computes offset *relative to the target P rank's first head*
+        so it works regardless of how many heads P has.
+        When neither side replicates, falls back to tp_rank % tp_ratio.
+        Returns 0 when D does not index into P's block.
+        """
+        if self.use_mla or self.tp_ratio <= 0:
+            return 0
+        if self.is_d_replicated:
+            d_head = self.d_rank * self.K // self.d_tp
+            p_rank = self.fa_read_targets[0]
+            p_start = p_rank * self.K // self.p_tp
+            return (d_head - p_start) * remote_kv_block_len
+        return self.d_rank % self.tp_ratio * remote_kv_block_len
+
+    @property
+    def needs_split_handles(self) -> bool:
+        """Whether per-P-rank split handles are needed.
+
+        True when FA and mamba have different read counts, requiring
+        different splitting factors in the local handle.
+        """
+        return self.tp_ratio < 0 and not self.use_mla and len(self.transfer_targets) > 1
+
+    def compute_split_handle_data(
+        self,
+        src_blocks_data: list[tuple[int, int, int]],
+        num_fa_descs: int,
+        abs_tp: int,
+    ) -> list[list[tuple[int, int, int]]]:
+        """Compute per-P-rank (addr, len, tp) triples for Mamba-HMA split handles.
+
+        FA descriptors (indices < num_fa_descs) are sliced by
+        ``physical_fa_num_reads``; mamba descriptors are sliced uniformly
+        by ``abs_tp``.
+
+        Returns one list of triples per transfer target.
+        """
+        all_handle_data: list[list[tuple[int, int, int]]] = []
+        for p_idx, p_rank in enumerate(self.transfer_targets):
+            handle_data: list[tuple[int, int, int]] = []
+            skip_fa = self.should_skip_fa(p_rank)
+            fa_slot = self.fa_head_slot(p_rank) if not skip_fa else 0
+
+            for j, (addr, local_len, tp) in enumerate(src_blocks_data):
+                if j < num_fa_descs:
+                    assert self.physical_fa_num_reads >= 1
+                    fa_chunk = local_len // self.physical_fa_num_reads
+                    handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, tp))
+                else:
+                    mamba_chunk = local_len // abs_tp
+                    handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, tp))
+            all_handle_data.append(handle_data)
+        return all_handle_data
+
+    def filter_block_ids_for_rank(
+        self,
+        remote_rank: int,
+        local_ids: BlockIds,
+        remote_ids: BlockIds,
+        is_mamba_group: list[bool],
+    ) -> tuple[BlockIds, BlockIds]:
+        """Zero out FA groups for P ranks outside fa_read_targets.
+
+        Returns (filtered_local_ids, filtered_remote_ids).  When the
+        remote rank carries FA data for this D rank, returns the inputs
+        unchanged.
+        """
+        if not self.should_skip_fa(remote_rank):
+            return local_ids, remote_ids
+        num_groups = len(local_ids)
+        filtered_local: list[list[int]] = [
+            [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups)
+        ]
+        filtered_remote: list[list[int]] = [
+            [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups)
+        ]
+        return filtered_local, filtered_remote
+
+    def describe(self) -> str:
+        """One-line summary for logging."""
+        return (
+            f"HeteroTPTransferConfig("
+            f"tp_ratio={self.tp_ratio}, K={self.K}, "
+            f"d_tp={self.d_tp}, p_tp={self.p_tp}, d_rank={self.d_rank}, "
+            f"physical_fa_reads={self.physical_fa_num_reads}, "
+            f"mamba_reads={self.mamba_num_reads}, "
+            f"fa_targets={self.fa_read_targets}, "
+            f"transfer_targets={self.transfer_targets}, "
+            f"fa_entry_size={self.fa_entry_size}, "
+            f"d_block_len={self.d_block_len}, p_block_len={self.p_block_len})"
+        )
+
+
 def get_current_attn_backends(
     vllm_config: VllmConfig, layer_names: list[str] | None = None
 ) -> list[type[AttentionBackend]]:
@@ -559,3 +891,50 @@ def get_current_attn_backend(
 ) -> type[AttentionBackend]:
     """Get the first attention backend for the given layers."""
     return get_current_attn_backends(vllm_config, layer_names)[0]
+
+
+# TODO (ZhanqiuHu): Consolidate TpKVTopology and HeteroTPTransferConfig
+# into a single engine-agnostic TransferTopology class.
+# 6 of 9 HeteroTPTransferConfig init fields duplicate TpKVTopology data.
+#
+# @dataclass
+# class EngineTransferInfo:
+#     """Per-remote-engine transfer state, computed at handshake."""
+#     p_tp: int
+#     tp_ratio: int
+#     p_block_len: int
+#     block_size: int
+#     # Mamba-specific (None for non-mamba models)
+#     fa_read_targets: list[int] | None = None
+#     transfer_targets: list[int] | None = None
+#     physical_fa_num_reads: int | None = None
+#     mamba_num_reads: int | None = None
+#     fa_entry_size: int | None = None
+#
+# class TransferTopology:
+#     """Single source of truth for TP topology + transfer sizing."""
+#     # Shared (set once at init, replaces duplicate fields)
+#     tp_rank: int          # == TpKVTopology.tp_rank == HeteroTP.d_rank
+#     tp_size: int          # == TpKVTopology.tp_size == HeteroTP.d_tp
+#     total_num_kv_heads: int  # == HeteroTP.K
+#     is_mla: bool          # == HeteroTP.use_mla
+#     is_mamba: bool
+#     is_blocks_first: bool # == HeteroTP.is_blocks_first
+#     d_block_len: int
+#
+#     # Per-engine (populated via register_engine() at handshake)
+#     _engines: dict[EngineId, EngineTransferInfo]
+#
+#     def register_engine(self, engine_id, p_tp, p_block_len, ...): ...
+#
+#     # General (from TpKVTopology)
+#     def tp_ratio(self, engine_id) -> int: ...
+#     def target_remote_ranks(self, engine_id) -> list[int]: ...
+#     def is_kv_replicated(self, engine_id) -> bool: ...
+#
+#     # Mamba-specific (from HeteroTPTransferConfig, gated by is_mamba)
+#     def fa_rank_offset(self, engine_id, block_len) -> int: ...
+#     def physical_fa_num_reads(self, engine_id) -> int: ...
+#     def transfer_targets(self, engine_id) -> list[int]: ...
+#     def should_skip_fa(self, engine_id, p_rank) -> bool: ...
+#     def filter_block_ids_for_rank(self, engine_id, ...) -> ...: ...

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

@@ -25,6 +25,7 @@ from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     BlockIds,
     EngineId,
+    HeteroTPTransferConfig,
     TpKVTopology,
     get_current_attn_backend,
     get_current_attn_backends,
@@ -47,12 +48,18 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
     PromMetric,
     PromMetricT,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
+    MambaConvSplitInfo,
+    compute_mamba_phys_ratio,
+    derive_mamba_conv_split,
+)
 from vllm.distributed.parallel_state import (
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
 from vllm.forward_context import ForwardContext
 from vllm.logger import init_logger
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
@@ -1038,7 +1045,7 @@ class NixlConnectorWorker:
         }
         self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
 
-        # Mamba metadata
+        # ---- Mamba model state (derived from model config) ----
         self._is_mamba_group = [
             isinstance(group.kv_cache_spec, MambaSpec)
             for group in kv_cache_config.kv_cache_groups
@@ -1065,6 +1072,17 @@ class NixlConnectorWorker:
                 ssm_shape.numel() * ssm_nbytes,
             )
         self._mamba_ssm_size = mamba_ssm_size
+        # Conv state sub-projection decomposition (None when no Mamba).
+        # The 3-read transfer requires DS (dim, state_len) conv layout so
+        # that x/B/C sub-projections are contiguous in memory.
+        self._conv_decomp: MambaConvSplitInfo | None = None
+        if self._has_mamba:
+            assert is_conv_state_dim_first(), (
+                "3-read Mamba conv transfer requires DS conv state layout. "
+                "Set VLLM_SSM_CONV_STATE_LAYOUT=DS"
+            )
+            local_tp = vllm_config.parallel_config.tensor_parallel_size
+            self._conv_decomp = derive_mamba_conv_split(mamba_spec, local_tp)
 
         # Agent.
         non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
@@ -1175,6 +1193,16 @@ class NixlConnectorWorker:
         self.dst_num_blocks: dict[EngineId, int] = {}
         self._registered_descs: list[Any] = []
 
+        # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ----
+        # Per-engine transfer config (source of truth for FA/mamba sizing).
+        self._transfer_configs: dict[str, HeteroTPTransferConfig] = {}
+        # NOTE (ZhanqiuHu): _mamba_phys_ratio MUST be per-engine.
+        # compute_mamba_phys_ratio = ceil((conv_bytes + ssm_bytes) / block_len)
+        # where conv/ssm bytes are per-TP-rank (dimension-sharded).  With
+        # heterogeneous TP the per-rank sizes differ, so the ratio differs:
+        #   e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261.
+        self._mamba_phys_ratio: dict[EngineId, int] = {}
+
         # In progress transfers.
         # [req_id -> list[handle]]
         self._recving_metadata: dict[ReqId, ReqMeta] = {}
@@ -1701,8 +1729,7 @@ class NixlConnectorWorker:
             # then duplicate it logically to be able to index SSM/Conv separately.
             self.num_regions *= 2
 
-        # TODO (NickLucche) Adapt to different descs views (engine_id->tp_rank) to
-        # support heterogeneous TP.
+        # Total local FA descriptors (boundary between FA and mamba descs).
         self.num_descs = self.num_regions * self.num_blocks
 
         descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
@@ -1715,6 +1742,9 @@ class NixlConnectorWorker:
         self.dst_num_blocks[self.engine_id] = self.num_blocks
 
         if self._has_mamba:
+            self._mamba_phys_ratio[self.engine_id] = (
+                self._physical_blocks_per_logical_kv_block
+            )
             logger.info(
                 "Hybrid SSM registration: num_blocks=%s, "
                 "logical_num_blocks=%s, ratio=%s, num_regions=%s, "
@@ -1755,6 +1785,149 @@ class NixlConnectorWorker:
             agent_metadata_bytes=encoder.encode(agent_metadata),
         )
 
+    def _build_mamba_local(
+        self,
+        base_addresses: list[int],
+        block_size_ratio: int,
+    ) -> list[tuple[int, int, int]]:
+        """Build 4 desc regions (x, B, C, ssm) per layer for local mamba
+        blocks, enabling the 3-read transfer with DS conv layout."""
+        assert block_size_ratio == 1, (
+            "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
+            f"Got block_size_ratio={block_size_ratio}."
+        )
+        assert self._conv_decomp is not None
+        conv_offsets = self._conv_decomp.local_conv_offsets
+        conv_size, ssm_size = self._mamba_ssm_size
+        num_blocks = self._logical_num_blocks * block_size_ratio
+        phys_ratio = self._physical_blocks_per_logical_kv_block
+
+        result: list[tuple[int, int, int]] = []
+        for i, base_addr in enumerate(base_addresses):
+            page_stride = self.block_len_per_layer[i] // block_size_ratio * phys_ratio
+            for off, sz in conv_offsets:
+                for blk in range(num_blocks):
+                    result.append(
+                        (base_addr + blk * page_stride + off, sz, self.device_id)
+                    )
+            # SSM temporal state follows the conv state.
+            for blk in range(num_blocks):
+                result.append(
+                    (
+                        base_addr + blk * page_stride + conv_size,
+                        ssm_size,
+                        self.device_id,
+                    )
+                )
+        return result
+
+    def _build_fa_remote_for_mamba(
+        self,
+        nixl_agent_meta: NixlAgentMetadata,
+        transfer_cfg: HeteroTPTransferConfig,
+        block_size_ratio: int,
+        kv_topo: TpKVTopology,
+    ) -> list[tuple[int, int, int]]:
+        """Build remote FA descriptors for mamba models.
+
+        Uses transfer_cfg for GQA-aware FA divisor and head-based rank offset
+        instead of the standard uniform tp_ratio split.
+        """
+        assert block_size_ratio == 1, (
+            "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
+            f"Got block_size_ratio={block_size_ratio}."
+        )
+        # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA
+        # hetero-TP logic stabilizes.
+        tp_ratio = transfer_cfg.tp_ratio
+        result: list[tuple[int, int, int]] = []
+        for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
+            local_block_len = self.get_backend_aware_kv_block_len(
+                layer_idx=i, first_split=True, mamba_view=False
+            )
+            remote_kv_block_len = local_block_len // block_size_ratio
+            if block_size_ratio > 1:
+                local_block_len = remote_kv_block_len
+
+            if tp_ratio < 0 and not self.use_mla:
+                local_block_len = local_block_len // transfer_cfg.physical_fa_num_reads
+
+            rank_offset = transfer_cfg.fa_rank_offset(remote_kv_block_len)
+
+            num_blocks = nixl_agent_meta.num_blocks
+            page_size = nixl_agent_meta.block_lens[i]
+            for block_id in range(num_blocks):
+                block_offset = block_id * page_size
+                addr = base_addr + block_offset + rank_offset
+                result.append((addr, local_block_len, nixl_agent_meta.device_id))
+
+            if kv_topo.is_kv_layout_blocks_first:
+                second_split = self.get_backend_aware_kv_block_len(
+                    layer_idx=i, first_split=False, mamba_view=False
+                )
+                if tp_ratio < 0 and not self.use_mla:
+                    second_split = second_split // transfer_cfg.physical_fa_num_reads
+                for block_id in range(num_blocks):
+                    block_offset = block_id * page_size
+                    addr = base_addr + block_offset + rank_offset
+                    v_addr = addr + nixl_agent_meta.block_lens[i] // 2
+                    result.append((v_addr, second_split, nixl_agent_meta.device_id))
+        return result
+
+    def _build_mamba_remote(
+        self,
+        nixl_agent_meta: NixlAgentMetadata,
+        tp_ratio: int,
+    ) -> list[tuple[int, int, int]]:
+        """Build 4 remote desc regions (x, B, C, ssm) per layer for
+        the 3-read transfer.  For hetero-TP, each D rank reads only its
+        sub-projection slice from the P rank."""
+        assert self._conv_decomp is not None
+        effective_ratio = max(tp_ratio, 1)
+        # Mamba conv state is always TP-sharded, even when attention KV
+        # is replicated (num_kv_heads < tp_size).
+        local_offset = self.tp_rank % effective_ratio
+        conv_size_remote = nixl_agent_meta.ssm_sizes[0]
+
+        if tp_ratio >= 1:
+            # D_TP >= P_TP: P page is larger, D reads its slice.
+            conv_offsets = self._conv_decomp.remote_conv_offsets(
+                local_offset, effective_ratio
+            )
+            ssm_read_size = self._mamba_ssm_size[1]
+        else:
+            # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages
+            # are smaller than D's.  self._conv_decomp has D-sized dimensions,
+            # but we need P-sized offsets.  Scale down by |tp_ratio|.
+            abs_ratio = -tp_ratio
+            xb_p = self._conv_decomp.x_bytes // abs_ratio
+            bb_p = self._conv_decomp.b_bytes // abs_ratio
+            conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)]
+            ssm_read_size = nixl_agent_meta.ssm_sizes[1]
+
+        remote_ratio = self._mamba_phys_ratio[nixl_agent_meta.engine_id]
+        num_blocks = nixl_agent_meta.num_blocks // remote_ratio
+        device_id = nixl_agent_meta.device_id
+
+        result: list[tuple[int, int, int]] = []
+        # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
+        # block lengths vary across layers (e.g. MLA).
+        for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
+            page_stride = nixl_agent_meta.block_lens[i] * remote_ratio
+            for off, sz in conv_offsets:
+                for blk in range(num_blocks):
+                    result.append((base_addr + blk * page_stride + off, sz, device_id))
+            # SSM temporal state is also TP-sharded on the heads dimension.
+            for blk in range(num_blocks):
+                ssm_addr = (
+                    base_addr
+                    + blk * page_stride
+                    + conv_size_remote
+                    + local_offset * ssm_read_size
+                )
+                result.append((ssm_addr, ssm_read_size, device_id))
+        return result
+
     def register_local_xfer_handler(
         self,
         block_size: int,
@@ -1823,13 +1996,22 @@ class NixlConnectorWorker:
                 self.device_id,
             )
 
+        # NOTE (ZhanqiuHu): mamba=True path in register_blocks is not used
+        # right now — we use _build_mamba_local instead for the 3-read
+        # approach. However, we might still need this as a fallback for homogeneous TP.
         register_blocks(blocks_data, mamba=False)
         if self._has_mamba:
             assert self.num_descs == len(blocks_data)
-            logger.debug(
-                "Registering additional %s local Mamba blocks", len(blocks_data)
+            # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-read split is
+            # unnecessary — a single conv desc per block suffices.  Consider
+            # adding a fast path that falls back to the standard 2-region
+            # registration (register_blocks mamba=True) when no hetero-TP
+            # remote has been seen.  Currently we always register 4 regions
+            # because local descs are created before knowing the remote TP.
+            logger.debug("Registering local Mamba descriptors (4 regions/layer)")
+            blocks_data.extend(
+                self._build_mamba_local(local_base_addresses, block_size_ratio)
             )
-            register_blocks(blocks_data, mamba=True)
 
         descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
         # NIXL_INIT_AGENT to be used for preparations of local descs.
@@ -1880,6 +2062,9 @@ class NixlConnectorWorker:
 
         Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0
         so that the whole cache is shared by "tp_ratio" D TP workers.
+
+        For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and
+        tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer.
         """  # noqa: E501
         engine_id = nixl_agent_meta.engine_id
         # TODO re-evaluate refreshing for scaling/recovery
@@ -1915,6 +2100,10 @@ class NixlConnectorWorker:
 
         if engine_id not in self.dst_num_blocks:
             self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
+        if self._has_mamba:
+            self._mamba_phys_ratio[engine_id] = compute_mamba_phys_ratio(
+                nixl_agent_meta.ssm_sizes, nixl_agent_meta.block_lens[0]
+            )
 
         # Keep track of remote agent kv caches base addresses.
         self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
@@ -1931,6 +2120,21 @@ class NixlConnectorWorker:
             not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0
         )
 
+        # Create transfer config (single source of truth for descriptor sizes).
+        if self._has_mamba and engine_id not in self._transfer_configs:
+            self._transfer_configs[engine_id] = HeteroTPTransferConfig(
+                tp_ratio=tp_ratio,
+                K=kv_topo.total_num_kv_heads,
+                d_tp=self.world_size,
+                p_tp=remote_tp_size,
+                d_rank=self.tp_rank,
+                use_mla=self.use_mla,
+                d_block_len=self.block_len_per_layer[0],
+                p_block_len=nixl_agent_meta.block_lens[0],
+                is_blocks_first=kv_topo.is_kv_layout_blocks_first,
+            )
+            logger.info("Created %s", self._transfer_configs[engine_id].describe())
+
         logger.debug(
             "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
             engine_id,
@@ -1947,21 +2151,48 @@ class NixlConnectorWorker:
             # Remote tp_size > local tp_size: read from multiple remote ranks.
             # Logically "split" own regions into |tp_ratio| chunks. Mind that
             # we only do this once per remote tp_size (replica-friendly).
+            abs_tp = -tp_ratio
             self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
-            for i in range(-tp_ratio):
-                blocks_data = []
-                for memory_region in self.src_blocks_data:
-                    addr, local_block_len, own_tp_rank = memory_region
-                    # Computing block len layer by layer allows for different
-                    # block sizes to be used.
-                    remote_block_len = local_block_len // (-tp_ratio)
-                    addr = addr + i * remote_block_len
-                    blocks_data.append((addr, remote_block_len, own_tp_rank))
-                descs = self.nixl_wrapper.get_xfer_descs(
-                    blocks_data, self.nixl_memory_type
-                )
-                handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
-                self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
+
+            if self._has_mamba:
+                transfer_cfg = self._transfer_configs.get(engine_id)
+                assert transfer_cfg is not None
+                if transfer_cfg.needs_split_handles:
+                    # Mamba-HMA: FA and Mamba use different split factors.
+                    for handle_data in transfer_cfg.compute_split_handle_data(
+                        self.src_blocks_data, self.num_descs, abs_tp
+                    ):
+                        descs = self.nixl_wrapper.get_xfer_descs(
+                            handle_data, self.nixl_memory_type
+                        )
+                        handle = self.nixl_wrapper.prep_xfer_dlist(
+                            "NIXL_INIT_AGENT", descs
+                        )
+                        self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
+
+                    logger.info(
+                        "Mamba-HMA split handles: targets=%s, fa_reads=%s, "
+                        "fa_entry=%s, mamba_reads=%s, num_descs=%s",
+                        transfer_cfg.transfer_targets,
+                        transfer_cfg.physical_fa_num_reads,
+                        transfer_cfg.fa_entry_size,
+                        transfer_cfg.mamba_num_reads,
+                        self.num_descs,
+                    )
+            else:
+                # Original path: uniform divide by abs_tp (non-Mamba-HMA).
+                for i in range(abs_tp):
+                    blocks_data = []
+                    for memory_region in self.src_blocks_data:
+                        addr, local_block_len, own_tp_rank = memory_region
+                        remote_block_len = local_block_len // abs_tp
+                        addr = addr + i * remote_block_len
+                        blocks_data.append((addr, remote_block_len, own_tp_rank))
+                    descs = self.nixl_wrapper.get_xfer_descs(
+                        blocks_data, self.nixl_memory_type
+                    )
+                    handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
+                    self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
 
         ### Register remote agent memory regions
         blocks_data = []
@@ -2044,13 +2275,33 @@ class NixlConnectorWorker:
                 self.tp_rank,
             )
 
-        register_remote_blocks(blocks_data, mamba=False)
         if self._has_mamba:
-            # Create extra descs for the Mamba "view" of the same KV cache tensors.
+            # Mamba-HMA: separate FA registration with GQA-aware sizing,
+            # plus mamba 3-read registration for the Mamba "view" of the
+            # same KV cache tensors.
             logger.debug(
-                "Registering additional %s remote Mamba blocks", len(blocks_data)
+                "Registering remote Mamba blocks for engine %s rank %s",
+                engine_id,
+                remote_tp_rank,
             )
-            register_remote_blocks(blocks_data, mamba=True)
+            transfer_cfg = self._transfer_configs.get(engine_id)
+            assert transfer_cfg is not None
+            blocks_data.extend(
+                self._build_fa_remote_for_mamba(
+                    nixl_agent_meta,
+                    transfer_cfg,
+                    block_size_ratio,
+                    kv_topo,
+                )
+            )
+            blocks_data.extend(
+                self._build_mamba_remote(
+                    nixl_agent_meta,
+                    tp_ratio,
+                )
+            )
+        else:
+            register_remote_blocks(blocks_data, mamba=False)
 
         # Register with NIXL.
         descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
@@ -2083,17 +2334,17 @@ class NixlConnectorWorker:
         block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
             remote_engine_id
         )
-        # Num kv_heads > tp_size and P TP > D TP case, not supported
-        assert not (tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id))
+        # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
+        # Mamba models can have replicated FA KV with tp_ratio < 0.
+        if not self._has_mamba:
+            assert not (
+                tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id)
+            )
 
         if self._is_hma_required:
             assert block_size_ratio == 1, (
                 "HMA does not support different remote block size yet"
             )
-        # Mamba additional constraints
-        if self._has_mamba:
-            assert tp_ratio == 1, "Mamba does not support heterogeneous TP yet"
-
         kv_cache_layout = (
             self.kv_cache_layout
             if not self.use_host_buffer
@@ -2138,11 +2389,14 @@ class NixlConnectorWorker:
         remote_block_len = nixl_agent_meta.block_lens[0]
         if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id):
             # With replicated KV cache, only the number of blocks can differ.
-            for i in range(len(self.block_len_per_layer)):
-                assert (
-                    self.block_len_per_layer[i] // block_size_ratio
-                    == nixl_agent_meta.block_lens[i]
-                ), "KV cache sizes must match between P and D when replicated"
+            # TODO (ZhanqiuHu): For mamba models, validate FA and mamba
+            # block_lens separately.
+            if not self._has_mamba:
+                for i in range(len(self.block_len_per_layer)):
+                    assert (
+                        self.block_len_per_layer[i] // block_size_ratio
+                        == nixl_agent_meta.block_lens[i]
+                    ), "KV cache sizes must match between P and D when replicated"
         else:
             # When MLA is not used, this is a list of the same block length
             for block_len in nixl_agent_meta.block_lens:
@@ -2150,25 +2404,31 @@ class NixlConnectorWorker:
                     "All remote layers must have the same block size"
                 )
 
-            if tp_ratio > 0:
-                # Remote tp is smaller: remote block_len size is bigger
-                assert (
-                    remote_block_len
-                    == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio
-                ), (
-                    "Remote P worker KV layer cache must be of shape [2, N, "
-                    "local_kv_heads*tp_ratio, page_size, head_dim] and same dtype."
-                )  # noqa: E501
-            else:
-                assert block_size_ratio == 1, (
-                    "Different local/remote block sizes are not supported when"
-                    " P TP > D TP."
-                )
-                # Remote tp is bigger: remote block_len size is smaller
-                assert remote_block_len == self.block_len_per_layer[0] // (-tp_ratio), (
-                    "Remote P worker KV layer cache must be of shape [2, N, "
-                    "local_kv_heads/tp_ratio, page_size, head_dim] and same dtype."
-                )  # noqa: E501
+            # HMA hybrid models (mamba+attention) pad block_len to
+            # max(attn_page, mamba_page), so the linear tp_ratio scaling
+            # assumption only holds for pure-attention models.
+            if not self._has_mamba:
+                if tp_ratio > 0:
+                    assert (
+                        remote_block_len
+                        == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio
+                    ), (
+                        "Remote P worker KV layer cache must be of shape [2, N,"
+                        " local_kv_heads*tp_ratio, page_size, head_dim] and "
+                        "same dtype."
+                    )
+                else:
+                    assert block_size_ratio == 1, (
+                        "Different local/remote block sizes are not supported"
+                        " when P TP > D TP."
+                    )
+                    assert remote_block_len == self.block_len_per_layer[0] // (
+                        -tp_ratio
+                    ), (
+                        "Remote P worker KV layer cache must be of shape [2, N,"
+                        " local_kv_heads/tp_ratio, page_size, head_dim] and "
+                        "same dtype."
+                    )
 
         # TP workers that handhshake with same remote have same #blocks.
         assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
@@ -2471,9 +2731,8 @@ class NixlConnectorWorker:
                 meta.local_block_ids
             )
             assert meta.remote is not None
-            meta.remote.block_ids = self._logical_to_kernel_block_ids(
-                meta.remote.block_ids
-            )
+            # Remote block IDs are kept logical here; expanded in
+            # _read_blocks_for_req using the remote engine's phys ratio.
             remote_engine_id = meta.remote.engine_id
             logger.debug(
                 "start_load_kv for request %s from remote engine %s. "
@@ -2525,6 +2784,13 @@ class NixlConnectorWorker:
             meta.remote.engine_id
         )
         tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id)
+
+        if self._has_mamba:
+            # Expand remote logical → kernel block IDs.
+            meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
+                meta.remote.block_ids,
+                self._mamba_phys_ratio[meta.remote.engine_id],
+            )
         # D may have to perform multiple reads from different remote ranks.
         for i, remote_rank in enumerate(remote_ranks):
             if self.use_mla and tp_ratio < 0 and i > 0:
@@ -2558,12 +2824,26 @@ class NixlConnectorWorker:
             remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
                 remote_rank
             ]
+
+            local_ids: BlockIds = meta.local_physical_block_ids
+            remote_ids: BlockIds = meta.remote.block_ids
+            if self._has_mamba:
+                # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets.
+                transfer_cfg = self._transfer_configs.get(meta.remote.engine_id)
+                assert transfer_cfg is not None
+                local_ids, remote_ids = transfer_cfg.filter_block_ids_for_rank(
+                    remote_rank,
+                    local_ids,
+                    remote_ids,
+                    self._is_mamba_group,
+                )
+
             self._read_blocks(
                 request_id=req_id,
                 dst_engine_id=meta.remote.engine_id,
                 remote_request_id=meta.remote.request_id,
-                local_block_ids=meta.local_physical_block_ids,
-                remote_block_ids=meta.remote.block_ids,
+                local_block_ids=local_ids,
+                remote_block_ids=remote_ids,
                 remote_rank=remote_rank,
                 local_xfer_side_handle=local_xfer_side_handle,
                 remote_xfer_side_handle=remote_xfer_side_handle,
@@ -2663,9 +2943,12 @@ class NixlConnectorWorker:
         for i, remote_group in enumerate(remote_block_ids):
             num_remote_blocks = len(remote_group)
             num_local_blocks = len(local_block_ids[i])
-            assert num_local_blocks <= num_remote_blocks
+            if not self._is_mamba_group[i]:
+                assert num_local_blocks <= num_remote_blocks
             # Partial prefix cache hit: just read uncomputed blocks.
-            if num_local_blocks < num_remote_blocks:
+            # Skip mamba groups — their blocks represent full state (conv+ssm),
+            # not per-token data, so trimming would corrupt the transfer.
+            if num_local_blocks < num_remote_blocks and not self._is_mamba_group[i]:
                 remote_block_ids[i] = remote_group[-num_local_blocks:]
 
         # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
@@ -2781,16 +3064,22 @@ class NixlConnectorWorker:
             # This is like having two "low-level views" of the same storage.
             # `num_fa_descs` offset must be computed per-engine since P and D can
             # have different num_blocks (and thus different FA descs counts).
-            ratio = self._physical_blocks_per_logical_kv_block
-            # SSM may register fewer num_blocks than FA
+            ratio = self._mamba_phys_ratio[engine_id]
             logical_blocks = num_blocks // ratio
             num_fa_descs = self.num_regions * num_blocks
+            # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm).
+            mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None]
             all_descs = []
             for i, group in enumerate(block_ids):
-                stride = logical_blocks if self._is_mamba_group[i] else num_blocks
                 group_arr = np.asarray(group)[None, :]
-                offset = num_fa_descs if self._is_mamba_group[i] else 0
-                all_descs.append((region_ids * stride + group_arr + offset).flatten())
+                if self._is_mamba_group[i]:
+                    all_descs.append(
+                        (
+                            mamba_region_ids * logical_blocks + group_arr + num_fa_descs
+                        ).flatten()
+                    )
+                else:
+                    all_descs.append((region_ids * num_blocks + group_arr).flatten())
             return np.concatenate(all_descs)
 
     def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
@@ -2818,6 +3107,36 @@ class NixlConnectorWorker:
             for i, group in enumerate(block_ids)
         ]
 
+    def _logical_to_remote_kernel_block_ids(
+        self, block_ids: BlockIds, remote_ratio: int
+    ) -> BlockIds:
+        """Map logical block IDs to physical kernel block IDs on the remote.
+
+        Args:
+            block_ids: per-group lists of logical block IDs.
+            remote_ratio: remote engine's physical blocks per logical block.
+
+        Returns:
+            Same structure with FA groups expanded (each logical block L
+            becomes kernel blocks [L*remote_ratio .. L*remote_ratio +
+            local_ratio - 1]).  Mamba groups are passed through unchanged.
+        """
+        local_ratio = self._physical_blocks_per_logical_kv_block
+        if remote_ratio == 1:
+            return block_ids
+        local_arange = np.arange(local_ratio).reshape(1, -1)
+        group_specs = self.kv_cache_config.kv_cache_groups
+        result: list[list[int]] = []
+        for i, group in enumerate(block_ids):
+            if not isinstance(group_specs[i].kv_cache_spec, MambaSpec):
+                arr = np.array(group).reshape(-1, 1)
+                expanded = (arr * remote_ratio + local_arange).flatten()
+                result.append(expanded.tolist())
+            else:
+                # Mamba blocks are 1:1 logical-to-physical (no expansion).
+                result.append(group)
+        return result
+
     def get_backend_aware_kv_block_len(
         self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
     ) -> int:

vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Mamba conv-state sub-projection decomposition for the 3-read transfer.
+
+With DS conv state layout (dim, state_len), x/B/C sub-projections are
+contiguous in memory.  Each D rank reads its x, B, C slices via 3
+separate RDMA transfers — no P-side permutation needed.
+"""
+
+import math
+from dataclasses import dataclass
+
+import torch
+
+from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
+from vllm.v1.kv_cache_interface import MambaSpec
+
+
+@dataclass(frozen=True)
+class MambaConvSplitInfo:
+    """Per-rank byte sizes of x, B, C sub-projections in the Mamba conv state.
+
+    Used by both P and D sides for NIXL descriptor registration.
+    All fields are LOCAL to this engine's TP (already divided by TP size).
+
+    DS memory layout within one page (contiguous in memory):
+        |--- x (x_local * conv_rows) ---|- B (b_local * conv_rows) -|- C -|
+    """
+
+    conv_rows: int  # conv_kernel - 1 (typically 3)
+    x_local: int  # intermediate_size / TP  (columns for x)
+    b_local: int  # groups_ss / TP  (columns for B; C is same size)
+    conv_dtype_size: int  # bytes per element (e.g. 2 for float16)
+
+    @property
+    def conv_dim_local(self) -> int:
+        """Total conv columns per rank: x + B + C."""
+        return self.x_local + 2 * self.b_local
+
+    @property
+    def x_bytes(self) -> int:
+        """Byte size of the x sub-projection for one rank."""
+        return self.x_local * self.conv_rows * self.conv_dtype_size
+
+    @property
+    def b_bytes(self) -> int:
+        """Byte size of the B (or C) sub-projection for one rank."""
+        return self.b_local * self.conv_rows * self.conv_dtype_size
+
+    @property
+    def local_conv_offsets(self) -> list[tuple[int, int]]:
+        """(byte_offset, byte_size) of x, B, C within this engine's page.
+
+        Used by both P and D for local descriptor registration.
+        """
+        xb = self.x_bytes
+        bb = self.b_bytes
+        return [(0, xb), (xb, bb), (xb + bb, bb)]
+
+    def remote_conv_offsets(
+        self, local_rank_offset: int, tp_ratio: int
+    ) -> list[tuple[int, int]]:
+        """(byte_offset, byte_size) of this D rank's x, B, C slice within
+        one P page.
+
+        Used by D side only, during remote descriptor registration.
+
+        Args:
+            local_rank_offset: which slice this D rank reads.
+                tp_ratio > 0: tp_rank % tp_ratio (selects slice of P's page).
+                tp_ratio < 0: always 0 (read P's full page).
+            tp_ratio: effective ratio (>= 1 when D_TP > P_TP, 1 when
+                P_TP > D_TP since each P rank is read in full).
+        """
+        xb = self.x_bytes
+        bb = self.b_bytes
+        xr = xb * tp_ratio  # full remote x section in bytes
+        br = bb * tp_ratio  # full remote B section in bytes
+        return [
+            (local_rank_offset * xb, xb),
+            (xr + local_rank_offset * bb, bb),
+            (xr + br + local_rank_offset * bb, bb),
+        ]
+
+
+def derive_mamba_conv_split(
+    mamba_spec: MambaSpec,
+    local_tp: int,
+) -> MambaConvSplitInfo:
+    """Derive per-rank x/B/C byte sizes from a MambaSpec.
+
+    Called once at init on both P and D.  Decomposes the conv dimension
+    (= intermediate_size + 2 * groups_ss) into its x, B, C parts.
+
+    Args:
+        mamba_spec: MambaSpec whose shapes are:
+            shapes[0] = conv state: (conv_dim_local, conv_rows) in DS layout.
+            shapes[1] = SSM temporal: (local_num_heads, head_dim).
+        local_tp: this engine's tensor-parallel size.
+
+    Returns:
+        MambaConvSplitInfo with per-rank x_local, b_local, conv_rows, and
+        conv_dtype_size.
+    """
+    if mamba_spec.mamba_type != "mamba2":
+        raise NotImplementedError(
+            f"3-read conv transfer only supports Mamba2 models, "
+            f"got mamba_type={mamba_spec.mamba_type!r}.  "
+            f"Mamba1 SSM temporal shape is (intermediate_size // tp, state_size) "
+            f"which cannot be used to reconstruct intermediate_size."
+        )
+
+    conv_shape = mamba_spec.shapes[0]
+    assert len(conv_shape) == 2, f"Expected 2D conv state shape, got {conv_shape}"
+
+    # NOTE (ZhanqiuHu): 3-read requires DS layout, which is already asserted
+    # in nixl_connector __init__.  Use it directly instead of heuristic detection.
+    assert is_conv_state_dim_first(), "3-read requires DS conv state layout"
+    local_conv_dim = conv_shape[0]  # DS: (conv_dim_local, conv_rows)
+    conv_rows = conv_shape[1]
+
+    # NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored
+    # in MambaSpec, so we reconstruct it from the SSM temporal state shape:
+    #   shapes[1] = (local_num_heads, head_dim), already divided by TP.
+    head_dim = mamba_spec.shapes[1][1]
+    local_num_heads = mamba_spec.shapes[1][0]
+    intermediate_size = local_num_heads * local_tp * head_dim
+
+    # NOTE (ZhanqiuHu): global conv dim = intermediate_size + 2 * groups_ss,
+    # where groups_ss is the B (= C) dimension.  B and C are always the same
+    # size, so we recover groups_ss from the remainder after subtracting x.
+    remainder = local_conv_dim * local_tp - intermediate_size
+    assert remainder > 0 and remainder % 2 == 0, (
+        f"Conv dim ({local_conv_dim}*tp={local_tp}) doesn't decompose into "
+        f"intermediate_size={intermediate_size} + 2*groups_ss. "
+        f"remainder={remainder}"
+    )
+    groups_ss = remainder // 2
+
+    conv_dtype_size = torch.tensor(
+        [],
+        dtype=mamba_spec.dtypes[0],  # type: ignore[misc]
+    ).element_size()
+
+    # Divide by TP to get per-rank column counts.
+    return MambaConvSplitInfo(
+        conv_rows=conv_rows,
+        x_local=intermediate_size // local_tp,
+        b_local=groups_ss // local_tp,
+        conv_dtype_size=conv_dtype_size,
+    )
+
+
+def compute_mamba_phys_ratio(ssm_sizes: tuple[int, ...], block_len: int) -> int:
+    """Derive _physical_blocks_per_logical_kv_block from remote metadata.
+
+    The remote engine's ratio is not sent directly in the handshake, so we
+    reconstruct it: total mamba state per logical block / block_len.
+
+    Args:
+        ssm_sizes: (conv_state_bytes, ssm_state_bytes) from NixlAgentMetadata.
+        block_len: the engine's block_len in bytes (from block_lens[0]).
+    """
+    return math.ceil((ssm_sizes[0] + ssm_sizes[1]) / block_len)

vllm/env_override.py

@@ -500,7 +500,7 @@ if is_torch_equal("2.9.0"):
 # This mirrors the fix in https://github.com/pytorch/pytorch/pull/177558
 # and can be removed once torch >=2.12 is the minimum supported version.
 
-if not is_torch_equal_or_newer("2.12.0"):
+if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0"):
     import builtins as _builtins
     import pickle
 

vllm/envs.py

@@ -1464,6 +1464,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # - "flashinfer-trtllm": use flashinfer trtllm GEMM backend
     # - "flashinfer-cutlass": use flashinfer cutlass GEMM backend
     # - "marlin": use marlin GEMM backend (for GPUs without native FP4 support)
+    # - "emulation":
+    #     use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations.
+    #     This is only meant for research purposes to run on devices where NVFP4
+    #     GEMM kernels are not available.
     # - <none>: automatically pick an available backend
     "VLLM_NVFP4_GEMM_BACKEND": env_with_choices(
         "VLLM_NVFP4_GEMM_BACKEND",
@@ -1474,6 +1478,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
             "flashinfer-cutlass",
             "cutlass",
             "marlin",
+            "emulation",
         ],
     ),
     # Controls garbage collection during CUDA graph capture.

vllm/ir/ops/layernorm.py

@@ -16,7 +16,6 @@ def rms_norm(
     x_var = x if variance_size is None else x[..., :variance_size]
     variance = x_var.pow(2).mean(dim=-1, keepdim=True)
     x = x * torch.rsqrt(variance + epsilon)
-    x = x.to(orig_dtype)
     if weight is not None:
-        x = x * weight
-    return x
+        x = x.to(weight.dtype) * weight
+    return x.to(orig_dtype)

vllm/kernels/aiter_ops.py

@@ -36,13 +36,11 @@ AITER_SUPPORTED = is_aiter_found()
 
 rms_no_var_16bit_only = (
     lambda x, weight, epsilon, variance_size=None: variance_size is None
-    and x.dtype
-    in (
-        torch.float16,
-        torch.bfloat16,
-    )
+    and x.dtype in (torch.float16, torch.bfloat16)
+    and (weight is None or weight.dtype == x.dtype)
 )
-"""AITER rms_norm only supports float16 and bfloat16 acts and no var_size override."""
+"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
+and requires weight dtype to match x dtype."""
 
 
 @ir.ops.rms_norm.register_impl(

vllm/kernels/vllm_c.py

@@ -11,8 +11,11 @@ current_platform.import_kernels()
 CUDA_ALIKE = current_platform.is_cuda_alike()
 """Most kernels in this file are supported on all CUDA-alike platforms."""
 
-rms_no_var_size = lambda x, weight, epsilon, variance_size=None: variance_size is None
-"""vLLM kernel does not support variance_size parameter."""
+rms_no_var_size = (
+    lambda x, weight, epsilon, variance_size=None: variance_size is None
+    and (weight is None or weight.dtype == x.dtype)
+)
+"""vLLM kernel requires no variance_size override and matching input/weight dtype."""
 
 
 @ir.ops.rms_norm.register_impl(

vllm/kernels/xpu_ops.py

@@ -18,7 +18,9 @@ def is_xpu_kernels_found() -> bool:
 XPU_KERNELS_SUPPORTED = is_xpu_kernels_found()
 """Kernels in this file are supported if vLLM XPU kernels are installed."""
 
-rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
+rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None and (
+    weight is None or weight.dtype == x.dtype
+)
 
 
 @ir.ops.rms_norm.register_impl(

vllm/model_executor/kernels/linear/scaled_mm/triton.py

@@ -31,8 +31,6 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
 
     @classmethod
     def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
-        if not c.input_symmetric:
-            return False, "supports symmetric input only."
         return True, None
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -62,17 +60,59 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
         # INPUT SCALE
         if self.config.is_static_input_scheme:
             assert i_s is not None
-            replace_parameter(
-                layer,
-                i_s_name,
-                torch.nn.Parameter(i_s.max(), requires_grad=False),
-            )
-            setattr(layer, i_zp_name, None)
+
+            if self.config.input_symmetric:
+                replace_parameter(
+                    layer,
+                    i_s_name,
+                    torch.nn.Parameter(i_s.max(), requires_grad=False),
+                )
+                setattr(layer, i_zp_name, None)
+            else:
+                input_zero_point = getattr(layer, i_zp_name)
+
+                # Reconstruct the ranges to find a single scale and azp
+                int8_traits = torch.iinfo(torch.int8)
+                azps = input_zero_point.to(dtype=torch.int32)
+                range_max = (i_s * (int8_traits.max - azps)).max()
+                range_min = (i_s * (int8_traits.min - azps)).min()
+
+                scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
+                replace_parameter(
+                    layer,
+                    i_s_name,
+                    torch.nn.Parameter(scale, requires_grad=False),
+                )
+
+                # AZP loaded as int8 but used as int32
+                azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
+                replace_parameter(
+                    layer,
+                    i_zp_name,
+                    torch.nn.Parameter(azp, requires_grad=False),
+                )
         else:
             setattr(layer, i_s_name, None)
             setattr(layer, i_zp_name, None)
 
-        setattr(layer, azp_adj_name, None)
+        # azp_adj is the AZP adjustment term, used to account for weights.
+        # It does not depend on scales or azp, so it is the same for
+        # static and dynamic quantization.
+        # See csrc/quantization/w8a8/cutlass/Epilogues.md for the math.
+        if not self.config.input_symmetric:
+            weight = getattr(layer, w_q_name)
+            # weight is already transposed to [K, N], sum over K (dim=0)
+            azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
+            if self.config.is_static_input_scheme:
+                # Fold azp into azp_adj for the per-tensor case
+                azp_adj = getattr(layer, i_zp_name) * azp_adj
+            setattr(
+                layer,
+                azp_adj_name,
+                torch.nn.Parameter(azp_adj, requires_grad=False),
+            )
+        else:
+            setattr(layer, azp_adj_name, None)
 
     def apply_weights(
         self,
@@ -80,14 +120,33 @@ class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        w_q, w_s, i_s, i_zp, _ = self._get_layer_params(layer)
+        w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
 
+        symmetric = azp_adj is None
         x_q, x_s, x_zp = ops.scaled_int8_quant(
-            x.contiguous(), i_s, i_zp, symmetric=True
+            x.contiguous(), i_s, i_zp, symmetric=symmetric
         )
 
-        assert x_zp is None, "Triton kernel only supports symmetric quantization"
-
-        return triton_scaled_mm(
+        out = triton_scaled_mm(
             x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
         )
+
+        if azp_adj is not None:
+            # Asymmetric quantization: subtract the zero-point correction.
+            # D = scale_a * scale_b * (A_q @ B_q - azp * azp_adj) + bias
+            # triton_scaled_mm already computed scale_a * scale_b * (A_q @ B_q) + bias
+            # so we subtract scale_a * scale_b * azp * azp_adj
+            #
+            # x_s: [M, 1] or scalar, w_s: [N, 1] or scalar, azp_adj: [1, N]
+            # Reshape w_s from [N, 1] to [1, N] for proper broadcasting.
+            w_s_row = w_s.view(1, -1) if w_s.dim() > 0 else w_s
+            static = i_zp is not None
+            if not static and x_zp is not None:
+                # Dynamic per-token: azp is per-token, azp_adj is per-channel
+                # x_zp: [M, 1], azp_adj: [1, N]
+                out -= x_s * w_s_row * (x_zp * azp_adj).to(x.dtype)
+            else:
+                # Static per-tensor: azp already folded into azp_adj
+                out -= (x_s * w_s_row * azp_adj).to(x.dtype)
+
+        return out

vllm/model_executor/layers/attention/attention.py

@@ -131,9 +131,6 @@ def _init_kv_cache_quant(
         quant_config: Optional quantization configuration.
         prefix: Layer name prefix for quantization method lookup.
     """
-    quant_method = (
-        quant_config.get_quant_method(layer, prefix=prefix) if quant_config else None
-    )
 
     # Note [Register q/k/v/prob scales in state dict]
     # When calling model.to(device), only parameters/buffers in state dict are

vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_batched_moe.py

@@ -0,0 +1,353 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm import envs
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
+    TopKWeightAndReduceDelegate,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kNvfp4Dynamic,
+    kNvfp4Static,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import (
+    flashinfer_cutedsl_grouped_gemm_nt_masked,
+    has_flashinfer_cutedsl_grouped_gemm_nt_masked,
+    scaled_fp4_grouped_quantize,
+    silu_and_mul_scaled_nvfp4_experts_quantize,
+)
+
+logger = init_logger(__name__)
+
+
+class FlashInferCuteDSLBatchedExperts(mk.FusedMoEExpertsModular):
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+        max_num_tokens: int,
+        num_dispatchers: int,
+    ):
+        super().__init__(
+            moe_config=moe_config,
+            quant_config=quant_config,
+            max_num_tokens=max_num_tokens,
+            num_dispatchers=num_dispatchers,
+        )
+        assert quant_config.quant_dtype == "nvfp4", (
+            "Only nvfp4 quantization are currently supported."
+        )
+        self.out_dtype = moe_config.in_dtype
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
+        layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale)
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.BatchedExperts
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        p = current_platform
+        return (
+            p.is_cuda()
+            and p.is_device_capability_family(100)
+            and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
+        )
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        SUPPORTED_W_A = [
+            (kNvfp4Static, kNvfp4Dynamic),
+        ]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation == MoEActivation.SILU
+
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return True
+
+    def supports_expert_map(self) -> bool:
+        return False
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        # Let PrepareAndFinalize::finalize() decide the impl.
+        return TopKWeightAndReduceDelegate()
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        """
+        Compute the shapes for the temporary and final outputs of the two gemms
+        and activation in the fused expert function.  Since the gemms are
+        independent, the workspace for the first gemm can be shared with the
+        workspace for the last gemm.
+
+        Returns a tuple of:
+        - workspace13 shape tuple: must be large enough to hold the
+          result of either expert gemm.
+        - workspace2 shape tuple: must be large enough to hold the
+          result of the activation function.
+        - output shape tuple: must be exact size of the final gemm output.
+        - Workspace type: The dtype to use for the workspace tensors.
+        - Note: in order for activation chunking to work, the first dimension
+          of each tuple must be the number of tokens.
+        """
+
+        # We use global_num_experts due to how moe_align_block_size handles
+        # expert_maps.
+        K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
+        output_shape = (local_num_experts, M, K_dim)
+        workspace2 = (local_num_experts, M, N)
+        workspace1 = output_shape
+        return (workspace1, workspace2, output_shape)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,  # Not used
+        workspace13: torch.Tensor | None,
+        workspace2: torch.Tensor | None,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool | None,
+    ):
+        assert self.quant_dtype == "nvfp4", (
+            "Only nvfp4 quantization are currently supported."
+        )
+        # Ensure w1_scale and w2_scale are not None before calling view
+        assert self.w1_scale is not None and self.w2_scale is not None, (
+            "w1_scale and w2_scale must not be None for FlashInferExperts"
+        )
+        assert expert_tokens_meta is not None
+        expert_num_tokens = expert_tokens_meta.expert_num_tokens
+        assert hidden_states.ndim == 3
+        assert self.w1_scale.ndim == 3
+        assert self.w2_scale.ndim == 3
+
+        input_global_scale = (
+            None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
+        )
+        flashinfer_hidden_states = (
+            (hidden_states, a1q_scale)
+            if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
+            else hidden_states
+        )
+        flashinfer_cutedsl_moe_masked(
+            hidden_states=flashinfer_hidden_states,
+            input_global_scale=input_global_scale,
+            w1=w1,
+            w1_blockscale=self.w1_scale,
+            w1_alpha=self.g1_alphas,
+            w2=w2,
+            a2_global_scale=self.a2_gscale,
+            w2_blockscale=self.w2_scale,
+            w2_alpha=self.g2_alphas,
+            masked_m=expert_num_tokens,
+            workspace=workspace2,
+            out=output,
+        )
+
+
+def get_cute_dtype(input: torch.Tensor) -> str:
+    if input.dtype == torch.bfloat16:
+        return "bfloat16"
+    elif input.dtype == torch.float16:
+        return "float16"
+    elif input.dtype == torch.float32:
+        return "float32"
+    else:
+        raise ValueError(f"Unsupported cute dtype {input.dtype}")
+
+
+def flashinfer_cutedsl_moe_masked(
+    hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    input_global_scale: torch.Tensor,
+    w1: torch.Tensor,
+    w1_blockscale: torch.Tensor,
+    w1_alpha,
+    w2: torch.Tensor,
+    a2_global_scale: torch.Tensor,
+    w2_blockscale: torch.Tensor,
+    w2_alpha,
+    masked_m: torch.Tensor,
+    workspace: torch.Tensor,
+    out: torch.Tensor,
+):
+    """
+    Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
+    kernels.
+
+    Args:
+        hidden_states: Either of the following case
+            * torch.Tensor: [num_experts, m, k], bf16
+            * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2],
+                  uint8, [num_experts, m, k // 16], float8_e4m3fn
+        input_global_scale (torch.Tensor): (l,)
+        w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8
+        w1_blockscale (torch.Tensor): blockscale factors, e4m3,
+        w1_alpha (torch.Tensor): (l,)
+        w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8
+        a2_global_scale (torch.Tensor): (l,)
+        w2_blockscale (torch.Tensor): blockscale factors, e4m3,
+        w2_alpha (torch.Tensor): (l,)
+        masked_m (torch.Tensor): Masked dimension indices
+        workspace (torch.Tensor): For gateup_output
+
+    Notes:
+        - Assumes max(masked_m) <= m.
+    """
+
+    # === Assertions on dtypes ===
+    assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
+    assert w1_blockscale.dtype == torch.float8_e4m3fn, (
+        f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}"
+    )
+    assert w1_alpha.dtype == torch.float32, (
+        f"w1_alpha must be float32, got {w1_alpha.dtype}"
+    )
+    assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}"
+    assert a2_global_scale.dtype == torch.float32, (
+        f"a2_global_scale must be float32, got {a2_global_scale.dtype}"
+    )
+    assert w2_blockscale.dtype == torch.float8_e4m3fn, (
+        f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}"
+    )
+    assert w2_alpha.dtype == torch.float32, (
+        f"w2_alpha must be float32, got {w2_alpha.dtype}"
+    )
+
+    # === Assertions on shapes ===
+    n = w2.shape[-1] * 2  # intermediate dimension
+    if isinstance(hidden_states, tuple):
+        assert input_global_scale is None, (
+            "input_global_scale is needed when input needs quant"
+        )
+
+        aq = hidden_states[0].view(torch.uint8)
+        aq_sf = hidden_states[1].view(torch.float8_e4m3fn)
+        # m, k_by_2, num_experts = aq.shape
+        num_experts, m, k_by_2 = aq.shape
+        k = k_by_2 * 2
+        aq = aq.permute(1, 2, 0)
+    else:
+        num_experts, m, k = hidden_states.shape
+
+        assert input_global_scale.dtype == torch.float32, (
+            f"input_global_scale must be float32, got {input_global_scale.dtype}"
+        )
+        assert input_global_scale.shape == (num_experts,), (
+            f"input_global_scale must be (l,), got {input_global_scale.shape}"
+        )
+
+        aq, aq_sf = scaled_fp4_grouped_quantize(
+            hidden_states,
+            masked_m,
+            input_global_scale,
+        )
+
+    assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
+    assert w1.shape[-1] * 2 == k, (
+        f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}"
+    )
+    assert w2.shape[-2:] == (
+        k,
+        n // 2,
+    ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}"
+
+    assert w1_alpha.shape == (num_experts,), (
+        f"w1_alpha must be (l,), got {w1_alpha.shape}"
+    )
+    assert a2_global_scale.shape == (num_experts,), (
+        f"a2_global_scale must be (l,), got {a2_global_scale.shape}"
+    )
+    assert w2_alpha.shape == (num_experts,), (
+        f"w2_alpha must be (l,), got {w2_alpha.shape}"
+    )
+
+    workspace = workspace.permute(1, 2, 0)  # requirement of kernel
+    sf_vec_size = 16
+    assert aq_sf.dtype == torch.float8_e4m3fn
+    assert aq.dtype == torch.uint8
+    ab_dtype = "float4_e2m1fn"
+    sf_dtype = "float8_e4m3fn"
+
+    if isinstance(hidden_states, tuple):
+        c_dtype = "bfloat16"
+    else:
+        c_dtype = get_cute_dtype(hidden_states)
+
+    # Gemm1
+    flashinfer_cutedsl_grouped_gemm_nt_masked(
+        (aq, aq_sf),
+        (w1.permute(1, 2, 0), w1_blockscale),
+        workspace,
+        masked_m,
+        ab_dtype=ab_dtype,
+        sf_dtype=sf_dtype,
+        c_dtype=c_dtype,
+        sf_vec_size=sf_vec_size,
+        alpha=w1_alpha.view(1, 1, num_experts),
+        alpha_dtype=get_cute_dtype(w1_alpha),
+    )  # in logical [m, n, l]
+
+    # SILU and quantization
+    diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
+        workspace.permute(2, 0, 1),
+        masked_m,
+        a2_global_scale,
+    )
+
+    # Gemm2
+    out = out.permute(1, 2, 0)  # requirement of kernel
+    flashinfer_cutedsl_grouped_gemm_nt_masked(
+        (diq, diq_sf),
+        (w2.permute(1, 2, 0), w2_blockscale),
+        out,
+        masked_m,
+        ab_dtype=ab_dtype,
+        sf_dtype=sf_dtype,
+        c_dtype=c_dtype,
+        sf_vec_size=sf_vec_size,
+        alpha=w2_alpha.view(1, 1, num_experts),
+        alpha_dtype=get_cute_dtype(w2_alpha),
+    )  # in logical [m, k, l]
+    out = out.permute(2, 0, 1)

vllm/model_executor/layers/fused_moe/experts/flashinfer_cutedsl_moe.py

@@ -4,8 +4,6 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
-from vllm import envs
-from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
@@ -13,7 +11,7 @@ from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
-    TopKWeightAndReduceDelegate,
+    TopKWeightAndReduceNoOP,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     QuantKey,
@@ -22,33 +20,42 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import (
-    flashinfer_cutedsl_grouped_gemm_nt_masked,
-    has_flashinfer_cutedsl_grouped_gemm_nt_masked,
-    scaled_fp4_grouped_quantize,
-    silu_and_mul_scaled_nvfp4_experts_quantize,
+    flashinfer_cute_dsl_fused_moe_nvfp4,
+    has_flashinfer_cutedsl_moe_nvfp4,
 )
 
-logger = init_logger(__name__)
-
 
 class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
+    """
+    CuteDSL NvFP4 MoE experts using the FlashInfer functional API.
+
+    Uses Standard activation format (non-batched). The kernel handles
+    routing, expert computation, and reduction internally.
+    Supports expert parallelism natively.
+    """
+
     def __init__(
         self,
         moe_config: FusedMoEConfig,
         quant_config: FusedMoEQuantConfig,
-        max_num_tokens: int,
-        num_dispatchers: int,
     ):
         super().__init__(
             moe_config=moe_config,
             quant_config=quant_config,
-            max_num_tokens=max_num_tokens,
-            num_dispatchers=num_dispatchers,
         )
         assert quant_config.quant_dtype == "nvfp4", (
-            "Only nvfp4 quantization are currently supported."
+            "Only nvfp4 quantization is currently supported."
         )
         self.out_dtype = moe_config.in_dtype
+        self.hidden_dim = moe_config.hidden_dim
+        self.intermediate_size_per_partition = (
+            moe_config.intermediate_size_per_partition
+        )
+        self.topk = moe_config.experts_per_token
+        self.local_num_experts = moe_config.num_local_experts
+        self.global_num_experts = moe_config.num_experts
+        self.ep_rank = moe_config.moe_parallel_config.ep_rank
+        self.local_expert_offset = self.ep_rank * self.local_num_experts
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale)
@@ -56,7 +63,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
-        return mk.FusedMoEActivationFormat.BatchedExperts
+        return mk.FusedMoEActivationFormat.Standard
 
     @staticmethod
     def _supports_current_device() -> bool:
@@ -64,7 +71,7 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         return (
             p.is_cuda()
             and p.is_device_capability_family(100)
-            and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
+            and has_flashinfer_cutedsl_moe_nvfp4()
         )
 
     @staticmethod
@@ -86,15 +93,16 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         return activation == MoEActivation.SILU
 
     @staticmethod
-    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
         return True
 
     def supports_expert_map(self) -> bool:
         return False
 
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        # Let PrepareAndFinalize::finalize() decide the impl.
-        return TopKWeightAndReduceDelegate()
+        return TopKWeightAndReduceNoOP()
 
     def workspace_shapes(
         self,
@@ -107,29 +115,12 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        # We use global_num_experts due to how moe_align_block_size handles
-        # expert_maps.
-        """
-        Compute the shapes for the temporary and final outputs of the two gemms
-        and activation in the fused expert function.  Since the gemms are
-        independent, the workspace for the first gemm can be shared with the
-        workspace for the last gemm.
-
-        Returns a tuple of:
-        - workspace13 shape tuple: must be large enough to hold the
-          result of either expert gemm.
-        - workspace2 shape tuple: must be large enough to hold the
-          result of the activation function.
-        - output shape tuple: must be exact size of the final gemm output.
-        - Workspace type: The dtype to use for the workspace tensors.
-        - Note: in order for activation chunking to work, the first dimension
-          of each tuple must be the number of tokens.
-        """
-        K_dim = K * 2 if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else K
-        output_shape = (local_num_experts, M, K_dim)
-        workspace2 = (local_num_experts, M, N)
-        workspace1 = output_shape
-        return (workspace1, workspace2, output_shape)
+        workspace1 = (0,)
+        workspace2 = (0,)
+        # K is packed (K//2 for uint8), so output uses hidden_dim.
+        assert self.hidden_dim == K * 2
+        output = (M, self.hidden_dim)
+        return (workspace1, workspace2, output)
 
     def apply(
         self,
@@ -143,210 +134,39 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
         global_num_experts: int,
         expert_map: torch.Tensor | None,
         a1q_scale: torch.Tensor | None,
-        a2_scale: torch.Tensor | None,  # Not used
+        a2_scale: torch.Tensor | None,
         workspace13: torch.Tensor | None,
         workspace2: torch.Tensor | None,
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool | None,
     ):
-        assert self.quant_dtype == "nvfp4", (
-            "Only nvfp4 quantization are currently supported."
-        )
-        # Ensure w1_scale and w2_scale are not None before calling view
-        assert self.w1_scale is not None and self.w2_scale is not None, (
-            "w1_scale and w2_scale must not be None for FlashInferExperts"
-        )
-        assert expert_tokens_meta is not None
-        expert_num_tokens = expert_tokens_meta.expert_num_tokens
-        assert hidden_states.ndim == 3
-        assert self.w1_scale.ndim == 3
-        assert self.w2_scale.ndim == 3
+        assert self.quant_dtype == "nvfp4"
+        assert a1q_scale is not None
+        assert self.w1_scale is not None
+        assert self.w2_scale is not None
 
-        input_global_scale = (
-            None if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH else self.a1_gscale
-        )
-        flashinfer_hidden_states = (
-            (hidden_states, a1q_scale)
-            if envs.VLLM_DEEPEPLL_NVFP4_DISPATCH
-            else hidden_states
-        )
-        flashinfer_cutedsl_moe_masked(
-            hidden_states=flashinfer_hidden_states,
-            input_global_scale=input_global_scale,
-            w1=w1,
-            w1_blockscale=self.w1_scale,
-            w1_alpha=self.g1_alphas,
-            w2=w2,
-            a2_global_scale=self.a2_gscale,
-            w2_blockscale=self.w2_scale,
-            w2_alpha=self.g2_alphas,
-            masked_m=expert_num_tokens,
-            workspace=workspace2,
-            out=output,
-        )
+        # a1q_scale is (M, K//16) float8_e4m3fn from fp4_quantize.
+        # The functional API expects x_sf with trailing dim: (M, K//16, 1).
+        x_sf = a1q_scale.unsqueeze(-1)
 
+        from vllm.utils.flashinfer import _is_fi_autotuning, autotune
 
-def get_cute_dtype(input: torch.Tensor) -> str:
-    if input.dtype == torch.bfloat16:
-        return "bfloat16"
-    elif input.dtype == torch.float16:
-        return "float16"
-    elif input.dtype == torch.float32:
-        return "float32"
-    else:
-        raise ValueError(f"Unsupported cute dtype {input.dtype}")
-
-
-def flashinfer_cutedsl_moe_masked(
-    hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-    input_global_scale: torch.Tensor,
-    w1: torch.Tensor,
-    w1_blockscale: torch.Tensor,
-    w1_alpha,
-    w2: torch.Tensor,
-    a2_global_scale: torch.Tensor,
-    w2_blockscale: torch.Tensor,
-    w2_alpha,
-    masked_m: torch.Tensor,
-    workspace: torch.Tensor,
-    out: torch.Tensor,
-):
-    """
-    Perform masked Mixture-of-Experts computation with FlashInfer's CuteDSL
-    kernels.
-
-    Args:
-        hidden_states: Either of the following case
-            * torch.Tensor: [num_experts, m, k], bf16
-            * tuple[torch.Tensor, torch.Tensor]: [num_experts, m, k // 2],
-                  uint8, [num_experts, m, k // 16], float8_e4m3fn
-        input_global_scale (torch.Tensor): (l,)
-        w1 (torch.Tensor): fp4 weights, [l, 2 * n, k // 2], uint8
-        w1_blockscale (torch.Tensor): blockscale factors, e4m3,
-        w1_alpha (torch.Tensor): (l,)
-        w2 (torch.Tensor): fp4 weights, [l, k, n // 2], uint8
-        a2_global_scale (torch.Tensor): (l,)
-        w2_blockscale (torch.Tensor): blockscale factors, e4m3,
-        w2_alpha (torch.Tensor): (l,)
-        masked_m (torch.Tensor): Masked dimension indices
-        workspace (torch.Tensor): For gateup_output
-
-    Notes:
-        - Assumes max(masked_m) <= m.
-    """
-
-    # === Assertions on dtypes ===
-    assert w1.dtype == torch.uint8, f"w1 must be uint8, got {w1.dtype}"
-    assert w1_blockscale.dtype == torch.float8_e4m3fn, (
-        f"w1_blockscale must be float8_e4m3fn, got {w1_blockscale.dtype}"
-    )
-    assert w1_alpha.dtype == torch.float32, (
-        f"w1_alpha must be float32, got {w1_alpha.dtype}"
-    )
-    assert w2.dtype == torch.uint8, f"w2 must be uint8, got {w2.dtype}"
-    assert a2_global_scale.dtype == torch.float32, (
-        f"a2_global_scale must be float32, got {a2_global_scale.dtype}"
-    )
-    assert w2_blockscale.dtype == torch.float8_e4m3fn, (
-        f"w2_blockscale must be float8_e4m3fn, got {w2_blockscale.dtype}"
-    )
-    assert w2_alpha.dtype == torch.float32, (
-        f"w2_alpha must be float32, got {w2_alpha.dtype}"
-    )
-
-    # === Assertions on shapes ===
-    n = w2.shape[-1] * 2  # intermediate dimension
-    if isinstance(hidden_states, tuple):
-        assert input_global_scale is None, (
-            "input_global_scale is needed when input needs quant"
-        )
-
-        aq = hidden_states[0].view(torch.uint8)
-        aq_sf = hidden_states[1].view(torch.float8_e4m3fn)
-        # m, k_by_2, num_experts = aq.shape
-        num_experts, m, k_by_2 = aq.shape
-        k = k_by_2 * 2
-        aq = aq.permute(1, 2, 0)
-    else:
-        num_experts, m, k = hidden_states.shape
-
-        assert input_global_scale.dtype == torch.float32, (
-            f"input_global_scale must be float32, got {input_global_scale.dtype}"
-        )
-        assert input_global_scale.shape == (num_experts,), (
-            f"input_global_scale must be (l,), got {input_global_scale.shape}"
-        )
-
-        aq, aq_sf = scaled_fp4_grouped_quantize(
-            hidden_states,
-            masked_m,
-            input_global_scale,
-        )
-
-    assert w1.shape[-2] == 2 * n, f"w1 last-2 dim must be 2*n, got {w1.shape}"
-    assert w1.shape[-1] * 2 == k, (
-        f"w1 last dim * 2 must equal k, got {w1.shape[-1]} vs k={k}"
-    )
-    assert w2.shape[-2:] == (
-        k,
-        n // 2,
-    ), f"w2 shape mismatch, got {w2.shape[-2:]}, expected {(k, n // 2)}"
-
-    assert w1_alpha.shape == (num_experts,), (
-        f"w1_alpha must be (l,), got {w1_alpha.shape}"
-    )
-    assert a2_global_scale.shape == (num_experts,), (
-        f"a2_global_scale must be (l,), got {a2_global_scale.shape}"
-    )
-    assert w2_alpha.shape == (num_experts,), (
-        f"w2_alpha must be (l,), got {w2_alpha.shape}"
-    )
-
-    workspace = workspace.permute(1, 2, 0)  # requirement of kernel
-    sf_vec_size = 16
-    assert aq_sf.dtype == torch.float8_e4m3fn
-    assert aq.dtype == torch.uint8
-    ab_dtype = "float4_e2m1fn"
-    sf_dtype = "float8_e4m3fn"
-
-    if isinstance(hidden_states, tuple):
-        c_dtype = "bfloat16"
-    else:
-        c_dtype = get_cute_dtype(hidden_states)
-
-    # Gemm1
-    flashinfer_cutedsl_grouped_gemm_nt_masked(
-        (aq, aq_sf),
-        (w1.permute(1, 2, 0), w1_blockscale),
-        workspace,
-        masked_m,
-        ab_dtype=ab_dtype,
-        sf_dtype=sf_dtype,
-        c_dtype=c_dtype,
-        sf_vec_size=sf_vec_size,
-        alpha=w1_alpha.view(1, 1, num_experts),
-        alpha_dtype=get_cute_dtype(w1_alpha),
-    )  # in logical [m, n, l]
-
-    # SILU and quantization
-    diq, diq_sf = silu_and_mul_scaled_nvfp4_experts_quantize(
-        workspace.permute(2, 0, 1),
-        masked_m,
-        a2_global_scale,
-    )
-
-    # Gemm2
-    out = out.permute(1, 2, 0)  # requirement of kernel
-    flashinfer_cutedsl_grouped_gemm_nt_masked(
-        (diq, diq_sf),
-        (w2.permute(1, 2, 0), w2_blockscale),
-        out,
-        masked_m,
-        ab_dtype=ab_dtype,
-        sf_dtype=sf_dtype,
-        c_dtype=c_dtype,
-        sf_vec_size=sf_vec_size,
-        alpha=w2_alpha.view(1, 1, num_experts),
-        alpha_dtype=get_cute_dtype(w2_alpha),
-    )  # in logical [m, k, l]
-    out = out.permute(2, 0, 1)
+        with autotune(_is_fi_autotuning):
+            flashinfer_cute_dsl_fused_moe_nvfp4(
+                x=hidden_states,
+                x_sf=x_sf,
+                token_selected_experts=topk_ids.to(torch.int32),
+                token_final_scales=topk_weights.float(),
+                w1_weight=w1,
+                w1_weight_sf=self.w1_scale,
+                w1_alpha=self.g1_alphas,
+                fc2_input_scale=self.a2_gscale,
+                w2_weight=w2,
+                w2_weight_sf=self.w2_scale,
+                w2_alpha=self.g2_alphas,
+                num_experts=self.global_num_experts,
+                top_k=self.topk,
+                num_local_experts=self.local_num_experts,
+                local_expert_offset=self.local_expert_offset,
+                moe_output=output,
+            )

vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py

@@ -76,7 +76,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
         activation_key: QuantKey | None,
     ) -> bool:
         return routing_method in [
-            RoutingMethodType.Default,
             RoutingMethodType.DeepSeekV3,
             RoutingMethodType.Llama4,
             RoutingMethodType.Renormalize,

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

@@ -47,7 +47,6 @@ if has_triton_kernels():
             BIT,
             Bitmatrix,
         )
-        from triton_kernels.topk import topk
 
         try:
             from triton_kernels.tensor import (
@@ -89,6 +88,7 @@ def pack_bitmatrix(
     offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :]
     mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :]
     indices = tl.load(topk_ids + offsets, mask=mask, other=-1)
+    valid = indices >= 0
     div = indices // 32
     rem = indices % 32
     one = tl.cast(1, tl.uint32)
@@ -99,8 +99,13 @@ def pack_bitmatrix(
         offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32)
         # All topks that need to go into this column has the correct bit set.
         # Other bits are 0. x is a 2D tensor.
+        # Guard with `valid` to prevent negative indices from producing
+        # spurious bits (on HIP, -1 // 32 == 0 and 1 << (-1 % 32) sets
+        # bit 31).
         x = tl.where(
-            div[:, :, None] == offs[None, None, :], (one << rem)[:, :, None], 0
+            valid[:, :, None] & (div[:, :, None] == offs[None, None, :]),
+            (one << rem)[:, :, None],
+            0,
         )
         # Reduce x to get a single int32_t bitpack.
         y = tl.reduce_or(x, axis=1)
@@ -108,93 +113,6 @@ def pack_bitmatrix(
         tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows)
 
 
-def legacy_routing_from_bitmatrix(
-    bitmatrix: "Bitmatrix",
-    expt_scal: torch.Tensor,
-    expt_indx: torch.Tensor,
-    n_expts_tot: int,
-    n_expts_act: int,
-) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
-    """
-    Replacement for the removed triton_kernels.routing.routing_from_bitmatrix.
-    Creates routing data from a bitmatrix representation.
-    """
-    if use_legacy_triton_kernels:
-        from triton_kernels.routing import routing_from_bitmatrix
-
-        return routing_from_bitmatrix(
-            bitmatrix, expt_scal, expt_indx, n_expts_tot, n_expts_act
-        )
-    sparse_logits = SparseMatrix(indx=expt_indx, vals=expt_scal, mask=bitmatrix)
-    dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
-    combine_indx = sparse_logits.mask_metadata.col_sorted_indx
-    ragged_batch_metadata = make_ragged_tensor_metadata(
-        sparse_logits.mask_metadata.col_sum,
-        dispatch_indx.shape[0],
-    )
-    gate_scal = sparse_logits.vals.flatten()[combine_indx]
-    routing_data = RoutingData(
-        gate_scal,
-        ragged_batch_metadata.block_sizes,
-        n_expts_tot,
-        n_expts_act,
-        ragged_batch_metadata,
-    )
-    gather_idx = GatherIndx(combine_indx, dispatch_indx)
-    scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
-    return routing_data, gather_idx, scatter_idx
-
-
-def legacy_routing_from_sparsematrix(
-    sparse_logits: "SparseMatrix",
-    n_expts_tot: int,
-    n_expts_act: int,
-) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
-    """
-    Creates routing data from a SparseMatrix representation.
-    """
-    dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
-    combine_indx = sparse_logits.mask_metadata.col_sorted_indx
-    ragged_batch_metadata = make_ragged_tensor_metadata(
-        sparse_logits.mask_metadata.col_sum,
-        dispatch_indx.shape[0],
-    )
-    gate_scal = sparse_logits.vals.flatten()[combine_indx]
-    routing_data = RoutingData(
-        gate_scal,
-        ragged_batch_metadata.block_sizes,
-        n_expts_tot,
-        n_expts_act,
-        ragged_batch_metadata,
-    )
-    gather_idx = GatherIndx(combine_indx, dispatch_indx)
-    scatter_idx = ScatterIndx(dispatch_indx, combine_indx)
-    return routing_data, gather_idx, scatter_idx
-
-
-def legacy_routing(
-    logits: torch.Tensor,
-    n_expts_act: int,
-    sm_first: bool = False,
-) -> tuple["RoutingData", "GatherIndx", "ScatterIndx"]:
-    """
-    Replacement for the removed triton_kernels.routing.routing function.
-    Computes routing data from gating logits.
-    """
-    if use_legacy_triton_kernels:
-        from triton_kernels.routing import routing
-
-        return routing(logits, n_expts_act, sm_first=sm_first)
-    if sm_first:
-        logits = torch.softmax(logits, dim=-1)
-    sparse_logits = topk(logits, n_expts_act, apply_softmax=not sm_first)
-    return legacy_routing_from_sparsematrix(
-        sparse_logits,
-        logits.shape[-1],
-        n_expts_act,
-    )
-
-
 def triton_kernel_moe_forward(
     hidden_states: torch.Tensor,
     w1,  # Tensor or triton_kernels.Tensor
@@ -241,26 +159,22 @@ def triton_kernel_moe_forward(
             unpadded_K_w2=unpadded_K_w2,
         )
 
-    if expert_map is not None:
-        # With expert parallelism, legacy_routing produces routing data
-        # using global expert IDs which don't correspond to local weight
-        # indices.  Split the routing into topk selection + expert_map
-        # remapping + local routing data construction (matching the
-        # approach used by OAITritonExperts.apply).
-        from triton_kernels.topk import topk as topk_fn
+    from triton_kernels.topk import topk as topk_fn
 
-        sm_first = not renormalize
-        logits = gating_output
-        if sm_first:
-            logits = torch.softmax(logits, dim=-1)
-        topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
-        # topk may return a tuple (vals, indx, bitmatrix) or a
-        # SparseMatrix depending on the triton_kernels version.
-        if isinstance(topk_result, tuple):
-            topk_weights, topk_ids_raw, _ = topk_result
-        else:
-            topk_weights = topk_result.vals
-            topk_ids_raw = topk_result.indx
+    sm_first = not renormalize
+    logits = gating_output
+    if sm_first:
+        logits = torch.softmax(logits, dim=-1)
+    topk_result = topk_fn(logits, topk, apply_softmax=not sm_first)
+    # topk may return a tuple (vals, indx, bitmatrix) or a
+    # SparseMatrix depending on the triton_kernels version.
+    if isinstance(topk_result, tuple):
+        topk_weights, topk_ids_raw, _ = topk_result
+    else:
+        topk_weights = topk_result.vals
+        topk_ids_raw = topk_result.indx
+
+    if expert_map is not None:
         # topk_ids_raw contains global expert IDs - remap to local.
         topk_ids = expert_map[topk_ids_raw.to(torch.long)]
         local_num_experts = w1.shape[0]
@@ -271,8 +185,9 @@ def triton_kernel_moe_forward(
         effective_expert_map = None
         effective_global_num_experts = local_num_experts
     else:
-        routing_data, gather_idx, scatter_idx = legacy_routing(
-            gating_output, topk, sm_first=not renormalize
+        topk_ids = topk_ids_raw.to(torch.long)
+        routing_data, gather_idx, scatter_idx = make_routing_data(
+            topk_ids, topk_weights, gating_output.shape[-1]
         )
         effective_expert_map = expert_map
         effective_global_num_experts = global_num_experts
@@ -539,10 +454,31 @@ def make_routing_data(
 
     # matmul_ogs expects invalid topk_weights to be -1s
     topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
-    routing_data, gather_indx, scatter_indx = legacy_routing_from_bitmatrix(
-        bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
-    )
 
+    if use_legacy_triton_kernels:
+        from triton_kernels.routing import routing_from_bitmatrix
+
+        return routing_from_bitmatrix(
+            bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk
+        )
+
+    sparse_logits = SparseMatrix(indx=topk_ids, vals=topk_weights, mask=bitmatrix)
+    dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
+    combine_indx = sparse_logits.mask_metadata.col_sorted_indx
+    ragged_batch_metadata = make_ragged_tensor_metadata(
+        sparse_logits.mask_metadata.col_sum,
+        dispatch_indx.shape[0],
+    )
+    gate_scal = sparse_logits.vals.flatten()[combine_indx]
+    routing_data = RoutingData(
+        gate_scal,
+        ragged_batch_metadata.block_sizes,
+        num_local_experts,
+        num_topk,
+        ragged_batch_metadata,
+    )
+    gather_indx = GatherIndx(combine_indx, dispatch_indx)
+    scatter_indx = ScatterIndx(dispatch_indx, combine_indx)
     return routing_data, gather_indx, scatter_indx
 
 

vllm/model_executor/layers/fused_moe/layer.py

@@ -39,8 +39,8 @@ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
 from vllm.model_executor.layers.fused_moe.router.router_factory import (
     create_fused_moe_router,
 )
-from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (
-    DefaultMoERunner,
+from vllm.model_executor.layers.fused_moe.runner.moe_runner_factory import (
+    create_moe_runner,
 )
 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
     SharedExperts,
@@ -572,8 +572,8 @@ class FusedMoE(CustomOp):
         # Storing the runner in the FusedMoE is an intermediate state, eventually
         # the runner will own the FusedMoE layer and provide the execution interface
         # for MoE ops.
-        self.runner = DefaultMoERunner(
-            layer=self,
+        self.runner = create_moe_runner(
+            layer_name=self.layer_name,
             moe_config=self.moe_config,
             router=self.router,
             routed_input_transform=self._routed_input_transform,

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

@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
 )
 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
     prepare_nvfp4_moe_layer_for_fi_or_cutlass,
+    prepare_nvfp4_moe_layer_for_flashinfer_cutedsl,
 )
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     FlashinferMoeBackend,
@@ -41,6 +42,7 @@ class NvFp4MoeBackend(Enum):
     FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
     FLASHINFER_CUTLASS = "FLASHINFER_CUTLASS"
     FLASHINFER_CUTEDSL = "FLASHINFER_CUTEDSL"
+    FLASHINFER_CUTEDSL_BATCHED = "FLASHINFER_CUTEDSL_BATCHED"
     VLLM_CUTLASS = "VLLM_CUTLASS"
     MARLIN = "MARLIN"
 
@@ -49,6 +51,7 @@ FLASHINFER_NVFP4_MOE_BACKENDS = [
     NvFp4MoeBackend.FLASHINFER_TRTLLM,
     NvFp4MoeBackend.FLASHINFER_CUTLASS,
     NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+    NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
 ]
 
 fi_2_vllm_backend_map: dict[FlashinferMoeBackend, NvFp4MoeBackend] = {
@@ -95,6 +98,13 @@ def backend_to_kernel_cls(
 
         return [FlashInferCuteDSLExperts]
 
+    elif backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED:
+        from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_batched_moe import (  # noqa: E501
+            FlashInferCuteDSLBatchedExperts,
+        )
+
+        return [FlashInferCuteDSLBatchedExperts]
+
     elif backend == NvFp4MoeBackend.VLLM_CUTLASS:
         from vllm.model_executor.layers.fused_moe.cutlass_moe import (
             CutlassExpertsFp4,
@@ -143,6 +153,7 @@ def select_nvfp4_moe_backend(
     AVAILABLE_BACKENDS = [
         NvFp4MoeBackend.FLASHINFER_TRTLLM,
         NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+        NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
         NvFp4MoeBackend.FLASHINFER_CUTLASS,
         NvFp4MoeBackend.VLLM_CUTLASS,
         NvFp4MoeBackend.MARLIN,
@@ -198,6 +209,12 @@ def select_nvfp4_moe_backend(
     runner_backend = config.moe_backend
     if runner_backend != "auto":
         requested_backend = map_nvfp4_backend(runner_backend)
+        # For batched activation format, use batched variant if available.
+        if (
+            activation_format == mk.FusedMoEActivationFormat.BatchedExperts
+            and requested_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL
+        ):
+            requested_backend = NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED
         return _return_or_raise(
             requested_backend, config, weight_key, activation_key, activation_format
         )
@@ -288,7 +305,28 @@ def convert_to_nvfp4_moe_kernel_format(
     torch.Tensor,
     torch.Tensor,
 ]:
-    if (
+    if nvfp4_backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL:
+        (
+            w13,
+            w13_scale,
+            w13_scale_2,
+            a13_scale,
+            w2,
+            w2_scale,
+            w2_scale_2,
+            a2_scale,
+        ) = prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
+            layer=layer,
+            w13=w13,
+            w13_scale=w13_scale,
+            w13_scale_2=w13_scale_2,
+            a13_scale=a13_scale,
+            w2=w2,
+            w2_scale=w2_scale,
+            w2_scale_2=w2_scale_2,
+            a2_scale=a2_scale,
+        )
+    elif (
         nvfp4_backend in FLASHINFER_NVFP4_MOE_BACKENDS
         or nvfp4_backend == NvFp4MoeBackend.VLLM_CUTLASS
     ):
@@ -380,7 +418,13 @@ def make_nvfp4_moe_quant_config(
         # NOTE(rob): this is a hack until the MoE kernels
         # create their own quant configs. TRTLLM kernel
         # does not accept swizzled input quant scales.
-        is_nvfp4_scale_swizzled=(backend != NvFp4MoeBackend.FLASHINFER_TRTLLM),
+        is_nvfp4_scale_swizzled=(
+            backend
+            not in (
+                NvFp4MoeBackend.FLASHINFER_TRTLLM,
+                NvFp4MoeBackend.FLASHINFER_CUTEDSL,
+            )
+        ),
     )
 
 

vllm/model_executor/layers/fused_moe/runner/chunking_moe_runner.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.forward_context import (
+    get_forward_context,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+)
+from vllm.utils.math_utils import cdiv
+from vllm.v1.worker.ubatching import dbo_current_ubatch_id
+from vllm.v1.worker.workspace import current_workspace_manager
+
+
+class ChunkingMoERunner(MoERunnerBase):
+    """
+    MoE runner wrapper that adds chunked processing to any MoERunnerBase.
+
+    This runner wraps an inner MoERunnerBase and overrides _forward_impl to
+    process large batches by breaking them into smaller chunks. Each chunk
+    is delegated to the inner runner's _forward_impl, making chunking
+    composable with any runner implementation.
+
+    All MoERunnerBase state (moe_config, router, quant_method, etc.) is
+    transparently delegated to the inner runner via __getattr__.
+    ChunkingMoERunner only owns chunking-specific state: the pre-allocated
+    workspace buffers and the reduce_results override.
+
+    Key behaviors:
+    - Pre-allocates workspace tensors for CUDA graph compatibility
+    - Processes chunks via inner._forward_impl per chunk
+    - Never reduces results (reduce_results always returns False)
+    """
+
+    def __init__(self, inner: MoERunnerBase):
+        # Assert that _maybe_dispatch/_maybe_combine will be nops.
+        assert inner.moe_config.pcp_size == 1
+
+        # Skip MoERunnerBase.__init__ — all state is delegated to inner
+        # via __getattr__. Only chunking-specific state lives here.
+        self._inner = inner
+
+        # Pre-allocated staging buffers. These need to exist ahead of time
+        # due to CUDA graph construction needing fixed buffer addresses.
+        self.batched_hidden_states, self.batched_router_logits = (
+            self._init_dp_chunking()
+        )
+
+    def __getattr__(self, name):
+        # Delegate attribute access to the inner runner. This is only
+        # called when normal lookup (instance __dict__, class MRO) fails,
+        # so ChunkingMoERunner's own attributes and methods take priority.
+        return getattr(self._inner, name)
+
+    @property
+    def shared_experts(self) -> SharedExperts | None:
+        return self._inner.shared_experts
+
+    # TODO(bnell): temporary hack, do not call this method.
+    def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
+        self._inner._replace_quant_method(quant_method)
+        self.quant_method = quant_method
+
+    def is_internal_router(self) -> bool:
+        return self._inner.gate is not None
+
+    # Reducing results when chunking is handled by the MK finalize operations
+    # when DP chunking is enabled..
+    # This will be removed by #35949
+    @property
+    def reduce_results(self) -> bool:
+        return False
+
+    def _init_dp_chunking(self) -> list[torch.Tensor]:
+        states_shape: tuple[int, ...]
+        logits_shape: tuple[int, ...]
+
+        moe = self.moe_config
+
+        if self.enable_dbo:
+            states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim)
+            logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts)
+        else:
+            states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)
+            logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)
+
+        # Does this need some kind of profiling run check like modular_kernel.py?
+        return current_workspace_manager().get_simultaneous(
+            (states_shape, moe.in_dtype),
+            (logits_shape, moe.router_logits_dtype),
+        )
+
+    def _allocate_dp_chunking_outputs(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor]:
+        # Assert the inputs are of the proper type and shape.
+        assert self.batched_hidden_states is not None
+        assert self.batched_router_logits is not None
+
+        assert self.batched_hidden_states.dtype == hidden_states.dtype, (
+            f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}"
+        )
+        assert self.batched_router_logits.dtype == router_logits.dtype, (
+            f"{self.batched_router_logits.dtype} == {router_logits.dtype}"
+        )
+
+        # Check size compatibility.
+        assert self.batched_hidden_states.size(-1) == hidden_states.size(-1)
+        assert self.batched_router_logits.size(-1) == router_logits.size(-1)
+
+        final_fused_hidden_states = torch.empty_like(hidden_states)
+        if self.shared_experts is not None:
+            if shared_experts_input is not None:
+                final_shared_hidden_states = torch.empty_like(shared_experts_input)
+            else:
+                final_shared_hidden_states = torch.empty_like(hidden_states)
+        else:
+            final_shared_hidden_states = None
+
+        return final_shared_hidden_states, final_fused_hidden_states
+
+    def _slice_and_copy_input(
+        self,
+        out_slice: torch.Tensor,
+        orig: torch.Tensor | None,
+        start: int,
+        end: int,
+    ) -> torch.Tensor:
+        assert orig is not None
+        slice_size = end - start
+        orig_slice = orig[start:end, :]
+        if self.enable_dbo:
+            assert out_slice.dim() == 3
+            batch_buffer_idx = dbo_current_ubatch_id()
+            out_slice = out_slice[batch_buffer_idx, :]
+
+        assert out_slice.size(0) >= slice_size
+        out_slice = out_slice[:slice_size, :]
+        out_slice.copy_(orig_slice, non_blocking=True)
+        return out_slice
+
+    def _forward_impl(
+        self,
+        layer: torch.nn.Module,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        final_shared_hidden_states, final_fused_hidden_states = (
+            self._allocate_dp_chunking_outputs(
+                hidden_states, router_logits, shared_experts_input
+            )
+        )
+
+        ctx = get_forward_context()
+        # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
+        max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu
+        moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens
+
+        # If the input to the MoE is sequence parallel then divide by sp_size
+        # to find the maximum number of tokens for any individual dispatcher.
+        if self.moe_config.is_sequence_parallel:
+            max_tokens_across_dispatchers = cdiv(
+                max_tokens_across_dispatchers, self.moe_config.sp_size
+            )
+
+        num_tokens = hidden_states.size(0)
+        for chunk_idx, chunk_start_ in enumerate(
+            range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)
+        ):
+            chunk_start = chunk_start_
+            chunk_end = min(
+                chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers
+            )
+            # clamp start and end
+            chunk_start = min(chunk_start, num_tokens - 1)
+            chunk_end = min(chunk_end, num_tokens)
+            chunk_sizes = ctx.dp_metadata.chunked_sizes(
+                self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx
+            )
+            with chunk_sizes:
+                hidden_states_chunk = self._slice_and_copy_input(
+                    self.batched_hidden_states,
+                    hidden_states,
+                    chunk_start,
+                    chunk_end,
+                )
+
+                router_logits_chunk = self._slice_and_copy_input(
+                    self.batched_router_logits,
+                    router_logits,
+                    chunk_start,
+                    chunk_end,
+                )
+
+                shared_experts_input_chunk = (
+                    shared_experts_input[chunk_start:chunk_end, :]
+                    if shared_experts_input is not None
+                    else None
+                )
+
+                # Delegate per-chunk computation to the inner runner.
+                chunk_result = self._inner._forward_impl(
+                    layer=layer,
+                    hidden_states=hidden_states_chunk,
+                    router_logits=router_logits_chunk,
+                    shared_experts_input=shared_experts_input_chunk,
+                )
+
+                # Store outputs
+                # TODO(bnell): document when chunk_start >= num_tokens
+                if chunk_start < num_tokens:
+                    if self.shared_experts is not None:
+                        assert isinstance(chunk_result, tuple)
+                        shared_output_chunk, hidden_states_chunk = chunk_result
+                        final_fused_hidden_states[chunk_start:chunk_end, :].copy_(
+                            hidden_states_chunk, non_blocking=True
+                        )
+                        assert shared_output_chunk is not None
+                        assert final_shared_hidden_states is not None
+                        final_shared_hidden_states[chunk_start:chunk_end, :].copy_(
+                            shared_output_chunk, non_blocking=True
+                        )
+                    else:
+                        assert isinstance(chunk_result, torch.Tensor)
+                        final_fused_hidden_states[chunk_start:chunk_end, :].copy_(
+                            chunk_result, non_blocking=True
+                        )
+
+        if self.shared_experts is None:
+            return final_fused_hidden_states
+        else:
+            assert final_shared_hidden_states is not None
+            return (final_shared_hidden_states, final_fused_hidden_states)

vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py

@@ -1,516 +1,45 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections.abc import Callable
-from contextlib import nullcontext
-from typing import TYPE_CHECKING
 
 import torch
-import torch.nn.functional as F
 
 from vllm.distributed import (
     get_ep_group,
     get_pcp_group,
-    tensor_model_parallel_all_reduce,
 )
-from vllm.forward_context import (
-    ForwardContext,
-    get_forward_context,
-    is_forward_context_available,
-)
-from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.config import (
-    FusedMoEConfig,
-)
-from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
-    FusedMoEMethodBase,
-)
-from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
-    FusedMoERouter,
-)
-from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
-from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
-    SharedExperts,
-    SharedExpertsOrder,
-)
-from vllm.platforms import current_platform
-from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import (
-    HAS_OPAQUE_TYPE,
-    ModuleName,
-    direct_register_custom_op,
-)
-from vllm.v1.worker.ubatching import dbo_current_ubatch_id
-
-logger = init_logger(__name__)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner_base import MoERunnerBase
 
 
-def get_layer_from_name(layer_name: str) -> torch.nn.Module:
-    forward_context: ForwardContext = get_forward_context()
-    if layer_name == "from_forward_context":
-        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 `all_moe_layers` "
-                "to be equal to the number of "
-                "{vllm.moe_forward, vllm.moe_forward_shared} calls."
-            )
-        layer_name = all_moe_layers[moe_layer_index]
-        forward_context.moe_layer_index += 1
-    return forward_context.no_compile_layers[layer_name]
-
-
-# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object;
-# on older versions it remains a plain str.
-if TYPE_CHECKING:
-    from typing import TypeAlias
-
-    _layer_name_type: TypeAlias = str | ModuleName
-else:
-    _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str
-
-
-def _resolve_layer_name(layer_name: str | ModuleName) -> str:
-    return layer_name.value if isinstance(layer_name, ModuleName) else layer_name
-
-
-# Note: _moe_forward and _moe_forward_shared should not contain any
-# implementation details, They should merely pass along control to
-# the runner's 'forward_dispatch' method.
-def _moe_forward(
-    hidden_states: torch.Tensor,
-    router_logits: torch.Tensor,
-    shared_experts_input: torch.Tensor | None,
-    layer_name: _layer_name_type,
-) -> torch.Tensor:
-    layer = get_layer_from_name(_resolve_layer_name(layer_name))
-    return layer.runner.forward_dispatch(
-        layer,
-        hidden_states,
-        router_logits,
-        shared_experts_input,
-    )
-
-
-def _moe_forward_fake(
-    hidden_states: torch.Tensor,
-    router_logits: torch.Tensor,
-    shared_experts_input: torch.Tensor | None,
-    layer_name: _layer_name_type,
-) -> torch.Tensor:
-    return torch.empty_like(hidden_states)
-
-
-def _moe_forward_shared(
-    hidden_states: torch.Tensor,
-    router_logits: torch.Tensor,
-    shared_experts_input: torch.Tensor | None,
-    layer_name: _layer_name_type,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    layer = get_layer_from_name(_resolve_layer_name(layer_name))
-    return layer.runner.forward_dispatch(
-        layer,
-        hidden_states,
-        router_logits,
-        shared_experts_input,
-    )
-
-
-def _moe_forward_shared_fake(
-    hidden_states: torch.Tensor,
-    router_logits: torch.Tensor,
-    shared_experts_input: torch.Tensor | None,
-    layer_name: _layer_name_type,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    # Output shapes:
-    # - fused_out: same as hidden_states (routed experts use transformed size)
-    # - shared_out: same as shared_experts_input if provided, else same as
-    #               hidden_states
-    # (For latent MoE: shared experts use original hidden_size, not latent size)
-    fused_out = torch.empty_like(hidden_states)
-    if shared_experts_input is not None:
-        shared_out = torch.empty_like(shared_experts_input)
-    else:
-        shared_out = torch.empty_like(hidden_states)
-    return shared_out, fused_out
-
-
-direct_register_custom_op(
-    op_name="moe_forward",
-    op_func=_moe_forward,
-    mutates_args=["hidden_states"],  # is this still true?
-    fake_impl=_moe_forward_fake,
-    tags=(torch.Tag.needs_fixed_stride_order,),
-)
-
-
-direct_register_custom_op(
-    op_name="moe_forward_shared",
-    op_func=_moe_forward_shared,
-    fake_impl=_moe_forward_shared_fake,
-    tags=(torch.Tag.needs_fixed_stride_order,),
-)
-
-
-class DefaultMoERunner(MoERunner):
+class DefaultMoERunner(MoERunnerBase):
     """
-    Default implementation of the MoE runner for executing Mixture of Experts layers.
+    Standard MoE runner implementation for executing Mixture of Experts layers.
 
-    This class provides a comprehensive implementation for running MoE computations
-    with support for:
-    - Expert routing and token dispatching
+    This is the primary concrete implementation of MoE execution logic, providing
+    comprehensive support for standard MoE operations. It handles:
+    - Expert routing and token dispatching using various routing strategies
     - Shared experts computation with optional parallel execution using CUDA streams
-    - Data parallel (DP) chunking for large batch processing
     - Tensor model parallel and expert parallel operations
-    - Various quantization methods and custom operators
+    - Multiple quantization methods and optimized kernel selection
     - Both monolithic and decomposed expert execution paths
+    - Integration with various parallel execution modes (TP, EP, DP)
 
-    The runner handles the complete MoE forward pass including routing tokens to
-    experts, executing expert computations, and combining results. It supports
-    advanced features like overlapped execution of shared experts and optimized
-    kernels for different parallel execution modes.
+    The runner orchestrates the complete MoE forward pass including routing tokens
+    to experts, executing expert computations in parallel, and combining results.
+    It supports advanced features like overlapped execution of shared experts,
+    optimized kernels for different parallel configurations, and seamless
+    integration with vLLM's distributed execution framework.
 
-    Eventually, this class will be split up and specialized for different
-    configurations, e.g. the presence or absence of shared experts, a gate, etc.
+    This implementation is suitable for most standard MoE use cases. For specialized
+    scenarios like large batch chunking, alternative runners like ChunkingMoERunner
+    may be more appropriate.
+
+    Eventually, this class may be split into more specialized implementations
+    for different configurations (e.g., with/without shared experts, gates, etc.).
     """
 
-    def __init__(
-        self,
-        layer: torch.nn.Module,
-        moe_config: FusedMoEConfig,
-        router: FusedMoERouter,
-        routed_input_transform: torch.nn.Module | None,
-        gate: torch.nn.Module | None,
-        shared_experts: torch.nn.Module | None,
-        quant_method: FusedMoEMethodBase,
-        reduce_results: bool,
-        enable_dbo: bool,
-    ):
-        super().__init__()
-        self.moe_config = moe_config
-        self.router = router
-        self.routed_input_transform = routed_input_transform
-        self.gate = gate
-        self.quant_method = quant_method
-        self.reduce_results = reduce_results
-        self.enable_dbo = enable_dbo
-
-        self.shared_experts: SharedExperts | None = None
-        if shared_experts is not None:
-            self.shared_experts = SharedExperts(
-                shared_experts,
-                moe_config=moe_config,
-                # Note: For now we must pass quant_method along to SharedExperts so it
-                # can property determine where the shared experts are supposed to be
-                # called, i.e. by a MK or by the MoERunner.
-                # Once the MK can be created upfront, we can just pass in the proper
-                # flags derived from the quant_method's MK.
-                reduce_results=reduce_results,
-                quant_method=quant_method,
-                enable_dbo=enable_dbo,
-            )
-
-        # Chunked all2all staging tensor
-        # These need to exist ahead of time due to CUDAgraph construction
-        # needing a fixed buffer address.
-        self.use_dp_chunking = self.moe_config.moe_parallel_config.use_dp_chunking
-        self.batched_hidden_states: torch.Tensor | None = None
-        self.batched_router_logits: torch.Tensor | None = None
-        self._maybe_init_dp_chunking()
-
-        # Needed for string -> FusedMoE layer lookup in custom ops.
-        self.layer_name = layer.layer_name
-
-        self.forward_entry, self.forward_impl = self._select_forward(layer)
-
-    def _select_forward(self, layer: torch.nn.Module) -> tuple[Callable, Callable]:
-        # Select implementation based on presence of DP chunking.
-        forward_impl_fn = (
-            self._forward_impl_chunked if self.use_dp_chunking else self._forward_impl
-        )
-
-        if current_platform.is_tpu() or current_platform.is_cpu():
-            # TODO: Once the OOM issue for the TPU backend is resolved, we
-            # will switch to using the moe_forward custom op.
-            # Note: CPU doesn't require wrapped forward_impl.
-            return (
-                _moe_forward if self.shared_experts is None else _moe_forward_shared,
-                forward_impl_fn,
-            )
-
-        return (
-            torch.ops.vllm.moe_forward
-            if self.shared_experts is None
-            else torch.ops.vllm.moe_forward_shared,
-            forward_impl_fn,
-        )
-
-    # TODO(bnell): temporary hack, do not call this method.
-    def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
-        if self.shared_experts is not None:
-            self.shared_experts._quant_method = quant_method
-        self.quant_method = quant_method
-
-    def is_internal_router(self) -> bool:
-        return self.gate is not None
-
-    def _maybe_init_dp_chunking(self):
-        if not self.use_dp_chunking:
-            return
-
-        assert self.batched_hidden_states is None
-        states_shape: tuple[int, ...]
-        logits_shape: tuple[int, ...]
-
-        moe = self.moe_config
-
-        if self.enable_dbo:
-            states_shape = (2, moe.max_num_tokens, self.moe_config.hidden_dim)
-            logits_shape = (2, moe.max_num_tokens, self.moe_config.num_logical_experts)
-        else:
-            states_shape = (moe.max_num_tokens, self.moe_config.hidden_dim)
-            logits_shape = (moe.max_num_tokens, self.moe_config.num_logical_experts)
-
-        device = torch.accelerator.current_device_index()
-        self.batched_hidden_states = torch.zeros(
-            states_shape,
-            dtype=moe.in_dtype,
-            device=device,
-        )
-
-        self.batched_router_logits = torch.zeros(
-            logits_shape,
-            dtype=moe.router_logits_dtype,
-            device=device,
-        )
-
-    def must_reduce_shared_expert_outputs(self) -> bool:
-        """
-        The shared_experts are typically computed using the RowParallelLinear
-        layer. The result of this function is typically used as
-        the reduce_results argument to the module.
-        When just tensor-parallel is used, it is not required to reduce
-        the shared_experts results immediately. Instead we reduce at the
-        once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
-        With EP and all2all kernels - this is no longer viable as all
-        GPU ranks in DP, produce the complete set of hidden_states.
-        Therefore it is required that we reduce the shared_experts output
-        early.
-        """
-        return (
-            self.quant_method.moe_kernel is not None
-            and self.quant_method.moe_kernel.output_is_reduced()
-        )
-
-    def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):
-        """
-        Some combine kernels reduce across GPU ranks by default.
-        """
-        if self.must_reduce_shared_expert_outputs():
-            return final_hidden_states
-        else:
-            return tensor_model_parallel_all_reduce(final_hidden_states)
-
-    def apply_routed_input_transform(
-        self, hidden_states: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Apply transform for routed experts (e.g., latent projection).
-
-        This is called by FusedMoE.forward_native. The original hidden_states
-        is saved separately so shared experts get [S, hidden_size] while
-        routed experts get the transformed [S, moe_latent_size].
-
-        TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be
-        moved inside SharedFusedMoE to all-reduce on the smaller latent
-        dimension.
-
-        Returns (possibly transformed) hidden states and the input for shared
-        experts (or None if there are no shared experts).
-        """
-        if self.routed_input_transform is not None:
-            result = self.routed_input_transform(hidden_states)
-            # ReplicatedLinear returns (output, extra_bias) tuple.
-            # We only need the output tensor; extra_bias is not used here.
-            if isinstance(result, tuple):
-                return result[0], hidden_states
-            return result, hidden_states
-
-        return (
-            hidden_states,
-            hidden_states if self.shared_experts is not None else None,
-        )
-
-    def _maybe_reduce_output(
-        self,
-        states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-        trunc_sizes: list[int],
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
-            return x[..., :trunc_size]
-
-        def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
-            return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size)
-
-        if (
-            not self.moe_config.is_sequence_parallel
-            and not self.use_dp_chunking
-            and self.reduce_results
-            and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
-        ):
-            func = reduce_and_trunc
-        else:
-            func = trunc
-
-        if isinstance(states, tuple):
-            return tuple(
-                [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)]
-            )
-        else:
-            assert len(trunc_sizes) == 1
-            return func(states, trunc_sizes[0])
-
-    def _encode_layer_name(self) -> str | ModuleName:
-        if HAS_OPAQUE_TYPE:
-            return ModuleName(self.layer_name)
-        # Can be unavailable or None in unittests
-        if (
-            is_forward_context_available()
-            and get_forward_context().all_moe_layers is not None
-        ):
-            return "from_forward_context"
-        return self.layer_name
-
-    def _maybe_pad_hidden_states(
-        self,
-        shared_experts_input: torch.Tensor | None,
-        hidden_states: torch.Tensor,
-    ) -> tuple[torch.Tensor, list[int]]:
-        shared_experts_hidden_dim = (
-            shared_experts_input.shape[-1] if shared_experts_input is not None else 0
-        )
-        transformed_hidden_dim = hidden_states.shape[-1]
-        if (
-            not self.quant_method.skip_forward_padding
-            and self.moe_config.hidden_dim != transformed_hidden_dim
-        ):
-            hidden_states = F.pad(
-                hidden_states,
-                (0, self.moe_config.hidden_dim - transformed_hidden_dim),
-                mode="constant",
-                value=0.0,
-            )
-
-        if self.shared_experts is not None:
-            orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim]
-        else:
-            orig_hidden_dims = [transformed_hidden_dim]
-
-        return hidden_states, orig_hidden_dims
-
-    def _maybe_apply_shared_experts(
-        self,
-        shared_experts_input: torch.Tensor | None,
-        order: SharedExpertsOrder,
-    ):
-        if self.shared_experts is not None:
-            assert shared_experts_input is not None
-            self.shared_experts.apply(shared_experts_input, order)
-
-    def _apply_quant_method(
-        self,
-        layer: torch.nn.Module,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-        shared_experts_input: torch.Tensor | None,
-    ) -> tuple[torch.Tensor | None, torch.Tensor]:
-        # Run this before quant_method to avoid inplace issues.
-        # TODO(bnell): probably not needed anymore since inplace is
-        # disabled when shared experts are present.
-        self._maybe_apply_shared_experts(
-            shared_experts_input, SharedExpertsOrder.NO_OVERLAP
-        )
-
-        if self.quant_method.is_monolithic:
-            fused_out = self.quant_method.apply_monolithic(
-                layer=layer,
-                x=hidden_states,
-                router_logits=router_logits,
-            )
-        else:
-            topk_weights, topk_ids = self.router.select_experts(
-                hidden_states=hidden_states,
-                router_logits=router_logits,
-            )
-
-            # Passing shared_experts_input in case SharedExpertsOrder is
-            # NO_OVERLAP or MK_INTERNAL_OVERLAPPED.
-            fused_out = self.quant_method.apply(
-                layer=layer,
-                x=hidden_states,
-                topk_weights=topk_weights,
-                topk_ids=topk_ids,
-                shared_experts_input=shared_experts_input,
-            )
-
-        self._maybe_apply_shared_experts(
-            shared_experts_input,
-            SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
-        )
-
-        return (
-            self.shared_experts.output if self.shared_experts is not None else None,
-            fused_out,
-        )
-
-    def _sequence_parallel_context(self):
-        ctx = get_forward_context()
-        return (
-            ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
-            if ctx.dp_metadata
-            else nullcontext()
-        )
-
-    def _allocate_dp_chunking_outputs(
-        self,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-    ) -> tuple[torch.Tensor | None, torch.Tensor]:
-        assert self.use_dp_chunking
-
-        # Assert the inputs are of the proper type and shape.
-        assert self.batched_hidden_states is not None
-        assert self.batched_router_logits is not None
-
-        assert self.batched_hidden_states.dtype == hidden_states.dtype, (
-            f"{self.batched_hidden_states.dtype} == {hidden_states.dtype}"
-        )
-        assert self.batched_router_logits.dtype == router_logits.dtype, (
-            f"{self.batched_router_logits.dtype} == {router_logits.dtype}"
-        )
-
-        # Check size compatibility.
-        assert self.batched_hidden_states.size(-1) == hidden_states.size(-1)
-        assert self.batched_router_logits.size(-1) == router_logits.size(-1)
-
-        final_fused_hidden_states = torch.empty_like(hidden_states)
-        if self.shared_experts is not None:
-            final_shared_hidden_states = torch.empty_like(hidden_states)
-        else:
-            final_shared_hidden_states = None
-
-        return final_shared_hidden_states, final_fused_hidden_states
-
-    def _maybe_sync_shared_experts_stream(
-        self,
-        shared_experts_input: torch.Tensor | None,
-    ):
-        # If router/gate provided, then apply it here.
-        # (Note: This code runs only when "overlapped mode" is on to allow
-        #        parallel execution of shared experts with the FusedMoE via
-        #        separate cuda stream)
-        if self.shared_experts is not None:
-            self.shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
+    @property
+    def reduce_results(self) -> bool:
+        return self._reduce_results
 
     @property
     def do_naive_dispatch_combine(self) -> bool:
@@ -572,195 +101,6 @@ class DefaultMoERunner(MoERunner):
         else:
             return hidden_states
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        """Invoke the fused moe layer.
-
-        Input:
-        - hidden_states
-        - router_logits
-
-        Output:
-        - The new hidden_states.
-        or
-        - A tuple of (shared experts output, new hidden_states).
-
-        Calling sequence
-        - forward
-          - self.forward_entry (_moe_forward or _moe_forward_shared custom op)
-            - forward_dispatch
-              - forward_impl (_forward_impl or _forward_impl_chunked)
-
-        Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
-        to the following reasons:
-        1. the chunking loop in _forward_impl_chunked cannot be compiled by
-           torch.compile
-        2. pytorch cannot handle union types in custom op signatures so _moe_forward
-           and _moe_forward_shared must be split.
-
-        If _forward_impl_chunked can be implemented via torch.scan we can potentially
-        get rid of _moe_forward and _moe_forward_shared and collapse the whole sequence
-        into the 'forward' method.
-        """
-
-        # Apply transform for routed experts (e.g., latent projection for latent MoE)
-        hidden_states, shared_experts_input = self.apply_routed_input_transform(
-            hidden_states
-        )
-
-        hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(
-            shared_experts_input,
-            hidden_states,
-        )
-
-        fused_output = self.forward_entry(
-            hidden_states,
-            router_logits,
-            shared_experts_input,
-            self._encode_layer_name(),
-        )
-
-        return self._maybe_reduce_output(fused_output, og_hidden_dims)
-
-    def forward_dispatch(
-        self,
-        layer: torch.nn.Module,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-        shared_experts_input: torch.Tensor | None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        # TODO(bnell): this can be removed after MK migration is complete.
-        layer.ensure_moe_quant_config_init()
-
-        # Sync aux and main stream for shared expert multi-stream overlap.
-        self._maybe_sync_shared_experts_stream(shared_experts_input)
-
-        # If the Runner holds the gate, apply it after the stream sync,
-        # so it can run overlapped with the
-        # NOTE: in future PR, MoE runner will always hold the gate.
-        if self.gate is not None:
-            router_logits, _ = self.gate(hidden_states)
-
-        self._maybe_apply_shared_experts(
-            shared_experts_input,
-            SharedExpertsOrder.EXTERNAL,
-        )
-
-        with self._sequence_parallel_context():
-            return self.forward_impl(
-                layer,
-                hidden_states,
-                router_logits,
-                shared_experts_input,
-            )
-
-    def _slice_and_copy_input(
-        self,
-        out_slice: torch.Tensor,
-        orig: torch.Tensor | None,
-        start: int,
-        end: int,
-    ) -> torch.Tensor:
-        assert orig is not None
-        slice_size = end - start
-        orig_slice = orig[start:end, :]
-        if self.enable_dbo:
-            assert out_slice.dim() == 3
-            batch_buffer_idx = dbo_current_ubatch_id()
-            out_slice = out_slice[batch_buffer_idx, :]
-
-        assert out_slice.size(0) >= slice_size
-        out_slice = out_slice[:slice_size, :]
-        out_slice.copy_(orig_slice, non_blocking=True)
-        return out_slice
-
-    def _forward_impl_chunked(
-        self,
-        layer: torch.nn.Module,
-        hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
-        shared_experts_input: torch.Tensor | None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        final_shared_hidden_states, final_fused_hidden_states = (
-            self._allocate_dp_chunking_outputs(hidden_states, router_logits)
-        )
-
-        ctx = get_forward_context()
-        # flashinfer_cutlass_kernels can handle: optional DP + TP/EP
-        max_tokens_across_dispatchers = ctx.dp_metadata.max_tokens_across_dp_cpu
-        moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens
-
-        # If the input to the MoE is sequence parallel then divide by sp_size
-        # to find the maximum number of tokens for any individual dispatcher.
-        if self.moe_config.is_sequence_parallel:
-            max_tokens_across_dispatchers = cdiv(
-                max_tokens_across_dispatchers, self.moe_config.sp_size
-            )
-
-        num_tokens = hidden_states.size(0)
-        for chunk_idx, chunk_start_ in enumerate(
-            range(0, max_tokens_across_dispatchers, moe_dp_chunk_size_per_rank)
-        ):
-            chunk_start = chunk_start_
-            chunk_end = min(
-                chunk_start + moe_dp_chunk_size_per_rank, max_tokens_across_dispatchers
-            )
-            # clamp start and end
-            chunk_start = min(chunk_start, num_tokens - 1)
-            chunk_end = min(chunk_end, num_tokens)
-            chunk_sizes = ctx.dp_metadata.chunked_sizes(
-                self.moe_config.sp_size, moe_dp_chunk_size_per_rank, chunk_idx
-            )
-            with chunk_sizes:
-                hidden_states_chunk = self._slice_and_copy_input(
-                    self.batched_hidden_states,
-                    hidden_states,
-                    chunk_start,
-                    chunk_end,
-                )
-
-                router_logits_chunk = self._slice_and_copy_input(
-                    self.batched_router_logits,
-                    router_logits,
-                    chunk_start,
-                    chunk_end,
-                )
-
-                shared_experts_input_chunk = (
-                    shared_experts_input[chunk_start:chunk_end, :]
-                    if shared_experts_input is not None
-                    else None
-                )
-
-                shared_output_chunk, hidden_states_chunk = self._apply_quant_method(
-                    layer=layer,
-                    hidden_states=hidden_states_chunk,
-                    router_logits=router_logits_chunk,
-                    shared_experts_input=shared_experts_input_chunk,
-                )
-
-                # Store outputs
-                # TODO(bnell): document when chunk_start >= num_tokens
-                if chunk_start < num_tokens:
-                    final_fused_hidden_states[chunk_start:chunk_end, :].copy_(
-                        hidden_states_chunk, non_blocking=True
-                    )
-                    if self.shared_experts is not None:
-                        assert shared_output_chunk is not None
-                        assert final_shared_hidden_states is not None
-                        final_shared_hidden_states[chunk_start:chunk_end, :].copy_(
-                            shared_output_chunk, non_blocking=True
-                        )
-
-        if self.shared_experts is None:
-            return final_fused_hidden_states
-        else:
-            assert final_shared_hidden_states is not None
-            return (final_shared_hidden_states, final_fused_hidden_states)
-
     def _forward_impl(
         self,
         layer: torch.nn.Module,

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -4,6 +4,13 @@ from abc import ABC, abstractmethod
 
 import torch
 
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+)
+
 
 class MoERunner(ABC):
     """
@@ -36,3 +43,13 @@ class MoERunner(ABC):
     @abstractmethod
     def is_internal_router(self) -> bool:
         raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def shared_experts(self) -> SharedExperts | None:
+        raise NotImplementedError
+
+    # TODO(bnell): temporary hack, do not call this method.
+    @abstractmethod
+    def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
+        raise NotImplementedError

vllm/model_executor/layers/fused_moe/runner/moe_runner_base.py

@@ -0,0 +1,527 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from abc import abstractmethod
+from collections.abc import Callable
+from contextlib import nullcontext
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F
+
+from vllm.distributed import (
+    tensor_model_parallel_all_reduce,
+)
+from vllm.forward_context import (
+    ForwardContext,
+    get_forward_context,
+    is_forward_context_available,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
+    FusedMoERouter,
+)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+    SharedExpertsOrder,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import (
+    HAS_OPAQUE_TYPE,
+    ModuleName,
+    direct_register_custom_op,
+)
+
+
+def get_layer_from_name(layer_name: str) -> torch.nn.Module:
+    forward_context: ForwardContext = get_forward_context()
+    if layer_name == "from_forward_context":
+        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 `all_moe_layers` "
+                "to be equal to the number of "
+                "{vllm.moe_forward, vllm.moe_forward_shared} calls."
+            )
+        layer_name = all_moe_layers[moe_layer_index]
+        forward_context.moe_layer_index += 1
+    return forward_context.no_compile_layers[layer_name]
+
+
+# On torch >= 2.11, layer_name is a hoisted ModuleName opaque object;
+# on older versions it remains a plain str.
+if TYPE_CHECKING:
+    from typing import TypeAlias
+
+    _layer_name_type: TypeAlias = str | ModuleName
+else:
+    _layer_name_type = ModuleName if HAS_OPAQUE_TYPE else str
+
+
+def _resolve_layer_name(layer_name: str | ModuleName) -> str:
+    return layer_name.value if isinstance(layer_name, ModuleName) else layer_name
+
+
+# Note: _moe_forward and _moe_forward_shared should not contain any
+# implementation details, They should merely pass along control to
+# the runner's 'forward_dispatch' method.
+def _moe_forward(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    shared_experts_input: torch.Tensor | None,
+    layer_name: _layer_name_type,
+) -> torch.Tensor:
+    layer = get_layer_from_name(_resolve_layer_name(layer_name))
+    return layer.runner.forward_dispatch(
+        layer,
+        hidden_states,
+        router_logits,
+        shared_experts_input,
+    )
+
+
+def _moe_forward_fake(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    shared_experts_input: torch.Tensor | None,
+    layer_name: _layer_name_type,
+) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
+def _moe_forward_shared(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    shared_experts_input: torch.Tensor | None,
+    layer_name: _layer_name_type,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    layer = get_layer_from_name(_resolve_layer_name(layer_name))
+    return layer.runner.forward_dispatch(
+        layer,
+        hidden_states,
+        router_logits,
+        shared_experts_input,
+    )
+
+
+def _moe_forward_shared_fake(
+    hidden_states: torch.Tensor,
+    router_logits: torch.Tensor,
+    shared_experts_input: torch.Tensor | None,
+    layer_name: _layer_name_type,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Output shapes:
+    # - fused_out: same as hidden_states (routed experts use transformed size)
+    # - shared_out: same as shared_experts_input if provided, else same as
+    #               hidden_states
+    # (For latent MoE: shared experts use original hidden_size, not latent size)
+    fused_out = torch.empty_like(hidden_states)
+    if shared_experts_input is not None:
+        shared_out = torch.empty_like(shared_experts_input)
+    else:
+        shared_out = torch.empty_like(hidden_states)
+    return shared_out, fused_out
+
+
+direct_register_custom_op(
+    op_name="moe_forward",
+    op_func=_moe_forward,
+    mutates_args=["hidden_states"],  # is this still true?
+    fake_impl=_moe_forward_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
+
+
+direct_register_custom_op(
+    op_name="moe_forward_shared",
+    op_func=_moe_forward_shared,
+    fake_impl=_moe_forward_shared_fake,
+    tags=(torch.Tag.needs_fixed_stride_order,),
+)
+
+
+class MoERunnerBase(MoERunner):
+    """
+    Abstract base class providing common functionality for MoE runner implementations.
+
+    This class serves as the foundation for concrete MoE runner implementations by
+    providing shared state management and common utilities. It handles:
+    - Common initialization and configuration management
+    - Shared expert output reduction logic for tensor parallel scenarios
+    - Base methods for tensor model parallel reductions
+    - Common properties and utility functions used across different runner types
+
+    Concrete subclasses must implement the abstract methods to define their specific
+    execution strategies, such as standard execution, chunked processing, or other
+    specialized approaches. The base class provides the infrastructure while
+    allowing flexibility in the actual MoE computation implementation.
+
+    Key abstract methods that subclasses must implement:
+    - reduce_results: Determines whether results should be reduced across ranks
+    - _forward_impl: The core MoE computation logic specific to each runner type
+    """
+
+    def __init__(
+        self,
+        layer_name: str,
+        moe_config: FusedMoEConfig,
+        router: FusedMoERouter,
+        routed_input_transform: torch.nn.Module | None,
+        gate: torch.nn.Module | None,
+        shared_experts: torch.nn.Module | None,
+        quant_method: FusedMoEMethodBase,
+        reduce_results: bool,
+        enable_dbo: bool,
+    ):
+        super().__init__()
+        self.moe_config = moe_config
+        self.router = router
+        self.routed_input_transform = routed_input_transform
+        self.gate = gate
+        self.quant_method = quant_method
+        self._reduce_results = reduce_results
+        self.enable_dbo = enable_dbo
+
+        self._shared_experts: SharedExperts | None = None
+        if shared_experts is not None:
+            self._shared_experts = SharedExperts(
+                shared_experts,
+                moe_config=moe_config,
+                # Note: For now we must pass quant_method along to SharedExperts so it
+                # can property determine where the shared experts are supposed to be
+                # called, i.e. by a MK or by the MoERunner.
+                # Once the MK can be created upfront, we can just pass in the proper
+                # flags derived from the quant_method's MK.
+                reduce_results=reduce_results,
+                quant_method=quant_method,
+                enable_dbo=enable_dbo,
+            )
+
+        # Needed for string -> FusedMoE layer lookup in custom ops.
+        self.layer_name = layer_name
+
+        self.forward_entry = self._select_forward()
+
+    def _select_forward(self) -> Callable:
+        if current_platform.is_tpu() or current_platform.is_cpu():
+            # TODO: Once the OOM issue for the TPU backend is resolved, we
+            # will switch to using the moe_forward custom op.
+            # Note: CPU doesn't require wrapped _forward_impl.
+            return _moe_forward if self._shared_experts is None else _moe_forward_shared
+
+        return (
+            torch.ops.vllm.moe_forward
+            if self._shared_experts is None
+            else torch.ops.vllm.moe_forward_shared
+        )
+
+    @property
+    def shared_experts(self) -> SharedExperts | None:
+        return self._shared_experts
+
+    # TODO(bnell): temporary hack, do not call this method.
+    def _replace_quant_method(self, quant_method: FusedMoEMethodBase):
+        if self._shared_experts is not None:
+            self._shared_experts._quant_method = quant_method
+        self.quant_method = quant_method
+
+    def is_internal_router(self) -> bool:
+        return self.gate is not None
+
+    @property
+    @abstractmethod
+    def reduce_results(self) -> bool:
+        raise NotImplementedError
+
+    def must_reduce_shared_expert_outputs(self) -> bool:
+        """
+        The shared_experts are typically computed using the RowParallelLinear
+        layer. The result of this function is typically used as
+        the reduce_results argument to the module.
+        When just tensor-parallel is used, it is not required to reduce
+        the shared_experts results immediately. Instead we reduce at the
+        once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
+        With EP and all2all kernels - this is no longer viable as all
+        GPU ranks in DP, produce the complete set of hidden_states.
+        Therefore it is required that we reduce the shared_experts output
+        early.
+        """
+        return (
+            self.quant_method.moe_kernel is not None
+            and self.quant_method.moe_kernel.output_is_reduced()
+        )
+
+    def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):
+        """
+        Some combine kernels reduce across GPU ranks by default.
+        """
+        if self.must_reduce_shared_expert_outputs():
+            return final_hidden_states
+        else:
+            return tensor_model_parallel_all_reduce(final_hidden_states)
+
+    def apply_routed_input_transform(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Apply transform for routed experts (e.g., latent projection).
+
+        This is called by FusedMoE.forward_native. The original hidden_states
+        is saved separately so shared experts get [S, hidden_size] while
+        routed experts get the transformed [S, moe_latent_size].
+
+        TODO: For latent MoE bandwidth optimization, fc2_latent_proj could be
+        moved inside SharedFusedMoE to all-reduce on the smaller latent
+        dimension.
+
+        Returns (possibly transformed) hidden states and the input for shared
+        experts (or None if there are no shared experts).
+        """
+        if self.routed_input_transform is not None:
+            result = self.routed_input_transform(hidden_states)
+            # ReplicatedLinear returns (output, extra_bias) tuple.
+            # We only need the output tensor; extra_bias is not used here.
+            if isinstance(result, tuple):
+                return result[0], hidden_states
+            return result, hidden_states
+
+        return (
+            hidden_states,
+            hidden_states if self._shared_experts is not None else None,
+        )
+
+    def _maybe_reduce_output(
+        self,
+        states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        trunc_sizes: list[int],
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        def trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
+            return x[..., :trunc_size]
+
+        def reduce_and_trunc(x: torch.Tensor, trunc_size: int) -> torch.Tensor:
+            return trunc(self.maybe_all_reduce_tensor_model_parallel(x), trunc_size)
+
+        if (
+            not self.moe_config.is_sequence_parallel
+            and self.reduce_results
+            and (self.moe_config.tp_size > 1 or self.moe_config.ep_size > 1)
+        ):
+            func = reduce_and_trunc
+        else:
+            func = trunc
+
+        if isinstance(states, tuple):
+            return tuple(
+                [func(s, trunc_size) for s, trunc_size in zip(states, trunc_sizes)]
+            )
+        else:
+            assert len(trunc_sizes) == 1
+            return func(states, trunc_sizes[0])
+
+    def _encode_layer_name(self) -> str | ModuleName:
+        if HAS_OPAQUE_TYPE:
+            return ModuleName(self.layer_name)
+        # Can be unavailable or None in unittests
+        if (
+            is_forward_context_available()
+            and get_forward_context().all_moe_layers is not None
+        ):
+            return "from_forward_context"
+        return self.layer_name
+
+    def _maybe_pad_hidden_states(
+        self,
+        shared_experts_input: torch.Tensor | None,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, list[int]]:
+        shared_experts_hidden_dim = (
+            shared_experts_input.shape[-1] if shared_experts_input is not None else 0
+        )
+        transformed_hidden_dim = hidden_states.shape[-1]
+        if (
+            not self.quant_method.skip_forward_padding
+            and self.moe_config.hidden_dim != transformed_hidden_dim
+        ):
+            hidden_states = F.pad(
+                hidden_states,
+                (0, self.moe_config.hidden_dim - transformed_hidden_dim),
+                mode="constant",
+                value=0.0,
+            )
+
+        if self._shared_experts is not None:
+            orig_hidden_dims = [shared_experts_hidden_dim, transformed_hidden_dim]
+        else:
+            orig_hidden_dims = [transformed_hidden_dim]
+
+        return hidden_states, orig_hidden_dims
+
+    def _maybe_apply_shared_experts(
+        self,
+        shared_experts_input: torch.Tensor | None,
+        order: SharedExpertsOrder,
+    ):
+        if self._shared_experts is not None:
+            assert shared_experts_input is not None
+            self._shared_experts.apply(shared_experts_input, order)
+
+    def _apply_quant_method(
+        self,
+        layer: torch.nn.Module,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor]:
+        # Run this before quant_method to avoid inplace issues.
+        # TODO(bnell): probably not needed anymore since inplace is
+        # disabled when shared experts are present.
+        self._maybe_apply_shared_experts(
+            shared_experts_input, SharedExpertsOrder.NO_OVERLAP
+        )
+
+        if self.quant_method.is_monolithic:
+            fused_out = self.quant_method.apply_monolithic(
+                layer=layer,
+                x=hidden_states,
+                router_logits=router_logits,
+            )
+        else:
+            topk_weights, topk_ids = self.router.select_experts(
+                hidden_states=hidden_states,
+                router_logits=router_logits,
+            )
+
+            # Passing shared_experts_input in case SharedExpertsOrder is
+            # NO_OVERLAP or MK_INTERNAL_OVERLAPPED.
+            fused_out = self.quant_method.apply(
+                layer=layer,
+                x=hidden_states,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                shared_experts_input=shared_experts_input,
+            )
+
+        self._maybe_apply_shared_experts(
+            shared_experts_input,
+            SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
+        )
+
+        return (
+            self._shared_experts.output if self._shared_experts is not None else None,
+            fused_out,
+        )
+
+    def _sequence_parallel_context(self):
+        ctx = get_forward_context()
+        return (
+            ctx.dp_metadata.sp_local_sizes(self.moe_config.sp_size)
+            if ctx.dp_metadata
+            else nullcontext()
+        )
+
+    def _maybe_sync_shared_experts_stream(
+        self,
+        shared_experts_input: torch.Tensor | None,
+    ):
+        # If router/gate provided, then apply it here.
+        # (Note: This code runs only when "overlapped mode" is on to allow
+        #        parallel execution of shared experts with the FusedMoE via
+        #        separate cuda stream)
+        if self._shared_experts is not None:
+            self._shared_experts.maybe_sync_shared_experts_stream(shared_experts_input)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Invoke the fused moe layer.
+
+        Input:
+        - hidden_states
+        - router_logits
+
+        Output:
+        - The new hidden_states.
+        or
+        - A tuple of (shared experts output, new hidden_states).
+
+        Calling sequence
+        - forward
+          - self.forward_entry (_moe_forward or _moe_forward_shared custom op)
+            - forward_dispatch
+              - _forward_impl
+
+        Note: The existence of _moe_forward and _moe_forward_shared custom ops are due
+        to the following reasons:
+        1. the chunking loop in ChunkingMoERunner._forward_impl cannot be compiled by
+           torch.compile
+        2. pytorch cannot handle union types in custom op signatures so _moe_forward
+           and _moe_forward_shared must be split.
+
+        If ChunkingMoERunner._forward_impl can be implemented via torch.scan we can
+        potentially get rid of _moe_forward and _moe_forward_shared and collapse the
+        whole sequence into the 'forward' method.
+        """
+
+        # Apply transform for routed experts (e.g., latent projection for latent MoE)
+        hidden_states, shared_experts_input = self.apply_routed_input_transform(
+            hidden_states
+        )
+
+        hidden_states, og_hidden_dims = self._maybe_pad_hidden_states(
+            shared_experts_input,
+            hidden_states,
+        )
+
+        fused_output = self.forward_entry(
+            hidden_states,
+            router_logits,
+            shared_experts_input,
+            self._encode_layer_name(),
+        )
+
+        return self._maybe_reduce_output(fused_output, og_hidden_dims)
+
+    def forward_dispatch(
+        self,
+        layer: torch.nn.Module,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        # TODO(bnell): this can be removed after MK migration is complete.
+        layer.ensure_moe_quant_config_init()
+
+        # Sync aux and main stream for shared expert multi-stream overlap.
+        self._maybe_sync_shared_experts_stream(shared_experts_input)
+
+        # If the Runner holds the gate, apply it after the stream sync,
+        # so it can run overlapped with the
+        # NOTE: in future PR, MoE runner will always hold the gate.
+        if self.gate is not None:
+            router_logits, _ = self.gate(hidden_states)
+
+        with self._sequence_parallel_context():
+            return self._forward_impl(
+                layer,
+                hidden_states,
+                router_logits,
+                shared_experts_input,
+            )
+
+    @abstractmethod
+    def _forward_impl(
+        self,
+        layer: torch.nn.Module,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError

vllm/model_executor/layers/fused_moe/runner/moe_runner_factory.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
+    FusedMoERouter,
+)
+from vllm.model_executor.layers.fused_moe.runner.chunking_moe_runner import (
+    ChunkingMoERunner,
+)
+from vllm.model_executor.layers.fused_moe.runner.default_moe_runner import (
+    DefaultMoERunner,
+)
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+)
+
+
+def create_moe_runner(
+    layer_name: str,
+    moe_config: FusedMoEConfig,
+    router: FusedMoERouter,
+    routed_input_transform: torch.nn.Module | None,
+    gate: torch.nn.Module | None,
+    shared_experts: SharedExperts | None,
+    quant_method: FusedMoEMethodBase,
+    reduce_results: bool,
+    enable_dbo: bool,
+) -> MoERunner:
+    runner = DefaultMoERunner(
+        layer_name,
+        moe_config,
+        router,
+        routed_input_transform,
+        gate,
+        shared_experts,
+        quant_method,
+        reduce_results,
+        enable_dbo,
+    )
+    if moe_config.moe_parallel_config.use_dp_chunking:
+        return ChunkingMoERunner(runner)
+    return runner

vllm/model_executor/layers/fused_moe/runner/shared_experts.py

@@ -32,19 +32,14 @@ class SharedExpertsOrder(IntEnum):
     # No shared experts.
     NONE = (0,)
 
-    # Get rid of this one?  combine with BEFORE?
-    # Note: this might be important for torch.compile reasons. Can
-    # get rid of it after _moe_forward is undone.
-    EXTERNAL = (1,)
-
     # No overlap - defensively called before MK.
-    NO_OVERLAP = (2,)
+    NO_OVERLAP = (1,)
 
     # Overlapped with dispatch/combine in DP/EP - called by the MK.
-    MK_INTERNAL_OVERLAPPED = (3,)
+    MK_INTERNAL_OVERLAPPED = (2,)
 
     # Overlapped with the gate, router, experts in aux stream.
-    MULTI_STREAM_OVERLAPPED = (4,)
+    MULTI_STREAM_OVERLAPPED = (3,)
 
 
 class SharedExperts:
@@ -110,9 +105,6 @@ class SharedExperts:
         self,
         hidden_states: torch.Tensor,
     ) -> SharedExpertsOrder:
-        if self._use_external_experts:
-            return SharedExpertsOrder.EXTERNAL
-
         if self._quant_method.mk_owns_shared_expert:
             return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
 
@@ -205,12 +197,4 @@ class SharedExperts:
         else:
             self._output[self._output_idx] = self._layer(shared_experts_input)
 
-        if order == SharedExpertsOrder.EXTERNAL:
-            # TODO: figure out how to combine this with maybe_reduce_output?
-            # or get rid of it completely.
-            assert self._output[self._output_idx] is not None
-            self._output[self._output_idx] = self._maybe_reduce_shared_out(
-                self._output[self._output_idx]
-            )
-
         assert self._output[self._output_idx] is not None

vllm/model_executor/layers/layernorm.py

@@ -376,77 +376,32 @@ class GemmaRMSNorm(CustomOp):
         self.weight = nn.Parameter(torch.zeros(hidden_size))
         self.variance_epsilon = eps
 
-    @staticmethod
-    def _forward_static_no_residual(
-        weight: torch.Tensor,
-        variance_epsilon: float,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        """PyTorch-native implementation equivalent to forward() without residual."""
-        orig_dtype = x.dtype
-        x = x.float()
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-        x = x * torch.rsqrt(variance + variance_epsilon)
-        x = x * (1.0 + weight.float())
-        x = x.to(orig_dtype)
-        return x
-
-    @staticmethod
-    def _forward_static_with_residual(
-        weight: torch.Tensor,
-        variance_epsilon: float,
-        x: torch.Tensor,
-        residual: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """PyTorch-native implementation equivalent to forward() with residual."""
-        orig_dtype = x.dtype
-        x = (
-            x.float() + residual.float()
-            if orig_dtype == torch.float16
-            else x + residual
-        )
-        residual = x
-
-        x = x.float()
-        variance = x.pow(2).mean(dim=-1, keepdim=True)
-        x = x * torch.rsqrt(variance + variance_epsilon)
-        # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16)
-        # See https://github.com/huggingface/transformers/pull/29402
-        x = x * (1.0 + weight.float())
-        x = x.to(orig_dtype)
-        return x, residual
-
     def forward_native(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """PyTorch-native implementation equivalent to forward()."""
-        if residual is None:
-            return self._forward_static_no_residual(
-                self.weight.data, self.variance_epsilon, x
-            )
-        else:
-            return self._forward_static_with_residual(
-                self.weight.data, self.variance_epsilon, x, residual
+        orig_dtype = x.dtype
+        weight = self.weight.data.float() + 1.0
+        if residual is not None:
+            x = (
+                x.float() + residual.float()
+                if orig_dtype == torch.float16
+                else x + residual
             )
+            residual = x
+        # ir.ops.rms_norm handles fp32 upcast internally
+        out = ir.ops.rms_norm(x, weight, self.variance_epsilon)
+        return (
+            out.to(orig_dtype) if residual is None else (out.to(orig_dtype), residual)
+        )
 
     def forward_cuda(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        if torch.compiler.is_compiling():
-            return self.forward_native(x, residual)
-
-        if not getattr(self, "_is_compiled", False):
-            self._forward_static_no_residual = torch.compile(  # type: ignore
-                self._forward_static_no_residual
-            )
-            self._forward_static_with_residual = torch.compile(  # type: ignore
-                self._forward_static_with_residual
-            )
-            self._is_compiled = True
         return self.forward_native(x, residual)
 
 

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

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe import (  # noqa: E501
+    CompressedTensorsMoEMethod,
+)
+
+__all__ = [
+    "CompressedTensorsMoEMethod",
+]

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

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+from compressed_tensors import CompressionFormat
+from compressed_tensors.quantization import (
+    ActivationOrdering,
+    QuantizationStrategy,
+)
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEMethodBase,
+    UnquantizedFusedMoEMethod,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import (  # noqa
+    WNA16_SUPPORTED_BITS,
+)
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    check_moe_marlin_supports_layer,
+)
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+class CompressedTensorsMoEMethod(FusedMoEMethodBase):
+    @staticmethod
+    def get_moe_method(
+        quant_config: "CompressedTensorsConfig",  # type: ignore # noqa E501
+        layer: torch.nn.Module,
+        layer_name: str,
+    ) -> FusedMoEMethodBase:
+        # FusedMoE was made by combining multiple Linears so need to
+        # make sure quantization config for Linear can target it
+        quant_config._add_fused_moe_to_target_scheme_map()
+        unfused_names = [
+            layer_name + proj_name
+            for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
+        ]
+        # TODO: refactor this to use expert_mapping and check all layer numbers
+        all_scheme_dicts = [
+            quant_config.get_scheme_dict(layer, name) for name in unfused_names
+        ]
+        scheme_dict = all_scheme_dicts.pop()
+
+        # multiple schemes found
+        if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]):
+            raise ValueError(
+                "All MoE projections need to have same "
+                "quantization scheme but found multiple"
+            )
+
+        if scheme_dict is None:  # ignored layer
+            return UnquantizedFusedMoEMethod(layer.moe_config)
+
+        # TODO: @dsikka: refactor this to use schemes as other kernels
+        # are supported + check if the layer is being ignored.
+        weight_quant = scheme_dict.get("weights")
+        input_quant = scheme_dict.get("input_activations")
+        format = scheme_dict.get("format")
+
+        if quant_config._is_mxfp4(weight_quant):
+            from .compressed_tensors_moe_w4a4_mxfp4 import (
+                CompressedTensorsW4A4Mxfp4MoEMethod,
+            )
+
+            return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config)
+
+        if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
+            # group_size=None means channelwise
+            group_size = weight_quant.group_size or -1
+
+            valid_format_and_bits = (
+                weight_quant.num_bits in WNA16_SUPPORTED_BITS
+                and format == CompressionFormat.pack_quantized.value
+            )
+
+            if not valid_format_and_bits:
+                raise ValueError(
+                    "For Fused MoE layers, only format: ",
+                    f"{CompressionFormat.pack_quantized.value} ",
+                    f" and bits: {WNA16_SUPPORTED_BITS} is supported ",
+                    f"but got format: {CompressionFormat.pack_quantized.value} "
+                    f" and bits: {weight_quant.num_bits}",
+                )
+
+            # Prefer to use the MarlinMoE kernel when it is supported.
+            if (
+                not check_moe_marlin_supports_layer(layer, group_size)
+                or current_platform.is_rocm()
+            ):
+                from .compressed_tensors_moe_wna16 import (
+                    CompressedTensorsWNA16MoEMethod,
+                )
+
+                if (
+                    weight_quant.strategy == QuantizationStrategy.GROUP
+                    and weight_quant.actorder
+                    in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
+                ):
+                    raise ValueError(
+                        "WNA16MoE is not supported with actorder=group/dynamic."
+                    )
+                logger.info_once("Using CompressedTensorsWNA16MoEMethod")
+                return CompressedTensorsWNA16MoEMethod(
+                    weight_quant, input_quant, layer.moe_config
+                )
+            else:
+                from .compressed_tensors_moe_wna16_marlin import (
+                    CompressedTensorsWNA16MarlinMoEMethod,
+                )
+
+                logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
+                return CompressedTensorsWNA16MarlinMoEMethod(
+                    weight_quant, input_quant, layer.moe_config
+                )
+        elif quant_config._is_nvfp4_format(weight_quant):
+            from .compressed_tensors_moe_w4a4_nvfp4 import (
+                CompressedTensorsW4A4Nvfp4MoEMethod,
+            )
+
+            _is_valid_nvfp4_activations = (
+                quant_config._is_nvfp4_format(input_quant) or input_quant is None
+            )
+            if not _is_valid_nvfp4_activations:
+                raise ValueError(
+                    "For NVFP4 weights, input quantization must also be NVFP4 format ",
+                    f"or None for NVFP4A16, found {input_quant}",
+                )
+            return CompressedTensorsW4A4Nvfp4MoEMethod(
+                layer.moe_config, layer_name, use_a16=(input_quant is None)
+            )
+        elif (
+            quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
+            or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
+            or quant_config._is_fp8_w8a8(weight_quant, input_quant)
+        ):
+            from .compressed_tensors_moe_w8a8_fp8 import (
+                CompressedTensorsW8A8Fp8MoEMethod,
+            )
+
+            return CompressedTensorsW8A8Fp8MoEMethod(
+                weight_quant, input_quant, layer.moe_config
+            )
+        elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
+            from .compressed_tensors_moe_w8a8_int8 import (
+                CompressedTensorsW8A8Int8MoEMethod,
+            )
+
+            return CompressedTensorsW8A8Int8MoEMethod(
+                weight_quant, input_quant, layer.moe_config
+            )
+        elif quant_config._is_fp8_w4a8_sm90(weight_quant, input_quant):
+            from .compressed_tensors_moe_w4a8_fp8 import (
+                CompressedTensorsW4A8Fp8MoEMethod,
+            )
+
+            logger.info_once("Using CompressedTensorsW4A8Fp8MoEMethod")
+            return CompressedTensorsW4A8Fp8MoEMethod(
+                weight_quant, input_quant, layer.moe_config
+            )
+        elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
+            from .compressed_tensors_moe_w4a8_int8 import (
+                CompressedTensorsW4A8Int8MoEMethod,
+            )
+
+            return CompressedTensorsW4A8Int8MoEMethod(
+                weight_quant, input_quant, layer.moe_config
+            )
+        else:
+            raise RuntimeError(
+                f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}"
+            )

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

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoE,
+    FusedMoeWeightScaleSupported,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
+    MarlinExperts,
+)
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+    Mxfp4MoeBackend,
+    make_mxfp4_moe_kernel,
+    make_mxfp4_moe_quant_config,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa E501
+    CompressedTensorsMoEMethod,
+)
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+    prepare_moe_fp4_layer_for_marlin,
+)
+from vllm.model_executor.utils import set_weight_attrs
+
+logger = init_logger(__name__)
+
+
+class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
+    def __init__(self, moe):
+        super().__init__(moe)
+        self.group_size = 32
+        self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
+        self.experts_cls = MarlinExperts
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        layer.num_experts = num_experts
+        layer.params_dtype = params_dtype
+
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                # 2 fp4 items are packed in the input dimension
+                hidden_size // 2,
+                requires_grad=False,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_packed", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                # 2 fp4 items are packed in the input dimension
+                intermediate_size_per_partition // 2,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_packed", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        w13_weight_scale = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                # 2 fp4 items are packed in the input dimension
+                hidden_size // self.group_size,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_scale", w13_weight_scale)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
+        )
+        set_weight_attrs(w13_weight_scale, extra_weight_attrs)
+
+        w2_weight_scale = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                # 2 fp4 items are packed in the input dimension
+                intermediate_size_per_partition // self.group_size,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_scale", w2_weight_scale)
+        set_weight_attrs(w2_weight_scale, extra_weight_attrs)
+
+    def get_fused_moe_quant_config(
+        self, layer: torch.nn.Module
+    ) -> FusedMoEQuantConfig | None:
+        return make_mxfp4_moe_quant_config(
+            mxfp4_backend=self.mxfp4_backend,
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+        )
+
+    def process_weights_after_loading(self, layer: FusedMoE) -> None:
+        layer.w13_weight = torch.nn.Parameter(
+            layer.w13_weight_packed.data, requires_grad=False
+        )
+        delattr(layer, "w13_weight_packed")
+
+        layer.w2_weight = torch.nn.Parameter(
+            layer.w2_weight_packed.data, requires_grad=False
+        )
+        delattr(layer, "w2_weight_packed")
+
+        logger.warning_once(
+            "Your GPU does not have native support for FP4 computation but "
+            "FP4 quantization is being used. Weight-only FP4 compression "
+            "will be used leveraging the Marlin kernel. This may degrade "
+            "performance for compute-heavy workloads."
+        )
+        prepare_moe_fp4_layer_for_marlin(layer)
+
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        if self.moe_quant_config is not None:
+            self.moe_kernel = make_mxfp4_moe_kernel(
+                moe_quant_config=self.moe_quant_config,
+                moe_config=self.moe,
+                experts_cls=self.experts_cls,
+                mxfp4_backend=self.mxfp4_backend,
+                shared_experts=layer.shared_experts,
+                routing_tables=layer._maybe_init_expert_routing_tables(),
+            )
+
+    def apply(
+        self,
+        layer: FusedMoE,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            topk_weights,
+            topk_ids,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            shared_experts_input=shared_experts_input,
+        )

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

@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoE,
+    FusedMoeWeightScaleSupported,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    convert_to_nvfp4_moe_kernel_format,
+    is_global_sf_supported_for_nvfp4_backend,
+    make_nvfp4_moe_kernel,
+    make_nvfp4_moe_quant_config,
+    select_nvfp4_moe_backend,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa E501
+    CompressedTensorsMoEMethod,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kNvfp4Dynamic,
+    kNvfp4Static,
+)
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+
+logger = init_logger(__name__)
+
+
+class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
+    def __init__(
+        self,
+        moe: FusedMoEConfig,
+        layer_name: str | None = None,
+        use_a16: bool = False,
+    ):
+        super().__init__(moe)
+        self.group_size = 16
+
+        # Select experts implementation.
+        self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
+            config=self.moe,
+            weight_key=kNvfp4Static,
+            activation_key=None if use_a16 else kNvfp4Dynamic,
+        )
+
+        self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
+            self.nvfp4_backend
+        )
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        layer.num_experts = num_experts
+        layer.params_dtype = params_dtype
+        w13_num_shards = 2 if self.moe.is_act_and_mul else 1
+
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                w13_num_shards * intermediate_size_per_partition,
+                # 2 fp4 items are packed in the input dimension
+                hidden_size // 2,
+                requires_grad=False,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_packed", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                # 2 fp4 items are packed in the input dimension
+                intermediate_size_per_partition // 2,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_packed", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # Weight Scales
+        w13_weight_scale = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                w13_num_shards * intermediate_size_per_partition,
+                # 2 fp4 items are packed in the input dimension
+                hidden_size // self.group_size,
+                dtype=torch.float8_e4m3fn,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_scale", w13_weight_scale)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
+        )
+        set_weight_attrs(w13_weight_scale, extra_weight_attrs)
+
+        w2_weight_scale = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                # 2 fp4 items are packed in the input dimension
+                intermediate_size_per_partition // self.group_size,
+                dtype=torch.float8_e4m3fn,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_scale", w2_weight_scale)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
+        )
+        set_weight_attrs(w2_weight_scale, extra_weight_attrs)
+
+        # Weight Global Scales
+        w13_weight_scale_2 = torch.nn.Parameter(
+            torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_global_scale", w13_weight_scale_2)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
+        )
+        set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
+
+        w2_weight_scale_2 = torch.nn.Parameter(
+            torch.empty(num_experts, dtype=torch.float32), requires_grad=False
+        )
+        layer.register_parameter("w2_weight_global_scale", w2_weight_scale_2)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
+        )
+        set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
+
+        # Input Global Scales
+        w13_input_scale = torch.nn.Parameter(
+            torch.empty(num_experts, w13_num_shards, dtype=torch.float32),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_input_global_scale", w13_input_scale)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
+        )
+        set_weight_attrs(w13_input_scale, extra_weight_attrs)
+
+        w2_input_scale = torch.nn.Parameter(
+            torch.empty(num_experts, dtype=torch.float32), requires_grad=False
+        )
+        layer.register_parameter("w2_input_global_scale", w2_input_scale)
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
+        )
+        set_weight_attrs(w2_input_scale, extra_weight_attrs)
+
+    def process_weights_after_loading(self, layer: FusedMoE) -> None:
+        """
+        Convert NVFP4 MoE weights into kernel format and setup the kernel.
+        """
+        # NOTE(rob): wN_weight_packed -> wN_weight is because ModularKernelMethod
+        # requires this naming convention. However, the name change breaks
+        # reloading because the state dict no longer matches disk. Once we
+        # remove MKM, we should revert this change to ensure compatibility.
+        layer.w13_weight = torch.nn.Parameter(
+            layer.w13_weight_packed.data, requires_grad=False
+        )
+        delattr(layer, "w13_weight_packed")
+
+        layer.w2_weight = torch.nn.Parameter(
+            layer.w2_weight_packed.data, requires_grad=False
+        )
+        delattr(layer, "w2_weight_packed")
+
+        # Use a single gscale for w13.
+        if self.moe.is_act_and_mul and not torch.allclose(
+            layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]
+        ):
+            logger.warning_once(
+                "w1_weight_global_scale must match w3_weight_global_scale. "
+                "Accuracy may be affected.",
+            )
+        w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()
+
+        # Shuffle weights into the NvFp4 kernel format.
+        (
+            w13,
+            w13_scale,
+            w13_scale_2,
+            a13_scale,
+            w2,
+            w2_scale,
+            w2_scale_2,
+            a2_scale,
+        ) = convert_to_nvfp4_moe_kernel_format(
+            nvfp4_backend=self.nvfp4_backend,
+            layer=layer,
+            w13=layer.w13_weight,
+            w13_scale=layer.w13_weight_scale,
+            w13_scale_2=(1.0 / w13_weight_global_scale),
+            a13_scale=(1.0 / layer.w13_input_global_scale),
+            w2=layer.w2_weight,
+            w2_scale=layer.w2_weight_scale,
+            w2_scale_2=(1.0 / layer.w2_weight_global_scale),
+            a2_scale=(1.0 / layer.w2_input_global_scale),
+            is_act_and_mul=self.moe.is_act_and_mul,
+        )
+
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w13_weight_scale", w13_scale)
+        replace_parameter(layer, "w2_weight", w2)
+        replace_parameter(layer, "w2_weight_scale", w2_scale)
+        layer.w13_weight_scale_2 = w13_scale_2
+        layer.w2_weight_scale_2 = w2_scale_2
+        layer.w13_input_scale = a13_scale
+        layer.w2_input_scale = a2_scale
+
+        # Setup modular kernel.
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.experts_cls is not None
+        self.moe_kernel = make_nvfp4_moe_kernel(
+            moe_quant_config=self.moe_quant_config,
+            moe_config=self.moe,
+            experts_cls=self.experts_cls,
+            shared_experts=layer.shared_experts,
+            routing_tables=layer._maybe_init_expert_routing_tables(),
+        )
+        self.moe_kernel.fused_experts.process_weights_after_loading(layer)
+
+    def maybe_make_prepare_finalize(
+        self,
+        routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+    ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
+        raise ValueError(
+            f"{self.__class__.__name__} uses the new modular kernel initialization "
+            "logic. This function should not be called."
+        )
+
+    def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
+        return make_nvfp4_moe_quant_config(
+            backend=self.nvfp4_backend,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w13_scale_2=layer.w13_weight_scale_2,
+            w2_scale_2=layer.w2_weight_scale_2,
+            a13_scale=layer.w13_input_scale,
+            a2_scale=layer.w2_input_scale,
+        )
+
+    def apply_monolithic(
+        self,
+        layer: FusedMoE,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply_monolithic(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            router_logits,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            num_expert_group=layer.num_expert_group,
+            topk_group=layer.topk_group,
+            e_score_correction_bias=layer.e_score_correction_bias,
+            routed_scaling_factor=layer.routed_scaling_factor,
+        )
+
+    def apply(
+        self,
+        layer: FusedMoE,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            topk_weights,
+            topk_ids,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            shared_experts_input=shared_experts_input,
+        )