[Model] Add support for BharatGen's Param2MoE model (#38000)

Signed-off-by: bhargav-patel-29 <bhargav.patel@tihiitb.org>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.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

.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/

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. | ✅︎ | ✅︎ |

tests/kernels/core/test_layernorm.py

@@ -7,12 +7,20 @@ 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.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)

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/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/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/spec_decode/test_extract_hidden_states.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 from unittest import mock
 
 import numpy as np
 import pytest
 import torch
+from transformers import CLIPVisionConfig, LlamaConfig, LlavaConfig, PretrainedConfig
 
 from tests.v1.attention.utils import (
     BatchSpec,
@@ -23,6 +25,10 @@ from vllm.config import (
 )
 from vllm.config.load import LoadConfig
 from vllm.platforms import current_platform
+from vllm.transformers_utils.config import get_hf_text_config
+from vllm.transformers_utils.configs.extract_hidden_states import (
+    ExtractHiddenStatesConfig,
+)
 from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
 
@@ -323,3 +329,160 @@ def test_propose_different_layer_counts(num_hidden_layers):
 
     assert draft_tokens.shape == (batch_size, 1)
     assert torch.equal(draft_tokens, sampled_token_ids)
+
+
+# ---------------------------------------------------------------------------
+# VLM / composite config tests for ExtractHiddenStatesConfig
+# ---------------------------------------------------------------------------
+
+
+class _DummyVLMConfig(PretrainedConfig):
+    """Minimal composite config that mimics VLMs like Kimi-K2.5 or LLaVA.
+
+    The text model's parameters (hidden_size, num_attention_heads, …) live
+    exclusively under ``text_config``; the top-level config has none of them.
+    """
+
+    model_type = "test_vlm"
+
+    def __init__(self, text_config: PretrainedConfig, **kwargs):
+        self.text_config = text_config
+        super().__init__(architectures=["LlamaForCausalLM"], **kwargs)
+
+    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:
+        del decoder
+        return self.text_config
+
+
+def test_extract_hidden_states_text_only_config_regression():
+    """Text-only models (no nested text_config) must keep working."""
+    model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
+
+    speculative_config = SpeculativeConfig(
+        target_model_config=model_config,
+        target_parallel_config=ParallelConfig(),
+        method="extract_hidden_states",
+        num_speculative_tokens=1,
+        draft_model_config={
+            "hf_config": {
+                "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4],
+            }
+        },
+    )
+
+    assert speculative_config.draft_model_config is not None
+    # For text-only models, hf_text_config should be the config itself.
+    assert speculative_config.draft_model_config.hf_text_config is (
+        speculative_config.draft_model_config.hf_config
+    )
+    assert (
+        speculative_config.draft_model_config.hf_text_config.num_attention_heads
+        == model_config.hf_text_config.num_attention_heads
+    )
+
+
+def test_extract_hidden_states_config_preserves_vlm_text_config():
+    """A real VLM config (LLaVA) with nested text_config must be preserved."""
+    text_config = LlamaConfig(
+        vocab_size=32000,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=8,
+    )
+    vlm_config = LlavaConfig(
+        vision_config=CLIPVisionConfig(),
+        text_config=text_config,
+    )
+
+    # Precondition: to_dict() flattens the nested config to a plain dict.
+    assert isinstance(vlm_config.to_dict()["text_config"], dict)
+
+    extract_config = ExtractHiddenStatesConfig(
+        vlm_config,
+        eagle_aux_hidden_state_layer_ids=[1, 2],
+    )
+
+    # The fix: text_config is still a PretrainedConfig, not a dict.
+    assert isinstance(extract_config.text_config, LlamaConfig)
+
+    extracted = get_hf_text_config(extract_config)
+    assert extracted is extract_config.text_config
+    assert extracted.num_attention_heads == text_config.num_attention_heads
+    assert extracted.hidden_size == text_config.hidden_size
+
+    # Serialization must still round-trip correctly.
+    serialized = extract_config.to_dict()
+    assert isinstance(serialized["text_config"], dict)
+    assert serialized["text_config"]["num_attention_heads"] == (
+        text_config.num_attention_heads
+    )
+
+    json_str = json.loads(extract_config.to_json_string())
+    assert json_str["text_config"]["num_attention_heads"] == (
+        text_config.num_attention_heads
+    )
+
+
+def test_extract_hidden_states_speculative_config_vlm():
+    """SpeculativeConfig with a VLM target must build without errors."""
+    nested_text_config = LlamaConfig(
+        vocab_size=32000,
+        hidden_size=128,
+        intermediate_size=256,
+        num_hidden_layers=2,
+        num_attention_heads=8,
+    )
+
+    target_model_config = ModelConfig(
+        model=model_dir,
+        runner="generate",
+        max_model_len=100,
+    )
+    # Replace the real text-only config with our composite VLM config.
+    target_model_config.hf_config = _DummyVLMConfig(
+        text_config=nested_text_config,
+    )
+    target_model_config.hf_text_config = nested_text_config
+
+    speculative_config = SpeculativeConfig(
+        target_model_config=target_model_config,
+        target_parallel_config=ParallelConfig(),
+        method="extract_hidden_states",
+        num_speculative_tokens=1,
+        draft_model_config={
+            "hf_config": {
+                "eagle_aux_hidden_state_layer_ids": [1, 2],
+            }
+        },
+    )
+
+    assert speculative_config.draft_model_config is not None
+    assert isinstance(
+        speculative_config.draft_model_config.hf_config.text_config,
+        LlamaConfig,
+    )
+    assert speculative_config.draft_model_config.hf_text_config is (
+        speculative_config.draft_model_config.hf_config.text_config
+    )
+    assert (
+        speculative_config.draft_model_config.hf_text_config.num_attention_heads
+        == nested_text_config.num_attention_heads
+    )
+
+
+def test_extract_hidden_states_config_invalid_text_config():
+    """A nested text_config missing required attrs must still be rejected."""
+    broken_text_config = PretrainedConfig(hidden_size=128)
+    vlm_config = _DummyVLMConfig(text_config=broken_text_config)
+
+    extract_config = ExtractHiddenStatesConfig(
+        vlm_config,
+        eagle_aux_hidden_state_layer_ids=[1],
+    )
+
+    # The object is preserved (not flattened), …
+    assert extract_config.text_config is broken_text_config
+    # … but validation still rejects the missing attribute.
+    with pytest.raises(ValueError, match="num_attention_heads"):
+        get_hf_text_config(extract_config)

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/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/fused_moe/experts/trtllm_fp8_moe.py

@@ -112,6 +112,24 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A
 
+    def moe_problem_size(
+        self,
+        a1: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> tuple[int, int, int, int, int]:
+        """Override to handle 4D BlockMajorK weights (E, K/bk, Mn, bk)."""
+        if w1.dim() == 4:
+            # BlockMajorK: (E, K/bk, Mn, bk)
+            E = w1.shape[0]
+            N = w1.shape[2]
+            K = a1.size(-1)
+            M = a1.size(0) if a1.dim() == 2 else a1.size(1)
+            topk = topk_ids.size(1)
+            return E, M, N, K, topk
+        return super().moe_problem_size(a1, w1, w2, topk_ids)
+
     def workspace_shapes(
         self,
         M: int,
@@ -152,7 +170,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
         apply_router_weight_on_input: bool,
     ):
         import flashinfer
-        from flashinfer.fused_moe import Fp8QuantizationType
+        from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
 
         # Pack topk ids and weights into format expected by the kernel.
         packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
@@ -170,10 +188,12 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
         if is_mxfp8:
             fp8_quant_type = Fp8QuantizationType.MxFp8
             use_shuffled_weight = True
+            weight_layout = WeightLayout.MajorK
             hidden_states_scale = a1q_scale
         else:
             fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
-            use_shuffled_weight = False
+            use_shuffled_weight = True
+            weight_layout = WeightLayout.BlockMajorK
             hidden_states_scale = a1q_scale.t().contiguous()
 
         # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the
@@ -199,7 +219,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             routed_scaling_factor=None,
             routing_method_type=1,
             use_shuffled_weight=use_shuffled_weight,
-            weight_layout=0,
+            weight_layout=weight_layout,
             fp8_quantization_type=fp8_quant_type,
             # output=output,
         )
@@ -322,7 +342,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
         topk_group: int | None = None,
     ) -> torch.Tensor:
         import flashinfer
-        from flashinfer.fused_moe import Fp8QuantizationType
+        from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
 
         assert not apply_router_weight_on_input
         assert activation == MoEActivation.SILU
@@ -342,10 +362,12 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
         if is_mxfp8:
             fp8_quant_type = Fp8QuantizationType.MxFp8
             use_shuffled_weight = True
+            weight_layout = WeightLayout.MajorK
             hidden_states_scale = a1q_scale
         else:
             fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
-            use_shuffled_weight = False
+            use_shuffled_weight = True
+            weight_layout = WeightLayout.BlockMajorK
             hidden_states_scale = a1q_scale.t().contiguous()
 
         return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(
@@ -367,6 +389,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
             use_shuffled_weight=use_shuffled_weight,
+            weight_layout=weight_layout,
             fp8_quantization_type=fp8_quant_type,
         )
 

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

@@ -93,24 +93,24 @@ class SharedExperts:
                 )
 
     @property
-    def _has_external_experts(self) -> bool:
+    def _use_external_experts(self) -> bool:
+        if self._use_dp_chunking:
+            return False
+
         # Disable shared expert overlap if:
         #   - we are using eplb with non-default backend, because of correctness issues
         #   - we are using flashinfer with DP, since there nothing to gain
         backend = self._moe_config.moe_parallel_config.all2all_backend
-        return not (
-            (
-                self._moe_config.moe_parallel_config.enable_eplb
-                and backend != "allgather_reducescatter"
-            )
-            or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels
-        )
+        return (
+            self._moe_config.moe_parallel_config.enable_eplb
+            and backend != "allgather_reducescatter"
+        ) or self._moe_config.moe_parallel_config.use_fi_nvl_two_sided_kernels
 
     def _determine_shared_experts_order(
         self,
         hidden_states: torch.Tensor,
     ) -> SharedExpertsOrder:
-        if self._has_external_experts and not self._use_dp_chunking:
+        if self._use_external_experts:
             return SharedExpertsOrder.EXTERNAL
 
         if self._quant_method.mk_owns_shared_expert:

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

@@ -305,6 +305,39 @@ def align_fp8_moe_weights_for_fi(
     return padded_w13, padded_w2, padded_intermediate
 
 
+def _shuffle_deepseek_fp8_moe_weights(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Preprocess DeepSeek FP8 block-scale weights for the FlashInfer TRT-LLM
+    kernel using the shuffle + BlockMajorK layout variant.
+
+    Returns 4D weight tensors in BlockMajorK layout
+    (E, K/block_k, Mn, block_k)
+    """
+    from flashinfer import shuffle_matrix_a
+    from flashinfer.fused_moe import convert_to_block_layout
+
+    epilogue_tile_m = 64
+    block_k = 128
+    num_experts = w13.shape[0]
+
+    w13_shuffled: list[torch.Tensor] = []
+    w2_shuffled: list[torch.Tensor] = []
+    for i in range(num_experts):
+        t13 = shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)
+        t13 = convert_to_block_layout(t13, block_k)
+        w13_shuffled.append(t13)
+
+        t2 = shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)
+        t2 = convert_to_block_layout(t2, block_k)
+        w2_shuffled.append(t2)
+
+    w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn)
+    w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn)
+    return w13_out, w2_out
+
+
 def _shuffle_mxfp8_moe_weights(
     w13: torch.Tensor,
     w2: torch.Tensor,
@@ -405,6 +438,7 @@ def prepare_fp8_moe_layer_for_fi(
         hasattr(layer, "weight_block_size") and layer.weight_block_size is not None
     )
     is_mxfp8 = block_quant and w13_scale.dtype == torch.uint8
+    is_deepseek_fp8 = block_quant and not is_mxfp8
     is_gated = layer.activation.is_gated
 
     # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle.
@@ -447,6 +481,10 @@ def prepare_fp8_moe_layer_for_fi(
         if block_quant:
             w13_scale = swap_w13_to_w31(w13_scale)
 
+    # DeepSeekFp8 TRT-LLM: shuffle weights into BlockMajorK layout.
+    if is_deepseek_fp8 and is_trtllm:
+        w13, w2 = _shuffle_deepseek_fp8_moe_weights(w13, w2)
+
     # FI TRT-LLM FP8 per-tensor MoE kernel requires weight shuffle
     # and registration of alpha scales.
     if is_trtllm and not block_quant:

vllm/model_executor/models/interfaces.py

@@ -46,7 +46,7 @@ if TYPE_CHECKING:
     from vllm.multimodal.inputs import MultiModalFeatureSpec
     from vllm.multimodal.registry import _ProcessorFactories
     from vllm.sequence import IntermediateTensors
-    from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
+    from vllm.v1.worker.encoder_cudagraph_defs import (
         EncoderCudaGraphCaptureInputs,
         EncoderCudaGraphConfig,
         EncoderCudaGraphReplayBuffers,

vllm/model_executor/models/minimax_m2.py

@@ -24,6 +24,7 @@
 """Inference-only MiniMaxM2 model."""
 
 from collections.abc import Iterable
+from itertools import islice
 from typing import Any
 
 import torch
@@ -59,7 +60,7 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 from vllm.sequence import IntermediateTensors
 
-from .interfaces import SupportsLoRA, SupportsPP
+from .interfaces import EagleModelMixin, SupportsEagle3, SupportsLoRA, SupportsPP
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -313,7 +314,7 @@ class MiniMaxM2DecoderLayer(nn.Module):
 
 
 @support_torch_compile
-class MiniMaxM2Model(nn.Module):
+class MiniMaxM2Model(nn.Module, EagleModelMixin):
     fall_back_to_pt_during_load = False
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -366,7 +367,7 @@ class MiniMaxM2Model(nn.Module):
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None,
         inputs_embeds: torch.Tensor | None = None,
-    ) -> torch.Tensor | IntermediateTensors:
+    ) -> torch.Tensor | IntermediateTensors | tuple[torch.Tensor, list[torch.Tensor]]:
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
                 hidden_states = inputs_embeds
@@ -378,14 +379,24 @@ class MiniMaxM2Model(nn.Module):
             hidden_states = intermediate_tensors["hidden_states"]
             residual = intermediate_tensors["residual"]
 
-        for layer in self.layers[self.start_layer : self.end_layer]:
+        aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
+        for idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer)
+        ):
             hidden_states, residual = layer(positions, hidden_states, residual)
+            self._maybe_add_hidden_state(
+                aux_hidden_states, idx + 1, hidden_states, residual
+            )
 
         if not get_pp_group().is_last_rank:
             return IntermediateTensors(
                 {"hidden_states": hidden_states, "residual": residual}
             )
         hidden_states, _ = self.norm(hidden_states, residual)
+
+        if len(aux_hidden_states) > 0:
+            return hidden_states, aux_hidden_states
+
         return hidden_states
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
@@ -496,7 +507,7 @@ class MiniMaxM2Model(nn.Module):
         return loaded_params
 
 
-class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
+class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",

vllm/model_executor/models/nano_nemotron_vl.py

@@ -1239,12 +1239,13 @@ class NemotronH_Nano_VL_V2(
             img_context_token_ids=self._img_context_token_ids,
             video_temporal_patch_size=video_temporal_patch_size,
         )
+        device = video_embeddings.device
 
         # video_repl.full is a list of token IDs
-        repl_token_ids = torch.tensor(video_repl.full)
+        repl_token_ids = torch.tensor(video_repl.full, device=device)
 
         # Get embedding token IDs for image context (use pre-tokenized version)
-        embed_token_ids = torch.tensor(self._img_context_token_ids)
+        embed_token_ids = torch.tensor(self._img_context_token_ids, device=device)
 
         # Create mask for video embedding positions
         is_video_embed = torch.isin(repl_token_ids, embed_token_ids)

vllm/model_executor/models/param2moe.py

@@ -0,0 +1,900 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Copyright 2026 BharatGen AI team. All rights reserved.
+#
+# This code has been modified to accommodate Param2MoE's GQA-based MoE architecture.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# limitations under the License.
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from itertools import islice
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_world_size,
+)
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.fused_moe import SharedFusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.sequence import IntermediateTensors
+
+from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
+from .utils import (
+    AutoWeightsLoader,
+    PPMissingLayer,
+    is_pp_missing_parameter,
+    make_empty_intermediate_tensors_factory,
+    make_layers,
+    maybe_prefix,
+)
+
+
+def _is_expert_bias_name(name: str) -> bool:
+    """True when the weight is the MoE router's per-expert score bias."""
+    return name.endswith(".mlp.gate.expert_bias")
+
+
+def _zero_mean_tensor(t: torch.Tensor) -> torch.Tensor:
+    if t.numel() == 0:
+        return t
+    return t - t.mean()
+
+
+def _rename_and_normalize_weights(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """
+    Translate HuggingFace Param2MoE weight names to vLLM internal names
+    and zero-mean the expert-bias tensor so the router stays balanced.
+
+    Mapping table (HF → vLLM):
+      model.word_embeddings.*              → model.embed_tokens.*
+      *.attention.query_key_value.*        → *.self_attn.qkv_proj.*
+      *.attention.dense.*                  → *.self_attn.o_proj.*
+      *.attention.query_layernorm.*        → *.self_attn.q_layernorm.*
+      *.attention.key_layernorm.*          → *.self_attn.k_layernorm.*
+      *.mlp.gate.expert_bias               → *.mlp.gate.e_score_correction_bias
+        (also zero-meant for load balance)
+    """
+    for name, w in weights:
+        # Embedding table
+        name = name.replace("model.word_embeddings.", "model.embed_tokens.")
+        # Fused QKV projection  (HF: query_key_value → vLLM: qkv_proj)
+        name = name.replace(".attention.query_key_value.", ".self_attn.qkv_proj.")
+        # Output projection  (HF: dense → vLLM: o_proj)
+        name = name.replace(".attention.dense.", ".self_attn.o_proj.")
+        # Per-head query norm
+        name = name.replace(".attention.query_layernorm.", ".self_attn.q_layernorm.")
+        # Per-head key norm
+        name = name.replace(".attention.key_layernorm.", ".self_attn.k_layernorm.")
+        # Catch any remaining .attention. → .self_attn. prefixes
+        # (e.g. future bias params on the projection layers)
+        name = name.replace(".attention.", ".self_attn.")
+
+        # Expert-score bias: rename + zero-mean
+        if name.endswith(".mlp.gate.expert_bias"):
+            name = name.replace(
+                ".mlp.gate.expert_bias",
+                ".mlp.gate.e_score_correction_bias",
+            )
+            w = _zero_mean_tensor(w)
+
+        yield name, w
+
+
+class Param2MoEAttention(nn.Module):
+    """
+    Grouped-Query Attention (GQA) for Param2MoE.
+
+    Notable differences from a vanilla GQA layer:
+      * The checkpoint fuses Q, K, V into a single ``query_key_value`` weight.
+        vLLM receives it already renamed to ``qkv_proj`` by the weight-name
+        translator and splits it during ``load_weights``.
+      * Optional per-head RMS norms on Q and K (``use_qk_norm=True``).
+    """
+
+    def __init__(
+        self,
+        config,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim or (self.hidden_size // self.num_heads)
+        self.use_qk_norm: bool = getattr(config, "use_qk_norm", False)
+
+        tp_size = get_tensor_model_parallel_world_size()
+        assert self.num_heads % tp_size == 0, (
+            f"num_attention_heads ({self.num_heads}) must be divisible "
+            f"by tensor-parallel world size ({tp_size})."
+        )
+        assert self.num_kv_heads % tp_size == 0, (
+            f"num_key_value_heads ({self.num_kv_heads}) must be divisible "
+            f"by tensor-parallel world size ({tp_size})."
+        )
+        self.num_local_heads = self.num_heads // tp_size
+        self.num_local_kv_heads = self.num_kv_heads // tp_size
+
+        # Sizes after TP split (used in forward to split qkv output)
+        self.q_size_local = self.num_local_heads * self.head_dim
+        self.kv_size_local = self.num_local_kv_heads * self.head_dim
+
+        self.scaling = self.head_dim**-0.5
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            total_num_kv_heads=self.num_kv_heads,
+            bias=getattr(config, "use_qkv_bias", False),
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+
+        self.o_proj = RowParallelLinear(
+            input_size=self.num_heads * self.head_dim,
+            output_size=self.hidden_size,
+            bias=getattr(config, "use_bias", False),
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        if self.use_qk_norm:
+            self.q_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.k_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+        # `partial_rotary_factor` defaults to 1.0 (full RoPE) if not in config
+        partial_rotary_factor: float = getattr(config, "partial_rotary_factor", 1.0)
+        rope_dim = int(self.head_dim * partial_rotary_factor)
+
+        rope_parameters: dict = {
+            "rope_type": "default",
+            "base": config.rope_theta,
+        }
+        if config.rope_scaling is not None:
+            rope_parameters.update(config.rope_scaling)
+            # Normalise key: some checkpoints use "type", vLLM wants "rope_type"
+            if "type" in rope_parameters and "rope_type" not in rope_parameters:
+                rope_parameters["rope_type"] = rope_parameters.pop("type")
+
+        self.rotary_emb = get_rope(
+            rope_dim,
+            max_position=config.max_position_embeddings,
+            rope_parameters=rope_parameters,
+            is_neox_style=True,
+        )
+
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            scale=self.scaling,
+            num_kv_heads=self.num_kv_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        # 1. Fused QKV projection → split into local Q / K / V
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split(
+            [self.q_size_local, self.kv_size_local, self.kv_size_local],
+            dim=-1,
+        )
+
+        # 2. Optional per-head QK norms
+        #    Reshape to (T, num_local_heads, head_dim), norm, reshape back.
+        if self.use_qk_norm:
+            T = q.shape[0]
+            q = self.q_layernorm(q.view(T, self.num_local_heads, self.head_dim)).view(
+                T, self.q_size_local
+            )
+            k = self.k_layernorm(
+                k.view(T, self.num_local_kv_heads, self.head_dim)
+            ).view(T, self.kv_size_local)
+
+        # 3. Rotary position embeddings
+        q, k = self.rotary_emb(positions, q, k)
+
+        # 4. Paged attention
+        attn_output = self.attn(q, k, v)
+
+        # 5. Output projection
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Param2MoEMLP(nn.Module):
+    """SwiGLU feed-forward block used for dense layers."""
+
+    def __init__(
+        self,
+        intermediate_size: int,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=config.hidden_size,
+            output_sizes=[intermediate_size, intermediate_size],
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=config.hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            reduce_results=reduce_results,
+            prefix=f"{prefix}.down_proj",
+        )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate_up, _ = self.gate_up_proj(x)
+        x = self.act_fn(gate_up)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Param2MoEMoEBlock(nn.Module):
+    """
+    Mixture-of-Experts block for Param2MoE.
+
+    Routing:
+      * Sigmoid scoring  (config.score_function = "sigmoid")
+      * Grouped top-k   (n_group, topk_group)
+      * Per-expert bias  (gate.expert_bias → e_score_correction_bias)
+      * routed_scaling_factor normalisation
+
+    One set of shared (always-active) experts is added on top.
+    """
+
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        self.config = config
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.hidden_size = config.hidden_size
+
+        self.num_experts: int = config.num_experts
+        self.top_k: int = config.num_experts_per_tok
+        self.routed_scaling_factor: float = getattr(
+            config, "routed_scaling_factor", 1.0
+        )
+
+        self.n_group: int | None = getattr(config, "n_group", None)
+        self.topk_group: int | None = getattr(config, "topk_group", None)
+        self.use_grouped_topk: bool = (
+            self.n_group is not None and self.topk_group is not None
+        )
+
+        self.norm_expert_prob: bool = getattr(config, "norm_topk_prob", True)
+        self.score_function: str = getattr(config, "score_function", "sigmoid")
+
+        self.gate = nn.Linear(
+            self.hidden_size,
+            self.num_experts,
+            bias=False,
+        )
+
+        if getattr(config, "moe_router_enable_expert_bias", True):
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.zeros(self.num_experts, dtype=torch.float32)
+            )
+        else:
+            self.gate.e_score_correction_bias = None  # type: ignore[assignment]
+
+        self.num_shared_experts: int = getattr(config, "num_shared_experts", 1)
+        if self.num_shared_experts > 0:
+            # If moe_shared_expert_intermediate_size is present in the config
+            # it already encodes the TOTAL intermediate size across all shared
+            # experts (i.e. it equals moe_intermediate_size * num_shared_experts).
+            # Do NOT multiply again.  Fall back to computing the product only
+            # when the dedicated field is absent.
+            if (
+                hasattr(config, "moe_shared_expert_intermediate_size")
+                and config.moe_shared_expert_intermediate_size is not None
+            ):
+                shared_int: int = config.moe_shared_expert_intermediate_size
+            else:
+                shared_int = config.moe_intermediate_size * self.num_shared_experts
+            self.shared_experts = Param2MoEMLP(
+                intermediate_size=shared_int,
+                config=config,
+                quant_config=quant_config,
+                reduce_results=False,
+                prefix=f"{prefix}.shared_experts",
+            )
+        else:
+            self.shared_experts = None  # type: ignore[assignment]
+
+        self.experts = SharedFusedMoE(
+            shared_experts=self.shared_experts,
+            num_experts=self.num_experts,
+            top_k=self.top_k,
+            hidden_size=self.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            reduce_results=False,
+            renormalize=self.norm_expert_prob,
+            quant_config=quant_config,
+            prefix=f"{prefix}.experts",
+            scoring_func=self.score_function,
+            e_score_correction_bias=self.gate.e_score_correction_bias,
+            num_expert_group=self.n_group,
+            topk_group=self.topk_group,
+            use_grouped_topk=self.use_grouped_topk,
+            routed_scaling_factor=self.routed_scaling_factor,
+        )
+
+    def maybe_get_fused_moe(self) -> SharedFusedMoE:
+        return self.experts
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+
+        # Router: both input and weight must be float32 for numerical
+        # stability (mirrors the original Param2MoEGate behaviour).
+        # The gate nn.Linear weight lives in the model dtype (bfloat16),
+        # so we must cast both explicitly via F.linear instead of calling
+        # self.gate() which would hit a dtype mismatch.
+        router_logits = F.linear(
+            hidden_states.float(),
+            self.gate.weight.float(),
+        ).to(hidden_states.dtype)
+
+        final_hidden = self.experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+        )
+
+        if self.shared_experts is not None:
+            shared_output, expert_output = final_hidden
+        else:
+            shared_output, expert_output = None, final_hidden
+
+        if shared_output is not None:
+            expert_output = expert_output + shared_output
+
+        if self.tp_size > 1:
+            expert_output = self.experts.maybe_all_reduce_tensor_model_parallel(
+                expert_output
+            )
+
+        return expert_output.view(num_tokens, hidden_dim)
+
+
+class Param2MoEDecoderLayer(nn.Module):
+    """
+    Single transformer decoder block.
+
+    Dense for the first ``first_k_dense_replace`` layers; MoE thereafter.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+
+        hidden_size = config.hidden_size
+        # Derive the layer index from the prefix (e.g. "model.layers.3")
+        layer_idx = int(prefix.split(".")[-1])
+
+        self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = Param2MoEAttention(
+            config=config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps)
+
+        first_k_dense: int = getattr(config, "first_k_dense_replace", 1)
+        is_moe_layer = config.num_experts is not None and layer_idx >= first_k_dense
+
+        if is_moe_layer:
+            self.mlp = Param2MoEMoEBlock(
+                config=config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = Param2MoEMLP(  # type: ignore[assignment]
+                intermediate_size=config.intermediate_size,
+                config=config,
+                quant_config=quant_config,
+                reduce_results=True,
+                prefix=f"{prefix}.mlp",
+            )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        positions: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Pre-norm + attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        # Pre-norm + MLP
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class Param2MoEModel(nn.Module):
+    def __init__(
+        self,
+        *,
+        vllm_config: VllmConfig,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.embed_dim = config.hidden_size
+        self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False)
+
+        # Embedding  (HF name: word_embeddings → vLLM name: embed_tokens)
+        if get_pp_group().is_first_rank or (
+            self.tie_word_embeddings and get_pp_group().is_last_rank
+        ):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.vocab_size,
+                self.embed_dim,
+                quant_config=quant_config,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            lambda prefix: Param2MoEDecoderLayer(
+                vllm_config=vllm_config,
+                prefix=prefix,
+            ),
+            prefix=f"{prefix}.layers",
+        )
+
+        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+            ["hidden_states", "residual"], config.hidden_size
+        )
+
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(self.embed_dim, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                hidden_states = self.embed_input_ids(input_ids)
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+
+        for layer in islice(self.layers, self.start_layer, self.end_layer):
+            hidden_states, residual = layer(hidden_states, positions, residual)
+
+        if not get_pp_group().is_last_rank:
+            return IntermediateTensors(
+                {"hidden_states": hidden_states, "residual": residual}
+            )
+
+        if residual is None:
+            hidden_states = self.norm(hidden_states)
+        else:
+            hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        """
+        Custom weight loader for the inner Param2MoEModel.
+
+        Receives weights that have already been renamed/normalised by the
+        outer model and whose ``model.`` prefix has been stripped by
+        ``AutoWeightsLoader``.  Handles:
+          1. Fused QKV split (query_key_value → qkv_proj q/k/v shards).
+          2. gate_proj + up_proj → gate_up_proj stacking (dense + shared-exp).
+          3. Routed-expert weights via the fused-MoE mapping.
+          4. All remaining weights via their default loader.
+        """
+        config = self.config
+        num_heads: int = config.num_attention_heads
+        num_kv_heads: int = config.num_key_value_heads
+        head_dim: int = config.head_dim or (config.hidden_size // num_heads)
+        q_split = num_heads * head_dim
+        kv_split = num_kv_heads * head_dim
+
+        stacked_params_mapping = [
+            # (vllm_param_name, ckpt_weight_name, shard_id)
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+        expert_params_mapping = self.get_expert_mapping()
+
+        for name, loaded_weight in weights:
+            # ------------------------------------------------------------------
+            # 1. Fused QKV: split into q / k / v shards for QKVParallelLinear
+            # ------------------------------------------------------------------
+            if name.endswith(".self_attn.qkv_proj.weight"):
+                if name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                q_w = loaded_weight[:q_split, :]
+                k_w = loaded_weight[q_split : q_split + kv_split, :]
+                v_w = loaded_weight[q_split + kv_split :, :]
+                weight_loader(param, q_w, "q")
+                weight_loader(param, k_w, "k")
+                weight_loader(param, v_w, "v")
+                loaded_params.add(name)
+                continue
+
+            # ------------------------------------------------------------------
+            # 2. gate_proj / up_proj → gate_up_proj (dense MLP + shared-exp.)
+            # ------------------------------------------------------------------
+            matched_stacked = False
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                if "mlp.experts" in name:  # routed experts handled below
+                    continue
+                new_name = name.replace(weight_name, param_name)
+                if new_name.endswith(".bias") and new_name not in params_dict:
+                    continue
+                if new_name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(new_name, self):
+                    continue
+
+                param = params_dict[new_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight, shard_id)
+                loaded_params.add(new_name)
+                matched_stacked = True
+                break
+
+            if matched_stacked:
+                continue
+
+            # ------------------------------------------------------------------
+            # 3. Routed expert weights → fused-MoE kernel layout
+            # ------------------------------------------------------------------
+            matched_expert = False
+            for (
+                param_name,
+                weight_name,
+                expert_id,
+                shard_id,
+            ) in expert_params_mapping:
+                if weight_name not in name:
+                    continue
+                new_name = name.replace(weight_name, param_name)
+                if is_pp_missing_parameter(new_name, self):
+                    continue
+                if new_name not in params_dict:
+                    continue
+
+                param = params_dict[new_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(
+                    param,
+                    loaded_weight,
+                    name,
+                    shard_id=shard_id,
+                    expert_id=expert_id,
+                )
+                loaded_params.add(new_name)
+                matched_expert = True
+                break
+
+            if matched_expert:
+                continue
+
+            # ------------------------------------------------------------------
+            # 4. All other weights: direct load (layernorms, embed_tokens, …)
+            # ------------------------------------------------------------------
+            if name.endswith(".bias") and name not in params_dict:
+                continue
+            if name not in params_dict:
+                continue
+            if is_pp_missing_parameter(name, self):
+                continue
+
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            try:
+                weight_loader(param, loaded_weight)
+            except Exception as e:
+                raise RuntimeError(
+                    f"[param2moe] Failed to load weight '{name}' "
+                    f"with shape {tuple(loaded_weight.shape)} "
+                    f"into param type {type(param).__name__}: {e}"
+                ) from e
+            loaded_params.add(name)
+
+        return loaded_params
+
+    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
+        return SharedFusedMoE.make_expert_params_mapping(
+            self,
+            ckpt_gate_proj_name="gate_proj",
+            ckpt_down_proj_name="down_proj",
+            ckpt_up_proj_name="up_proj",
+            num_experts=self.config.num_experts,
+        )
+
+
+class Param2MoEMixtureOfExperts(MixtureOfExperts):
+    """Implements the vLLM MixtureOfExperts protocol for Param2MoE."""
+
+    expert_weights: list[torch.Tensor]
+
+    def extract_moe_parameters(self, example_moe: Param2MoEMoEBlock | None) -> None:
+        if example_moe is None:
+            raise RuntimeError(
+                "No Param2MoEMoEBlock found in model.layers. "
+                "Check first_k_dense_replace and num_experts in config."
+            )
+        self.num_logical_experts = example_moe.num_experts
+        self.num_routed_experts = example_moe.num_experts
+        self.num_shared_experts = example_moe.num_shared_experts
+
+        self.num_physical_experts = self.num_logical_experts
+        self.num_local_physical_experts = self.num_logical_experts
+        self.num_redundant_experts = 0
+
+    def update_physical_experts_metadata(
+        self,
+        num_physical_experts: int,
+        num_local_physical_experts: int,
+    ) -> None:
+        self.num_physical_experts = num_physical_experts
+        self.num_local_physical_experts = num_local_physical_experts
+        self.num_redundant_experts = num_physical_experts - self.num_logical_experts
+
+        for moe in self.moe_mlp_layers:
+            moe.n_physical_experts = num_physical_experts
+            moe.n_local_physical_experts = num_local_physical_experts
+            moe.n_redundant_experts = self.num_redundant_experts
+
+            fused = moe.experts
+            if hasattr(fused, "n_local_physical_experts"):
+                fused.n_local_physical_experts = num_local_physical_experts
+            if hasattr(fused, "n_physical_experts"):
+                fused.n_physical_experts = num_physical_experts
+            if hasattr(fused, "n_redundant_experts"):
+                fused.n_redundant_experts = self.num_redundant_experts
+            if hasattr(fused, "update_expert_map"):
+                fused.update_expert_map()
+
+    def set_eplb_state(
+        self,
+        expert_load_view: torch.Tensor,
+        logical_to_physical_map: torch.Tensor,
+        logical_replica_count: torch.Tensor,
+    ) -> None:
+        self.expert_weights.clear()
+        for layer_idx, layer in enumerate(self.moe_layers):
+            if hasattr(layer, "get_expert_weights"):
+                self.expert_weights.append(layer.get_expert_weights())
+            if hasattr(layer, "set_eplb_state"):
+                layer.set_eplb_state(
+                    moe_layer_idx=layer_idx,
+                    expert_load_view=expert_load_view,
+                    logical_to_physical_map=logical_to_physical_map,
+                    logical_replica_count=logical_replica_count,
+                )
+
+
+class Param2MoEForCausalLM(
+    nn.Module, SupportsPP, SupportsLoRA, Param2MoEMixtureOfExperts
+):
+    """
+    vLLM-native Param2MoE CausalLM.
+
+    Uses Grouped-Query Attention (GQA) with a Sigmoid-scored,
+    grouped-topk Mixture-of-Experts MLP.
+    """
+
+    # LoRA packed-module mapping. The fused gate_up_proj handles
+    # gate_proj and up_proj from the checkpoint.
+    packed_modules_mapping = {
+        "qkv_proj": ["query_key_value"],
+        "gate_up_proj": ["gate_proj", "up_proj"],
+    }
+
+    # Modules eligible for LoRA adaptation.
+    supported_lora_modules = [
+        "qkv_proj",
+        "o_proj",
+        "gate_up_proj",
+        "down_proj",
+    ]
+
+    # Embedding layers and their weight-tying counterparts.
+    embedding_modules = {
+        "embed_tokens": "input_embeddings",
+        "lm_head": "output_embeddings",
+    }
+
+    # Modules that need vocab-size padding for LoRA.
+    embedding_padding_modules = ["lm_head"]
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.quant_config = quant_config
+
+        self.model = Param2MoEModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+        )
+
+        self.tie_word_embeddings: bool = getattr(config, "tie_word_embeddings", False)
+        if get_pp_group().is_last_rank:
+            if self.tie_word_embeddings:
+                self.lm_head = self.model.embed_tokens
+            else:
+                self.lm_head = ParallelLMHead(
+                    config.vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=maybe_prefix(prefix, "lm_head"),
+                )
+            self.logits_processor = LogitsProcessor(config.vocab_size)
+        else:
+            self.lm_head = PPMissingLayer()
+            self.logits_processor = None  # type: ignore[assignment]
+
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+        self.expert_weights: list[torch.Tensor] = []
+        self.num_moe_layers: int = 0
+        self.moe_layers: list = []
+        self.moe_mlp_layers: list = []
+
+        example_moe: Param2MoEMoEBlock | None = None
+        for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+            if isinstance(layer.mlp, Param2MoEMoEBlock):
+                example_moe = layer.mlp
+                self.moe_mlp_layers.append(layer.mlp)
+                self.moe_layers.append(layer.mlp.experts)
+                self.num_moe_layers += 1
+
+        if self.config.num_experts is not None:
+            self.extract_moe_parameters(example_moe)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if not get_pp_group().is_last_rank:
+            return None
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(_rename_and_normalize_weights(weights))

vllm/model_executor/models/qwen3_vl.py

@@ -1733,7 +1733,7 @@ class Qwen3VLForConditionalGeneration(
     # -- SupportsEncoderCudaGraph protocol methods --
 
     def get_encoder_cudagraph_config(self):
-        from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
+        from vllm.v1.worker.encoder_cudagraph_defs import (
             EncoderCudaGraphConfig,
         )
 
@@ -1818,7 +1818,7 @@ class Qwen3VLForConditionalGeneration(
         device: torch.device,
         dtype: torch.dtype,
     ):
-        from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
+        from vllm.v1.worker.encoder_cudagraph_defs import (
             EncoderCudaGraphCaptureInputs,
         )
 
@@ -1872,7 +1872,7 @@ class Qwen3VLForConditionalGeneration(
         mm_kwargs: dict[str, Any],
         max_batch_size: int,
     ):
-        from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
+        from vllm.v1.worker.encoder_cudagraph_defs import (
             EncoderCudaGraphReplayBuffers,
         )
 

vllm/model_executor/models/registry.py

@@ -182,6 +182,7 @@ _TEXT_GENERATION_MODELS = {
     "PanguEmbeddedForCausalLM": ("openpangu", "PanguEmbeddedForCausalLM"),
     "PanguProMoEV2ForCausalLM": ("openpangu", "PanguProMoEV2ForCausalLM"),
     "PanguUltraMoEForCausalLM": ("openpangu", "PanguUltraMoEForCausalLM"),
+    "Param2MoEForCausalLM": ("param2moe", "Param2MoEForCausalLM"),
     "PersimmonForCausalLM": ("persimmon", "PersimmonForCausalLM"),
     "PhiForCausalLM": ("phi", "PhiForCausalLM"),
     "Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
@@ -554,6 +555,7 @@ _SPECULATIVE_DECODING_MODELS = {
     "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
     "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
     "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
+    "Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

vllm/model_executor/warmup/deep_gemm_warmup.py

@@ -13,7 +13,7 @@ import vllm.envs as envs
 from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank
 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M
-from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEModularMethod
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
     TritonOrDeepGemmExperts,
 )
@@ -168,14 +168,12 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
     ):
         return False
 
-    if not isinstance(module.quant_method, FusedMoEModularMethod):
-        # modular kernels could invoke deep_gemm_moe_fp8
-        return True
+    moe_kernel = getattr(module.quant_method, "moe_kernel", None)
+    if moe_kernel is None:
+        return False
 
-    # Further check if the ModularKernel implementation uses the DeepGemmExperts
-    return isinstance(
-        module.quant_method.moe_kernel, (DeepGemmExperts, TritonOrDeepGemmExperts)
-    )
+    fused_experts = moe_kernel.impl.fused_experts
+    return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts))
 
 
 FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set()

vllm/platforms/rocm.py

@@ -182,6 +182,7 @@ _ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"])
 _ON_GFX12X = any(arch in _GCN_ARCH for arch in ["gfx12"])
 _ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"])
 _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
+_ON_GFX90A = "gfx90a" in _GCN_ARCH
 _ON_GFX942 = "gfx942" in _GCN_ARCH
 _ON_GFX950 = "gfx950" in _GCN_ARCH
 
@@ -273,6 +274,10 @@ def on_gfx9() -> bool:
     return _ON_GFX9
 
 
+def on_gfx90a() -> bool:
+    return _ON_GFX90A
+
+
 def on_gfx942() -> bool:
     return _ON_GFX942
 

vllm/tool_parsers/gemma4_tool_parser.py

@@ -675,10 +675,11 @@ class Gemma4ToolParser(ToolParser):
         current_args_json = json.dumps(current_args, ensure_ascii=False)
 
         # Withhold trailing closing characters that may shift as more
-        # tokens arrive. Strip trailing '}', '"', and ']' sequences
-        # to get the "safe prefix".
+        # tokens arrive. Strip trailing '}', '"', ']' and partial
+        # STRING_DELIM fragments ('<', '|', '\\', '>') to get the
+        # "safe prefix".
         safe_json = current_args_json
-        while safe_json and safe_json[-1] in ("}", '"', "]"):
+        while safe_json and safe_json[-1] in ("}", '"', "]", "<", "|", "\\", ">"):
             safe_json = safe_json[:-1]
 
         prev_streamed = self.streamed_args_for_tool[self.current_tool_id]

vllm/transformers_utils/configs/extract_hidden_states.py

@@ -23,10 +23,14 @@ class ExtractHiddenStatesConfig(PretrainedConfig):
 
         if isinstance(model, dict):
             model_dict = model
+            source_text_config = None
         elif isinstance(model, PretrainedConfig):
             model_dict = model.to_dict()
+            text_config = model.get_text_config()
+            source_text_config = text_config if text_config is not model else None
         else:
             model_dict = {}
+            source_text_config = None
 
         # Combine: model_dict first, then kwargs override
         combined = {**model_dict, **kwargs}
@@ -35,6 +39,12 @@ class ExtractHiddenStatesConfig(PretrainedConfig):
 
         combined["architectures"] = ["ExtractHiddenStatesModel"]
 
+        # to_dict() and kwargs both flatten text_config to a plain dict;
+        # downstream get_hf_text_config() needs it as a PretrainedConfig
+        # for attribute access. Re-insert the original object.
+        if source_text_config is not None:
+            combined["text_config"] = source_text_config
+
         super().__init__(**combined)
 
     @classmethod

vllm/v1/worker/encoder_cudagraph.py

@@ -16,7 +16,7 @@ from vllm.distributed import (
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
 from vllm.model_executor.models.vision import get_load_balance_assignment
-from vllm.v1.worker.gpu.mm.encoder_cudagraph_defs import (
+from vllm.v1.worker.encoder_cudagraph_defs import (
     EncoderCudaGraphConfig,
 )
 

vllm/v1/worker/gpu_model_runner.py

@@ -211,7 +211,7 @@ from .utils import (
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
     from vllm.v1.spec_decode.ngram_proposer import NgramProposer
-    from vllm.v1.worker.gpu.mm.encoder_cudagraph import EncoderCudaGraphManager
+    from vllm.v1.worker.encoder_cudagraph import EncoderCudaGraphManager
 
 logger = init_logger(__name__)
 
@@ -5988,7 +5988,7 @@ class GPUModelRunner(
                 SupportsEncoderCudaGraph,
                 supports_encoder_cudagraph,
             )
-            from vllm.v1.worker.gpu.mm.encoder_cudagraph import (
+            from vllm.v1.worker.encoder_cudagraph import (
                 EncoderCudaGraphManager,
             )