[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/hardware_tests/amd.yaml

@@ -5,7 +5,6 @@ steps:
     depends_on: []
     device: amd_cpu
     no_plugin: true
-    soft_fail: true
     commands:
     - >
       docker build

.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/scripts/hardware_ci/run-amd-test.sh

@@ -35,23 +35,6 @@ export PYTHONPATH=".."
 # Helper Functions
 ###############################################################################
 
-wait_for_clean_gpus() {
-  local timeout=${1:-300}
-  local start=$SECONDS
-  echo "--- Waiting for clean GPU state (timeout: ${timeout}s)"
-  while true; do
-    if grep -q clean /opt/amdgpu/etc/gpu_state; then
-      echo "GPUs state is \"clean\""
-      return
-    fi
-    if (( SECONDS - start >= timeout )); then
-      echo "Error: GPUs did not reach clean state within ${timeout}s" >&2
-      exit 1
-    fi
-    sleep 3
-  done
-}
-
 cleanup_docker() {
   # Get Docker's root directory
   docker_root=$(docker info -f '{{.DockerRootDir}}')
@@ -365,19 +348,12 @@ apply_rocm_test_overrides() {
 ###############################################################################
 
 # --- GPU initialization ---
-echo "--- Confirming Clean Initial State"
-wait_for_clean_gpus
-
 echo "--- ROCm info"
 rocminfo
 
 # --- Docker housekeeping ---
 cleanup_docker
 
-echo "--- Resetting GPUs"
-echo "reset" > /opt/amdgpu/etc/gpu_state
-wait_for_clean_gpus
-
 # --- Pull test image ---
 echo "--- Pulling container"
 image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}"

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

@@ -224,20 +224,6 @@ steps:
   commands:
     - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
 
-- label: MessageQueue TCP Multi-Node (2 GPUs)
-  timeout_in_minutes: 10
-  working_dir: "/vllm-workspace/tests"
-  num_devices: 1
-  num_nodes: 2
-  no_plugin: true
-  optional: true
-  source_file_dependencies:
-  - vllm/distributed/device_communicators/shm_broadcast.py
-  - vllm/distributed/parallel_state.py
-  - tests/distributed/test_mq_tcp_multinode.py
-  commands:
-    - ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 1 $IMAGE_TAG "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py" "torchrun --nnodes 2 --nproc-per-node=1 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_mq_tcp_multinode.py"
-
 - label: Distributed NixlConnector PD accuracy (4 GPUs)
   timeout_in_minutes: 30
   working_dir: "/vllm-workspace/tests"

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

@@ -78,7 +78,6 @@ steps:
     - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
     - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
 
-# These require fix https://github.com/vllm-project/vllm/pull/36280
 - label: Model Runner V2 Pipeline Parallelism (4 GPUs)
   timeout_in_minutes: 60
   working_dir: "/vllm-workspace/tests"

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

@@ -91,9 +91,9 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
 
   if (n == 0) return;
 
-  const int64_t* src_data = src_ptrs.data_ptr<int64_t>();
-  const int64_t* dst_data = dst_ptrs.data_ptr<int64_t>();
-  const int64_t* size_data = sizes.data_ptr<int64_t>();
+  int64_t* src_data = src_ptrs.mutable_data_ptr<int64_t>();
+  int64_t* dst_data = dst_ptrs.mutable_data_ptr<int64_t>();
+  int64_t* size_data = sizes.mutable_data_ptr<int64_t>();
 
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
@@ -107,15 +107,24 @@ void swap_blocks_batch(const torch::Tensor& src_ptrs,
   CUmemcpyAttributes attr = {};
   attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
   size_t attrs_idx = 0;
+  #if defined(CUDA_VERSION) && CUDA_VERSION >= 13000
+  CUresult result = cuMemcpyBatchAsync(
+      reinterpret_cast<CUdeviceptr*>(dst_data),
+      reinterpret_cast<CUdeviceptr*>(src_data),
+      reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
+      &attrs_idx, 1, static_cast<CUstream>(stream));
+  TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed with error ",
+              result);
+  #else
   size_t fail_idx = 0;
   CUresult result = cuMemcpyBatchAsync(
-      reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(dst_data)),
-      reinterpret_cast<CUdeviceptr*>(const_cast<int64_t*>(src_data)),
-      reinterpret_cast<size_t*>(const_cast<int64_t*>(size_data)),
-      static_cast<size_t>(n), &attr, &attrs_idx, 1, &fail_idx,
-      static_cast<CUstream>(stream));
+      reinterpret_cast<CUdeviceptr*>(dst_data),
+      reinterpret_cast<CUdeviceptr*>(src_data),
+      reinterpret_cast<size_t*>(size_data), static_cast<size_t>(n), &attr,
+      &attrs_idx, 1, &fail_idx, static_cast<CUstream>(stream));
   TORCH_CHECK(result == CUDA_SUCCESS, "cuMemcpyBatchAsync failed at index ",
               fail_idx, " with error ", result);
+  #endif
 #else
   // Fallback for CUDA < 12.8 and ROCm: individual async copies.
   // cudaMemcpyDefault lets the driver infer direction from pointer types.

docker/Dockerfile.rocm

@@ -390,20 +390,21 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
 RUN mkdir src && mv vllm src/vllm
 
 # This is a workaround to ensure pytest exits with the correct status code in CI tests.
-RUN cat << 'EOF' > /vllm-workspace/conftest.py
-import os
-
-_exit_code = 1
-
-def pytest_sessionfinish(session, exitstatus):
-    global _exit_code
-    _exit_code = int(exitstatus)
-
-def pytest_unconfigure(config):
-    sys.stdout.flush()
-    sys.stderr.flush()
-    os._exit(_exit_code)
-EOF
+RUN printf '%s\n' \
+    'import os' \
+    '' \
+    '_exit_code = 1' \
+    '' \
+    'def pytest_sessionfinish(session, exitstatus):' \
+    '    global _exit_code' \
+    '    _exit_code = int(exitstatus)' \
+    '' \
+    'def pytest_unconfigure(config):' \
+    '    import sys' \
+    '    sys.stdout.flush()' \
+    '    sys.stderr.flush()' \
+    '    os._exit(_exit_code)' \
+    > /vllm-workspace/conftest.py
 
 # -----------------------
 # Final vLLM image

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

requirements/rocm-test.in

@@ -1,3 +1,5 @@
+-r common.txt
+
 # testing
 pytest
 tensorizer==2.10.1

requirements/rocm-test.txt

@@ -15,6 +15,7 @@ aiohappyeyeballs==2.6.1
 aiohttp==3.13.3
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   aiohttp-cors
     #   fsspec
     #   gpt-oss
@@ -38,20 +39,31 @@ annotated-doc==0.0.4
     #   typer
 annotated-types==0.7.0
     # via pydantic
+anthropic==0.89.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 antlr4-python3-runtime==4.9.3
     # via
     #   hydra-core
     #   omegaconf
-anyio==4.6.2.post1
+anyio==4.13.0
     # via
+    #   anthropic
     #   httpx
+    #   mcp
+    #   openai
+    #   sse-starlette
     #   starlette
+    #   watchfiles
 arctic-inference==0.1.1
     # via -r requirements/rocm-test.in
 argcomplete==3.6.3
     # via datamodel-code-generator
 arrow==1.4.0
     # via isoduration
+astor==0.8.1
+    # via depyf
 attrs==26.1.0
     # via
     #   aiohttp
@@ -83,6 +95,8 @@ bitsandbytes==0.49.2
     #   lightning
 black==26.3.1
     # via datamodel-code-generator
+blake3==1.0.8
+    # via -r requirements/common.txt
 blobfile==3.0.0
     # via -r requirements/rocm-test.in
 bm25s==0.2.13
@@ -99,6 +113,10 @@ bounded-pool-executor==0.0.3
     # via pqdm
 buildkite-test-collector==0.1.9
     # via -r requirements/rocm-test.in
+cachetools==7.0.5
+    # via -r requirements/common.txt
+cbor2==5.9.0
+    # via -r requirements/common.txt
 certifi==2026.2.25
     # via
     #   fiona
@@ -132,6 +150,7 @@ click==8.3.1
     #   nltk
     #   rasterio
     #   ray
+    #   rich-toolkit
     #   schemathesis
     #   typer
     #   uvicorn
@@ -142,6 +161,8 @@ cligj==0.7.2
     # via
     #   fiona
     #   rasterio
+cloudpickle==3.1.2
+    # via -r requirements/common.txt
 colorama==0.4.6
     # via
     #   perceptron
@@ -151,6 +172,10 @@ colorful==0.5.8
     # via ray
 colorlog==6.10.1
     # via optuna
+compressed-tensors==0.14.0.1
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 contourpy==1.3.3
     # via matplotlib
 coverage==7.13.5
@@ -182,24 +207,42 @@ decorator==5.2.1
     # via librosa
 decord==0.6.0
     # via -r requirements/rocm-test.in
+depyf==0.20.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 diffusers==0.37.0
     # via terratorch
 dill==0.3.8
     # via
     #   datasets
+    #   depyf
     #   evaluate
     #   lm-eval
     #   multiprocess
+diskcache==5.6.3
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 distlib==0.4.0
     # via virtualenv
+distro==1.9.0
+    # via
+    #   anthropic
+    #   openai
+dnspython==2.8.0
+    # via email-validator
 docker==7.1.0
     # via gpt-oss
 docopt==0.6.2
     # via num2words
 docstring-parser==0.17.0
-    # via jsonargparse
+    # via
+    #   anthropic
+    #   jsonargparse
 einops==0.8.2
     # via
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   encodec
     #   terratorch
@@ -208,6 +251,10 @@ einops==0.8.2
     #   vocos
 einx==0.4.2
     # via vector-quantize-pytorch
+email-validator==2.3.0
+    # via
+    #   fastapi
+    #   pydantic
 encodec==0.1.1
     # via vocos
 et-xmlfile==2.0.0
@@ -217,7 +264,15 @@ evaluate==0.4.6
 fastapi==0.135.2
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   gpt-oss
+    #   model-hosting-container-standards
+fastapi-cli==0.0.24
+    # via fastapi
+fastapi-cloud-cli==0.15.1
+    # via fastapi-cli
+fastar==0.9.0
+    # via fastapi-cloud-cli
 fastparquet==2026.3.0
     # via genai-perf
 fastsafetensors==0.2.2
@@ -225,6 +280,7 @@ fastsafetensors==0.2.2
 filelock==3.25.2
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   blobfile
     #   datasets
     #   diffusers
@@ -264,6 +320,10 @@ genson==1.3.0
     # via datamodel-code-generator
 geopandas==1.1.3
     # via terratorch
+gguf==0.18.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 gitdb==4.0.12
     # via gitpython
 gitpython==3.1.46
@@ -290,7 +350,10 @@ google-crc32c==1.8.0
 google-resumable-media==2.8.0
     # via google-cloud-storage
 googleapis-common-protos==1.73.0
-    # via google-api-core
+    # via
+    #   google-api-core
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
 gpt-oss==0.0.8
     # via -r requirements/rocm-test.in
 graphql-core==3.2.8
@@ -302,6 +365,7 @@ grpcio==1.78.0
     #   -c requirements/rocm.txt
     #   -r requirements/rocm-test.in
     #   grpcio-reflection
+    #   opentelemetry-exporter-otlp-proto-grpc
     #   ray
     #   tensorboard
 grpcio-reflection==1.78.0
@@ -328,12 +392,22 @@ html2text==2025.4.15
     # via gpt-oss
 httpcore==1.0.9
     # via httpx
+httptools==0.7.1
+    # via uvicorn
 httpx==0.27.2
     # via
     #   -r requirements/rocm-test.in
+    #   anthropic
     #   diffusers
+    #   fastapi
+    #   fastapi-cloud-cli
+    #   mcp
+    #   model-hosting-container-standards
+    #   openai
     #   perceptron
     #   schemathesis
+httpx-sse==0.4.3
+    # via mcp
 huggingface-hub==0.36.2
     # via
     #   -r requirements/rocm-test.in
@@ -370,10 +444,13 @@ hypothesis-jsonschema==0.23.1
 idna==3.11
     # via
     #   anyio
+    #   email-validator
     #   httpx
     #   jsonschema
     #   requests
     #   yarl
+ijson==3.5.0
+    # via -r requirements/common.txt
 imagehash==4.3.2
     # via -r requirements/rocm-test.in
 imageio==2.37.3
@@ -390,6 +467,8 @@ iniconfig==2.3.0
     # via pytest
 instanttensor==0.1.6
     # via -r requirements/rocm-test.in
+interegular==0.3.3
+    # via lm-format-enforcer
 isodate==0.7.2
     # via azure-storage-blob
 isoduration==20.11.0
@@ -399,15 +478,21 @@ isort==8.0.1
 jinja2==3.1.6
     # via
     #   datamodel-code-generator
+    #   fastapi
     #   genai-perf
     #   lm-eval
     #   torch
+jiter==0.13.0
+    # via
+    #   anthropic
+    #   openai
 jiwer==4.0.0
     # via -r requirements/rocm-test.in
 jmespath==1.1.0
     # via
     #   boto3
     #   botocore
+    #   model-hosting-container-standards
 joblib==1.5.3
     # via
     #   librosa
@@ -426,6 +511,7 @@ jsonpointer==3.1.0
 jsonschema==4.26.0
     # via
     #   hypothesis-jsonschema
+    #   mcp
     #   mistral-common
     #   ray
     #   schemathesis
@@ -443,6 +529,10 @@ kornia==0.8.2
     # via torchgeo
 kornia-rs==0.1.10
     # via kornia
+lark==1.2.2
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 lazy-loader==0.4
     # via
     #   librosa
@@ -466,14 +556,24 @@ lightning-utilities==0.15.3
     #   lightning
     #   pytorch-lightning
     #   torchmetrics
+llguidance==1.3.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 llvmlite==0.44.0
     # via numba
 lm-eval==0.4.11
     # via -r requirements/rocm-test.in
+lm-format-enforcer==0.11.3
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 logistro==2.0.1
     # via
     #   choreographer
     #   kaleido
+loguru==0.7.3
+    # via compressed-tensors
 lxml==6.0.2
     # via
     #   blobfile
@@ -500,12 +600,19 @@ mbstrdecoder==1.1.4
     #   dataproperty
     #   pytablewriter
     #   typepy
+mcp==1.27.0
+    # via -r requirements/common.txt
 mdurl==0.1.2
     # via markdown-it-py
 mistral-common==1.10.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
+model-hosting-container-standards==0.1.14
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 more-itertools==10.8.0
     # via
     #   inflect
@@ -522,6 +629,8 @@ msgpack==1.1.2
     # via
     #   librosa
     #   ray
+msgspec==0.20.0
+    # via -r requirements/common.txt
 mteb==2.11.5
     # via -r requirements/rocm-test.in
 multidict==6.7.1
@@ -541,6 +650,8 @@ networkx==3.6.1
     # via
     #   scikit-image
     #   torch
+ninja==1.13.0
+    # via -r requirements/common.txt
 nltk==3.9.3
     # via rouge-score
 num2words==0.5.14
@@ -555,6 +666,7 @@ numkong==7.1.1
     # via albucore
 numpy==2.2.6
     # via
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   accelerate
     #   albucore
@@ -572,6 +684,7 @@ numpy==2.2.6
     #   fastparquet
     #   genai-perf
     #   geopandas
+    #   gguf
     #   h5py
     #   imagehash
     #   imageio
@@ -620,15 +733,21 @@ numpy==2.2.6
     #   tritonclient
     #   vocos
     #   xarray
+    #   xgrammar
 omegaconf==2.3.0
     # via
     #   hydra-core
     #   lightning
 open-clip-torch==2.32.0
     # via -r requirements/rocm-test.in
+openai==2.30.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 openai-harmony==0.0.8
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   gpt-oss
 opencensus==0.11.4
     # via ray
@@ -637,6 +756,7 @@ opencensus-context==0.1.3
 opencv-python-headless==4.13.0.92
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   albumentations
     #   mistral-common
@@ -645,26 +765,59 @@ openpyxl==3.1.5
 opentelemetry-api==1.40.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-exporter-prometheus
     #   opentelemetry-sdk
     #   opentelemetry-semantic-conventions
+opentelemetry-exporter-otlp==1.40.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
+opentelemetry-exporter-otlp-proto-common==1.40.0
+    # via
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
+opentelemetry-exporter-otlp-proto-grpc==1.40.0
+    # via opentelemetry-exporter-otlp
+opentelemetry-exporter-otlp-proto-http==1.40.0
+    # via opentelemetry-exporter-otlp
 opentelemetry-exporter-prometheus==0.61b0
     # via ray
 opentelemetry-proto==1.40.0
-    # via ray
+    # via
+    #   opentelemetry-exporter-otlp-proto-common
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
+    #   ray
 opentelemetry-sdk==1.40.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-exporter-prometheus
+    #   opentelemetry-semantic-conventions-ai
     #   ray
 opentelemetry-semantic-conventions==0.61b0
-    # via opentelemetry-sdk
+    # via
+    #   opentelemetry-sdk
+    #   opentelemetry-semantic-conventions-ai
+opentelemetry-semantic-conventions-ai==0.5.1
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 optuna==3.6.1
     # via genai-perf
 orjson==3.11.7
     # via
     #   genai-perf
     #   kaleido
+outlines-core==0.2.11
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 packaging==26.0
     # via
     #   -c requirements/rocm.txt
@@ -682,6 +835,7 @@ packaging==26.0
     #   lazy-loader
     #   lightning
     #   lightning-utilities
+    #   lm-format-enforcer
     #   matplotlib
     #   optuna
     #   peft
@@ -713,6 +867,8 @@ pandas==3.0.1
     #   tacoreader
     #   torchgeo
     #   xarray
+partial-json-parser==0.2.1.1.post7
+    # via -r requirements/common.txt
 pathspec==1.0.4
     # via black
 pathvalidate==3.3.1
@@ -727,6 +883,7 @@ perf-analyzer==0.1.0
     # via genai-perf
 pillow==12.1.1
     # via
+    #   -r requirements/common.txt
     #   diffusers
     #   genai-perf
     #   imagehash
@@ -768,8 +925,14 @@ pqdm==0.2.0
 prometheus-client==0.24.1
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   opentelemetry-exporter-prometheus
+    #   prometheus-fastapi-instrumentator
     #   ray
+prometheus-fastapi-instrumentator==7.1.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 propcache==0.4.1
     # via
     #   aiohttp
@@ -779,6 +942,7 @@ proto-plus==1.27.1
 protobuf==6.33.6
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   google-api-core
     #   googleapis-common-protos
     #   grpcio-reflection
@@ -791,11 +955,14 @@ protobuf==6.33.6
     #   wandb
 psutil==7.2.2
     # via
+    #   -r requirements/common.txt
     #   accelerate
     #   peft
     #   tensorizer
 py==1.11.0
     # via pytest-forked
+py-cpuinfo==9.0.0
+    # via -r requirements/common.txt
 py-spy==0.4.1
     # via ray
 pyarrow==23.0.1
@@ -808,6 +975,8 @@ pyasn1==0.6.3
     # via pyasn1-modules
 pyasn1-modules==0.4.2
     # via google-auth
+pybase64==1.4.3
+    # via -r requirements/common.txt
 pycocotools==2.0.11
     # via terratorch
 pycountry==26.2.16
@@ -819,26 +988,44 @@ pycryptodomex==3.23.0
 pydantic==2.12.5
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   albumentations
+    #   anthropic
+    #   compressed-tensors
     #   datamodel-code-generator
     #   fastapi
+    #   fastapi-cloud-cli
     #   gpt-oss
     #   lightly
+    #   lm-format-enforcer
+    #   mcp
     #   mistral-common
+    #   model-hosting-container-standards
     #   mteb
+    #   openai
     #   openai-harmony
     #   pydantic-extra-types
+    #   pydantic-settings
     #   ray
     #   wandb
+    #   xgrammar
 pydantic-core==2.41.5
     # via pydantic
 pydantic-extra-types==2.11.1
-    # via mistral-common
+    # via
+    #   fastapi
+    #   mistral-common
+pydantic-settings==2.13.1
+    # via
+    #   fastapi
+    #   mcp
 pygments==2.19.2
     # via rich
 pyjwt==2.12.1
-    # via msal
+    # via
+    #   mcp
+    #   msal
 pyogrio==0.12.1
     # via geopandas
 pyparsing==3.3.2
@@ -898,6 +1085,16 @@ python-dateutil==2.9.0.post0
     #   typepy
 python-discovery==1.2.0
     # via virtualenv
+python-dotenv==1.2.2
+    # via
+    #   pydantic-settings
+    #   uvicorn
+python-json-logger==4.1.0
+    # via -r requirements/common.txt
+python-multipart==0.0.22
+    # via
+    #   fastapi
+    #   mcp
 python-rapidjson==1.23
     # via tritonclient
 pytokens==0.4.1
@@ -914,14 +1111,17 @@ pywavelets==1.9.0
     # via imagehash
 pyyaml==6.0.3
     # via
+    #   -r requirements/common.txt
     #   accelerate
     #   albumentations
     #   datamodel-code-generator
     #   datasets
     #   genai-perf
+    #   gguf
     #   huggingface-hub
     #   jsonargparse
     #   lightning
+    #   lm-format-enforcer
     #   omegaconf
     #   optuna
     #   peft
@@ -931,8 +1131,13 @@ pyyaml==6.0.3
     #   schemathesis
     #   timm
     #   transformers
+    #   uvicorn
     #   vocos
     #   wandb
+pyzmq==27.1.0
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 rapidfuzz==3.12.1
     # via
     #   -r requirements/rocm-test.in
@@ -952,6 +1157,7 @@ referencing==0.37.0
     #   jsonschema-specifications
 regex==2026.2.28
     # via
+    #   -r requirements/common.txt
     #   diffusers
     #   nltk
     #   open-clip-torch
@@ -961,12 +1167,14 @@ regex==2026.2.28
 requests==2.32.5
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   azure-core
     #   buildkite-test-collector
     #   datasets
     #   diffusers
     #   docker
     #   evaluate
+    #   gguf
     #   google-api-core
     #   google-cloud-storage
     #   gpt-oss
@@ -976,6 +1184,7 @@ requests==2.32.5
     #   mistral-common
     #   msal
     #   mteb
+    #   opentelemetry-exporter-otlp-proto-http
     #   pooch
     #   ray
     #   responses
@@ -999,8 +1208,15 @@ rich==14.3.3
     #   lightning
     #   mteb
     #   perceptron
+    #   rich-toolkit
     #   terratorch
     #   typer
+rich-toolkit==0.19.7
+    # via
+    #   fastapi-cli
+    #   fastapi-cloud-cli
+rignore==0.7.6
+    # via fastapi-cloud-cli
 rioxarray==0.22.0
     # via terratorch
 rouge-score==0.1.2
@@ -1070,12 +1286,20 @@ sentence-transformers==5.3.0
     # via
     #   -r requirements/rocm-test.in
     #   mteb
+sentencepiece==0.2.1
+    # via -r requirements/common.txt
 sentry-sdk==2.55.0
-    # via wandb
+    # via
+    #   fastapi-cloud-cli
+    #   wandb
+setproctitle==1.3.7
+    # via -r requirements/common.txt
 setuptools==79.0.1
     # via
     #   -c requirements/common.txt
     #   -c requirements/rocm.txt
+    #   -r requirements/common.txt
+    #   model-hosting-container-standards
     #   pytablewriter
     #   tensorboard
     #   torch
@@ -1092,6 +1316,7 @@ simplejson==3.20.2
 six==1.17.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   junit-xml
     #   lightly
     #   opencensus
@@ -1104,8 +1329,9 @@ smmap==5.0.3
     # via gitdb
 sniffio==1.3.1
     # via
-    #   anyio
+    #   anthropic
     #   httpx
+    #   openai
 sortedcontainers==2.4.0
     # via hypothesis
 soundfile==0.13.1
@@ -1124,10 +1350,16 @@ sqlalchemy==2.0.48
     #   optuna
 sqlitedict==2.1.0
     # via lm-eval
+sse-starlette==3.3.4
+    # via mcp
 starlette==0.52.1
     # via
     #   fastapi
+    #   mcp
+    #   model-hosting-container-standards
+    #   prometheus-fastapi-instrumentator
     #   schemathesis
+    #   sse-starlette
     #   starlette-testclient
 starlette-testclient==0.4.1
     # via schemathesis
@@ -1137,6 +1369,8 @@ stringzilla==4.6.0
     # via albucore
 structlog==25.5.0
     # via gpt-oss
+supervisor==4.3.0
+    # via model-hosting-container-standards
 sympy==1.14.0
     # via
     #   einx
@@ -1180,6 +1414,7 @@ tifffile==2026.3.3
 tiktoken==0.12.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   gpt-oss
     #   lm-eval
     #   mistral-common
@@ -1194,6 +1429,7 @@ timm==1.0.17
 tokenizers==0.22.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
     #   transformers
 tomli==2.4.0
@@ -1212,8 +1448,10 @@ torchmetrics==1.9.0
     #   torchgeo
 tqdm==4.67.3
     # via
+    #   -r requirements/common.txt
     #   datasets
     #   evaluate
+    #   gguf
     #   huggingface-hub
     #   lightly
     #   lightning
@@ -1221,6 +1459,7 @@ tqdm==4.67.3
     #   mteb
     #   nltk
     #   open-clip-torch
+    #   openai
     #   optuna
     #   peft
     #   pqdm
@@ -1233,11 +1472,14 @@ tqdm==4.67.3
 transformers==4.57.5
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   -r requirements/rocm-test.in
+    #   compressed-tensors
     #   genai-perf
     #   peft
     #   sentence-transformers
     #   transformers-stream-generator
+    #   xgrammar
 transformers-stream-generator==0.0.5
     # via -r requirements/rocm-test.in
 tritonclient==2.66.0
@@ -1251,6 +1493,8 @@ typepy==1.3.4
     #   tabledata
 typer==0.24.1
     # via
+    #   fastapi-cli
+    #   fastapi-cloud-cli
     #   fastsafetensors
     #   perceptron
 typeshed-client==2.9.0
@@ -1258,9 +1502,12 @@ typeshed-client==2.9.0
 typing-extensions==4.15.0
     # via
     #   -c requirements/common.txt
+    #   -r requirements/common.txt
     #   aiosignal
     #   albumentations
     #   alembic
+    #   anthropic
+    #   anyio
     #   azure-core
     #   azure-identity
     #   azure-storage-blob
@@ -1272,9 +1519,13 @@ typing-extensions==4.15.0
     #   lightning
     #   lightning-utilities
     #   lm-eval
+    #   mcp
     #   mistral-common
     #   mteb
+    #   openai
     #   opentelemetry-api
+    #   opentelemetry-exporter-otlp-proto-grpc
+    #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-sdk
     #   opentelemetry-semantic-conventions
     #   pqdm
@@ -1283,6 +1534,7 @@ typing-extensions==4.15.0
     #   pydantic-extra-types
     #   pytorch-lightning
     #   referencing
+    #   rich-toolkit
     #   sentence-transformers
     #   sqlalchemy
     #   starlette
@@ -1292,10 +1544,13 @@ typing-extensions==4.15.0
     #   typeshed-client
     #   typing-inspection
     #   wandb
+    #   xgrammar
 typing-inspection==0.4.2
     # via
     #   fastapi
+    #   mcp
     #   pydantic
+    #   pydantic-settings
 tzdata==2025.3
     # via arrow
 uri-template==1.3.0
@@ -1311,7 +1566,14 @@ urllib3==2.6.3
     #   sentry-sdk
     #   tritonclient
 uvicorn==0.42.0
-    # via gpt-oss
+    # via
+    #   fastapi
+    #   fastapi-cli
+    #   fastapi-cloud-cli
+    #   gpt-oss
+    #   mcp
+uvloop==0.22.1
+    # via uvicorn
 vector-quantize-pytorch==1.28.0
     # via -r requirements/rocm-test.in
 virtualenv==21.2.0
@@ -1320,10 +1582,16 @@ vocos==0.1.0
     # via -r requirements/rocm-test.in
 wandb==0.25.1
     # via terratorch
+watchfiles==1.1.1
+    # via
+    #   -r requirements/common.txt
+    #   uvicorn
 wcwidth==0.6.0
     # via ftfy
 webcolors==25.10.0
     # via jsonschema
+websockets==16.0
+    # via uvicorn
 werkzeug==3.1.6
     # via
     #   schemathesis
@@ -1334,6 +1602,10 @@ wrapt==2.1.2
     # via smart-open
 xarray==2026.2.0
     # via rioxarray
+xgrammar==0.1.33
+    # via
+    #   -c requirements/common.txt
+    #   -r requirements/common.txt
 xxhash==3.6.0
     # via
     #   datasets

setup.py

@@ -1060,8 +1060,6 @@ setup(
         ],  # Required for audio processing
         "video": [],  # Kept for backwards compatibility
         "flashinfer": [],  # Kept for backwards compatibility
-        # Optional deps for AMD FP4 quantization support
-        "petit-kernel": ["petit-kernel"],
         # Optional deps for Helion kernel development
         # NOTE: When updating helion version, also update CI files:
         #   - .buildkite/test_areas/kernels.yaml

tests/distributed/test_mq_tcp_multinode.py

@@ -1,119 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Multi-node integration test for MessageQueue TCP fallback.
-
-Verifies that when writer and readers span separate nodes (Docker containers
-with isolated /dev/shm), `create_from_process_group` correctly detects
-cross-node ranks via `in_the_same_node_as()` and falls back to ZMQ TCP
-transport — and that data actually arrives.
-"""
-
-import numpy as np
-import torch.distributed as dist
-
-from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
-from vllm.distributed.parallel_state import in_the_same_node_as
-
-
-def main():
-    dist.init_process_group(backend="gloo")
-
-    rank = dist.get_rank()
-    world_size = dist.get_world_size()
-    assert world_size >= 2, (
-        f"Need at least 2 ranks across nodes, got world_size={world_size}"
-    )
-
-    # Verify that in_the_same_node_as detects cross-node correctly
-    status = in_the_same_node_as(dist.group.WORLD, source_rank=0)
-    local_count = sum(status)
-    print(
-        f"[Rank {rank}] in_the_same_node_as(source=0): {status}  "
-        f"(local={local_count}/{world_size})"
-    )
-    # With 2 Docker containers (1 proc each), rank 0 and rank 1
-    # should be on different nodes.
-    assert local_count < world_size, (
-        f"Expected cross-node ranks but all {world_size} ranks appear local."
-    )
-
-    # Create MessageQueue
-    writer_rank = 0
-    mq = MessageQueue.create_from_process_group(
-        dist.group.WORLD,
-        max_chunk_bytes=1024 * 1024,  # 1 MiB
-        max_chunks=10,
-        writer_rank=writer_rank,
-    )
-
-    # Verify the transport path selection
-    if rank == writer_rank:
-        print(
-            f"[Rank {rank}] Writer: n_local_reader={mq.n_local_reader}, "
-            f"n_remote_reader={mq.n_remote_reader}"
-        )
-        assert mq.n_remote_reader > 0, (
-            "Writer should have at least 1 remote (TCP) reader in a multi-node setup."
-        )
-    else:
-        if status[rank]:
-            assert mq._is_local_reader, (
-                f"Rank {rank} is on the same node as writer but is not a local reader."
-            )
-            print(f"[Rank {rank}] Reader: local (shared memory)")
-        else:
-            assert mq._is_remote_reader, (
-                f"Rank {rank} is on a different node but is not a remote (TCP) reader."
-            )
-            print(f"[Rank {rank}] Reader: remote (TCP)")
-
-    # Test data transfer: simple objects
-    dist.barrier()
-    if rank == writer_rank:
-        mq.enqueue("hello_from_node0")
-    else:
-        msg = mq.dequeue(timeout=10)
-        assert msg == "hello_from_node0"
-    dist.barrier()
-    print(f"[Rank {rank}] Simple object test passed")
-
-    # Test data transfer: numpy arrays
-    np.random.seed(42)
-    arrays = [
-        np.random.randint(0, 100, size=np.random.randint(100, 5000)) for _ in range(100)
-    ]
-
-    dist.barrier()
-    if rank == writer_rank:
-        for arr in arrays:
-            mq.enqueue(arr)
-    else:
-        for i, expected in enumerate(arrays):
-            received = mq.dequeue(timeout=10)
-            assert np.array_equal(expected, received), (
-                f"Array mismatch at index {i}: "
-                f"expected shape {expected.shape}, got shape {received.shape}"
-            )
-    dist.barrier()
-    print(f"[Rank {rank}] Numpy array test passed")
-
-    # Test data transfer: large payload (> max_chunk_bytes)
-    dist.barrier()
-    big_array = np.zeros(200_000, dtype=np.int64)  # ~1.6 MiB > 1 MiB chunk
-    if rank == writer_rank:
-        mq.enqueue(big_array)
-    else:
-        received = mq.dequeue(timeout=10)
-        assert np.array_equal(big_array, received)
-    dist.barrier()
-    print(f"[Rank {rank}] Large payload test passed")
-
-    # Done -- cleanup
-    dist.barrier()
-    print(f"[Rank {rank}] All MessageQueue TCP multi-node tests passed!")
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()

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/quantization/test_online.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests online quantization."""
+
+import pytest
+import torch
+
+from tests.quantization.utils import (
+    _test_online_quant_peak_mem_impl,
+    is_quant_method_supported,
+)
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PerBlockOnlineLinearMethod,
+    Fp8PerBlockOnlineMoEMethod,
+    Fp8PerTensorOnlineLinearMethod,
+    Fp8PerTensorOnlineMoEMethod,
+)
+from vllm.platforms import current_platform
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+@pytest.mark.parametrize(
+    "quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls",
+    [
+        # simple case - quantization='fp8_per_tensor'
+        (
+            "fp8_per_tensor",
+            None,
+            Fp8PerTensorOnlineLinearMethod,
+            Fp8PerTensorOnlineMoEMethod,
+        ),
+        # simple case - quantization='fp8_per_block'
+        (
+            "fp8_per_block",
+            None,
+            Fp8PerBlockOnlineLinearMethod,
+            Fp8PerBlockOnlineMoEMethod,
+        ),
+        # quantization='online with linear_scheme_override and
+        # moe_scheme_override
+        (
+            "online",
+            {
+                "linear_scheme_override": "fp8_per_block",
+                "moe_scheme_override": "fp8_per_tensor",
+            },
+            Fp8PerBlockOnlineLinearMethod,
+            Fp8PerTensorOnlineMoEMethod,
+        ),
+        # ignore with direct layer name
+        (
+            "fp8_per_tensor",
+            # qkv_proj is fused from q_proj/k_proj/v_proj, so currently the
+            # ignore regex must match the unfused shard names
+            # TODO(future PR): also make 're:.*qkv_proj.*' work
+            {"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]},
+            Fp8PerTensorOnlineLinearMethod,
+            Fp8PerTensorOnlineMoEMethod,
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
+)
+def test_online_quantization(
+    vllm_runner,
+    quant_scheme: str,
+    online_quant_args: dict | None,
+    expected_linear_cls,
+    expected_moe_cls,
+    use_rocm_aiter: bool,
+    monkeypatch,
+) -> None:
+    """
+    Tests that online quantization frontend configuration works -
+    selecting quant schemes, overriding quant schemes by type, ignoring
+    layers.
+
+    Does not test performance, peak memory usage, etc.
+    """
+
+    if use_rocm_aiter:
+        monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
+
+    # `LLM.apply_model` requires pickling a function.
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+    # a tiny model with both dense and MoE layers
+    model_name = "ibm-granite/granite-3.0-1b-a400m-base"
+
+    runner_kwargs = dict(
+        quantization=quant_scheme,
+        enforce_eager=True,
+    )
+    if online_quant_args is not None:
+        runner_kwargs["quantization_config"] = online_quant_args
+
+    with vllm_runner(
+        model_name,
+        **runner_kwargs,
+    ) as llm:
+
+        def check_model(model):
+            # checks further down in the test case are hardcoded for this
+            # model
+            assert model_name == "ibm-granite/granite-3.0-1b-a400m-base"
+
+            o_proj = model.model.layers[0].self_attn.o_proj
+            moe = model.model.layers[0].block_sparse_moe.experts
+
+            # o_proj and moe in layer 0 are always quantized (never ignored)
+            # because of how we craft the test case inputs
+            assert isinstance(o_proj.quant_method, expected_linear_cls)
+            if moe is not None:
+                assert isinstance(moe.quant_method, expected_moe_cls)
+
+            if current_platform.is_cuda():
+                assert o_proj.weight.dtype == torch.float8_e4m3fn
+            elif current_platform.is_rocm():
+                assert o_proj.weight.dtype == current_platform.fp8_dtype()
+            else:
+                pytest.skip("Only runs on CUDA and ROCm.")
+
+            # Verify ignored layers are unquantized.
+            if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
+                # only .*1.self_attn_o_proj is skipped
+                for layer_idx in range(len(model.model.layers)):
+                    o_proj = model.model.layers[layer_idx].self_attn.o_proj
+                    if layer_idx == 1:
+                        assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
+                    else:
+                        assert isinstance(o_proj.quant_method, expected_linear_cls)
+
+                # every .*self_attn.qkv_proj is skipped
+                for layer_idx in range(len(model.model.layers)):
+                    qkv_proj = model.model.layers[layer_idx].self_attn.qkv_proj
+                    assert isinstance(qkv_proj.quant_method, UnquantizedLinearMethod)
+
+        llm.apply_model(check_model)
+
+        outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
+        print(outputs[0][1])
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+def test_online_quant_peak_mem(
+    vllm_runner,
+    caplog_mp_spawn,
+    monkeypatch,
+) -> None:
+    _test_online_quant_peak_mem_impl(
+        "fp8_per_tensor", vllm_runner, caplog_mp_spawn, monkeypatch
+    )
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("fp8"),
+    reason="FP8 is not supported on this GPU type.",
+)
+def test_online_quant_load_format_dummy(
+    vllm_runner,
+    monkeypatch,
+    caplog,
+) -> None:
+    with vllm_runner(
+        "ibm-granite/granite-3.0-1b-a400m-base",
+        quantization="fp8_per_tensor",
+        enforce_eager=True,
+        load_format="dummy",
+    ) as llm:
+        outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
+        print(outputs[0][1])

tests/quantization/utils.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import logging
+
+import regex as re
+
 from vllm.model_executor.layers.quantization import get_quantization_config
 from vllm.platforms import current_platform
 
@@ -21,3 +25,74 @@ def is_quant_method_supported(quant_method: str) -> bool:
     min_capability = get_quantization_config(quant_method).get_min_capability()
 
     return capability.to_int() >= min_capability
+
+
+def _test_online_quant_peak_mem_impl(
+    quantization_arg_value,
+    vllm_runner,
+    caplog_mp_spawn,
+    monkeypatch,
+) -> None:
+    # Note: `allenai/OLMoE-1B-7B-0125-Instruct` was selected because:
+    # 1. it covers both Linear and MoE paths
+    # 2. it is already used by other tests in CI, so adding it here
+    #    does not increase disk space for CI runners
+    # I really wanted to use `ibm-granite/granite-3.0-1b-a400m-base`
+    # which I think is the smallest MoE model in vLLM (2.5 GiB bf16,
+    # 1.3 GiB fp8), but could not as adding one more model makes CI
+    # run out of disk space.
+    model_name = "allenai/OLMoE-1B-7B-0125-Instruct"
+
+    # Force spawn to ensure caplog_mp_spawn works consistently
+    # (it relies on VLLM_LOGGING_CONFIG_PATH which spawn reads but fork ignores)
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    with (
+        caplog_mp_spawn(logging.DEBUG) as log_holder,
+        vllm_runner(
+            model_name,
+            quantization=quantization_arg_value,
+            enforce_eager=True,
+        ) as llm,
+    ):
+        outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4)
+        print(outputs[0][1])
+
+    log_text = log_holder.text
+
+    # Parse memory usage from captured logs
+    model_memory_gib = None
+    peak_memory_gib = None
+    for line in log_text.splitlines():
+        if model_memory_gib is None:
+            match = re.search(r"Model loading took ([\d.]+) GiB memory", line)
+            if match:
+                model_memory_gib = float(match.group(1))
+        if peak_memory_gib is None:
+            match = re.search(
+                r"Peak GPU memory after loading weights: ([\d.]+) GiB", line
+            )
+            if match:
+                peak_memory_gib = float(match.group(1))
+
+    assert model_memory_gib is not None, "Could not find model loading memory log"
+    assert peak_memory_gib is not None, "Could not find peak memory log"
+    print(f"GPU memory used after loading weights: {model_memory_gib} GiB")
+    print(f"Peak GPU memory usage while loading weights: {peak_memory_gib} GiB")
+
+    # model specific, allenai/OLMoE-1B-7B-0125-Instruct fp8 online quant
+    # uses 6.65 GiB for weight loading (bf16 checkpoint is ~12.89 GiB)
+    expected_model_memory_gib = 6.7
+
+    # for allenai/OLMoE-1B-7B-0125-Instruct the number we see today is 9.06
+    # GiB, which is 1.36x above model_memory_gib. A slightly higher number is
+    # expected as when we load and quantize weights in a streaming fashion we
+    # need to have individual weights in bf16 + fp8 alive at the same time.
+    expected_peak_memory_gib = expected_model_memory_gib * 1.4
+
+    assert model_memory_gib < expected_model_memory_gib, (
+        f"{model_memory_gib=} higher than {expected_model_memory_gib}"
+    )
+    assert peak_memory_gib < expected_peak_memory_gib, (
+        f"{peak_memory_gib=} higher than {expected_peak_memory_gib}"
+    )

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

@@ -68,10 +68,13 @@ class IrOpPriorityConfig:
     def set_priority(self):
         """
         Context manager to set the IR op priority for all op members.
-        It also imports vllm.kernels to ensure all implementations are made available.
+        It also imports IR kernel implementations for the current platform
+        to ensure all implementations are made available.
         """
-        import vllm.kernels  # noqa: F401, registers IR op implementations
         from vllm.ir.op import IrOp
+        from vllm.platforms import current_platform
+
+        current_platform.import_ir_kernels()
 
         with contextlib.ExitStack() as stack:
             for field in fields(self):

vllm/config/model.py

@@ -21,6 +21,7 @@ from vllm.config.multimodal import (
     MultiModalConfig,
 )
 from vllm.config.pooler import PoolerConfig
+from vllm.config.quantization import OnlineQuantizationConfigArgs
 from vllm.config.scheduler import RunnerType
 from vllm.config.utils import config, getattr_iter
 from vllm.logger import init_logger
@@ -199,6 +200,10 @@ class ModelConfig:
     `quantization_config` attribute in the model config file. If that is
     `None`, we assume the model weights are not quantized and use `dtype` to
     determine the data type of the weights."""
+    quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None = None
+    """Arguments for online quantization.
+    Auto-created when `quantization` equals to one of the string values of
+    the `OnlineQuantScheme` enum."""
     allow_deprecated_quantization: bool = False
     """Whether to allow deprecated quantization methods."""
     enforce_eager: bool = False
@@ -943,7 +948,6 @@ class ModelConfig:
                 "modelopt_fp4",
                 "modelopt_mxfp8",
                 "modelopt_mixed",
-                "petit_nvfp4",
                 # Ensure heavy backends are probed last to avoid unnecessary
                 # imports during override detection (e.g., MXFP4 imports Triton)
                 "mxfp4",

vllm/config/quantization.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from enum import Enum
+from typing import Any
+
+from pydantic import Field, field_validator
+
+from vllm.config.utils import config
+
+
+class OnlineQuantScheme(Enum):
+    """Supported online quantization schemes."""
+
+    # fp8, weights and activations scaled per-tensor
+    FP8_PER_TENSOR = "fp8_per_tensor"
+
+    # fp8, activations scaled in blocks of 1x128 elements, weights scaled in
+    # blocks of 128x128 elements (popularized by DeepSeek)
+    FP8_PER_BLOCK = "fp8_per_block"
+
+    # TODO(future PRs): add more online quant schemes here: mxfp8, etc
+
+
+@config
+class OnlineQuantizationConfigArgs:
+    """Configuration for online quantization.
+
+    Controls how ``OnlineQuantizationConfig`` is applied to a model.
+    At least one of ``global_scheme``, ``linear_scheme_override``, or
+    ``moe_scheme_override`` must be set.
+    """
+
+    global_scheme: OnlineQuantScheme | None = None
+    """Quantization scheme applied to every supported layer."""
+
+    linear_scheme_override: OnlineQuantScheme | None = None
+    """Quantization scheme override for ``LinearBase`` layers."""
+
+    moe_scheme_override: OnlineQuantScheme | None = None
+    """Quantization scheme override for ``FusedMoE`` layers."""
+
+    ignore: list[str] = Field(default_factory=list)
+    """Layers to skip quantization for. Supports exact names and regex
+    patterns with ``re:`` prefix (e.g. ``re:.*attn.*``), consistent with
+    compressed_tensors layer skipping."""
+
+    @field_validator(
+        "global_scheme", "linear_scheme_override", "moe_scheme_override", mode="before"
+    )
+    @classmethod
+    def _coerce_scheme(
+        cls, v: str | OnlineQuantScheme | None
+    ) -> OnlineQuantScheme | None:
+        if isinstance(v, str):
+            return OnlineQuantScheme(v)
+        return v
+
+
+def resolve_online_quant_config(
+    quantization: str | None,
+    quantization_config: dict[str, Any] | OnlineQuantizationConfigArgs | None,
+) -> OnlineQuantizationConfigArgs | None:
+    """Resolve online quant scheme shorthand into a quantization config.
+
+    If ``quantization`` is an online quant scheme (e.g. ``'fp8_per_tensor'``),
+    ensures ``quantization_config`` has a matching ``global_scheme`` and casts
+    it to :class:`OnlineQuantizationConfigArgs` if needed.
+    """
+    online_quant_values = {s.value for s in OnlineQuantScheme}
+    valid_quantization_values = online_quant_values | {"online"}
+    if quantization not in valid_quantization_values:
+        if quantization_config is not None:
+            raise ValueError(
+                f"quantization_config is only supported when quantization "
+                f"is one of {sorted(valid_quantization_values)}, "
+                f"got quantization={quantization!r}"
+            )
+        return None
+
+    if quantization in online_quant_values:
+        scheme = OnlineQuantScheme(quantization)
+
+        if quantization_config is None:
+            quantization_config = {
+                "global_scheme": scheme.value,
+            }
+        elif isinstance(quantization_config, OnlineQuantizationConfigArgs):
+            if quantization_config.global_scheme is None:
+                quantization_config.global_scheme = scheme
+            elif quantization_config.global_scheme != scheme:
+                raise ValueError(
+                    f"quantization={quantization!r} conflicts with "
+                    f"quantization_config.global_scheme="
+                    f"{quantization_config.global_scheme.value!r}. "
+                    f"These must match when both are specified."
+                )
+        elif isinstance(quantization_config, dict):
+            existing = quantization_config.get("global_scheme")
+            if existing is None:
+                quantization_config["global_scheme"] = scheme.value
+            else:
+                # Coerce to enum for comparison
+                existing_scheme = (
+                    OnlineQuantScheme(existing)
+                    if isinstance(existing, str)
+                    else existing
+                )
+                if existing_scheme != scheme:
+                    raise ValueError(
+                        f"quantization={quantization!r} conflicts "
+                        f"with quantization_config"
+                        f"['global_scheme']={existing!r}. "
+                        f"These must match when both are specified."
+                    )
+
+    # Cast dict to OnlineQuantizationConfigArgs
+    if isinstance(quantization_config, dict):
+        quantization_config = OnlineQuantizationConfigArgs(**quantization_config)
+
+    return quantization_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/config/vllm.py

@@ -1106,6 +1106,9 @@ class VllmConfig:
             )
         current_platform.check_and_update_config(self)
 
+        if envs.VLLM_USE_V2_MODEL_RUNNER:
+            self._validate_v2_model_runner()
+
         # Re-compute compile ranges after platform-specific config updates
         # (e.g., XPU may lower max_num_batched_tokens when MLA is enabled)
         self._set_compile_ranges()
@@ -1713,6 +1716,7 @@ class VllmConfig:
             f"dcp_comm_backend={self.parallel_config.dcp_comm_backend}, "  # noqa
             f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, "  # noqa
             f"quantization={self.model_config.quantization}, "
+            f"quantization_config={self.model_config.quantization_config}, "  # noqa
             f"enforce_eager={self.model_config.enforce_eager}, "
             f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, "  # noqa
             f"kv_cache_dtype={self.cache_config.cache_dtype}, "
@@ -1728,6 +1732,49 @@ class VllmConfig:
             f"kernel_config={self.kernel_config!r}"
         )
 
+    def _validate_v2_model_runner(self) -> None:
+        """Check for features not yet supported by the V2 model runner."""
+        unsupported: list[str] = []
+
+        if self.model_config is not None and self.model_config.has_inner_state:
+            unsupported.append("hybrid/mamba models")
+
+        if self.parallel_config.prefill_context_parallel_size > 1:
+            unsupported.append("prefill context parallelism")
+
+        if (
+            self.speculative_config is not None
+            and self.speculative_config.method not in ("eagle", "eagle3", "mtp")
+        ):
+            unsupported.append(f"speculative method '{self.speculative_config.method}'")
+
+        if self.parallel_config.enable_dbo:
+            unsupported.append("dual batch overlap")
+
+        if (
+            self.model_config is not None
+            and self.model_config.enable_return_routed_experts
+        ):
+            # Will be added by https://github.com/vllm-project/vllm/pull/38163
+            unsupported.append("routed experts capture")
+
+        if self.model_config is not None and self.model_config.logits_processors:
+            unsupported.append("custom logits processors")
+
+        if self.cache_config.kv_sharing_fast_prefill:
+            # Will be added by https://github.com/vllm-project/vllm/pull/35045
+            unsupported.append("KV sharing fast prefill")
+
+        if self.ec_transfer_config is not None:
+            # Will be added by https://github.com/vllm-project/vllm/pull/38390
+            unsupported.append("EC transfer")
+
+        if unsupported:
+            raise ValueError(
+                "VLLM_USE_V2_MODEL_RUNNER does not yet support: "
+                + ", ".join(unsupported)
+            )
+
     def validate_block_size(self) -> None:
         """Validate block_size against DCP and mamba constraints.
 

vllm/engine/arg_utils.py

@@ -112,6 +112,7 @@ from vllm.v1.sample.logits_processor import LogitsProcessor
 from vllm.version import __version__ as VLLM_VERSION
 
 if TYPE_CHECKING:
+    from vllm.config.quantization import OnlineQuantizationConfigArgs
     from vllm.model_executor.layers.quantization import QuantizationMethods
     from vllm.model_executor.model_loader import LoadFormats
     from vllm.usage.usage_lib import UsageContext
@@ -483,6 +484,7 @@ class EngineArgs:
     hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")
     tokenizer_revision: str | None = ModelConfig.tokenizer_revision
     quantization: QuantizationMethods | str | None = ModelConfig.quantization
+    quantization_config: "dict[str, Any] | OnlineQuantizationConfigArgs | None" = None
     allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization
     enforce_eager: bool = ModelConfig.enforce_eager
     disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce
@@ -661,6 +663,12 @@ class EngineArgs:
         if isinstance(self.ir_op_priority, dict):
             self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority)
 
+        from vllm.config.quantization import resolve_online_quant_config
+
+        self.quantization_config = resolve_online_quant_config(
+            self.quantization, self.quantization_config
+        )
+
         # Setup plugins
         from vllm.plugins import load_general_plugins
 
@@ -1431,6 +1439,7 @@ class EngineArgs:
             tokenizer_revision=self.tokenizer_revision,
             max_model_len=self.max_model_len,
             quantization=self.quantization,
+            quantization_config=self.quantization_config,
             allow_deprecated_quantization=self.allow_deprecated_quantization,
             enforce_eager=self.enforce_eager,
             enable_return_routed_experts=self.enable_return_routed_experts,

vllm/entrypoints/llm.py

@@ -34,6 +34,9 @@ from vllm.config.model import (
     RunnerOption,
     TokenizerMode,
 )
+from vllm.config.quantization import (
+    OnlineQuantizationConfigArgs,
+)
 from vllm.distributed.weight_transfer.base import (
     WeightTransferInitRequest,
     WeightTransferUpdateRequest,
@@ -247,6 +250,9 @@ class LLM:
         attention_config: dict[str, Any] | AttentionConfig | None = None,
         kv_cache_memory_bytes: int | None = None,
         compilation_config: int | dict[str, Any] | CompilationConfig | None = None,
+        quantization_config: dict[str, Any]
+        | OnlineQuantizationConfigArgs
+        | None = None,
         logits_processors: list[str | type[LogitsProcessor]] | None = None,
         **kwargs: Any,
     ) -> None:
@@ -367,6 +373,7 @@ class LLM:
             profiler_config=profiler_config_instance,
             attention_config=attention_config_instance,
             compilation_config=compilation_config_instance,
+            quantization_config=quantization_config,
             logits_processors=logits_processors,
             **kwargs,
         )

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/fla/ops/utils.py

@@ -154,9 +154,13 @@ is_nvidia_hopper = is_nvidia and (
 )
 use_cuda_graph = is_nvidia and os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1"
 is_gather_supported = hasattr(triton.language, "gather")
-is_tma_supported = (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9) and (
-    hasattr(triton.language, "_experimental_make_tensor_descriptor")
-    or hasattr(triton.language, "make_tensor_descriptor")
+is_tma_supported = (
+    is_nvidia_hopper
+    and os.getenv("FLA_USE_TMA", "0") == "1"
+    and (
+        hasattr(triton.language, "_experimental_make_tensor_descriptor")
+        or hasattr(triton.language, "make_tensor_descriptor")
+    )
 )
 
 

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

@@ -79,11 +79,8 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
             RoutingMethodType.Default,
             RoutingMethodType.DeepSeekV3,
             RoutingMethodType.Llama4,
-            # NOTE: TRTLLM Kernel has issue with Qwen3.5 router.
-            # Re-enable once the issue is resolved.
-            # https://github.com/vllm-project/vllm/issues/37591
-            # RoutingMethodType.Renormalize,
-            # RoutingMethodType.RenormalizeNaive
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
         ]
 
     @staticmethod

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,
         )
@@ -277,13 +297,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        """Monolithic kernels need to express router support.
-        Renormalize/RenormalizeNaive are excluded: the monolithic kernel's
-        internal routing for these methods produces output uncorrelated
-        with the modular kernel's output and with Triton kernel's output
-        for Qwen3.5-35B-A3B-FP8.
-        See: https://github.com/vllm-project/vllm/issues/37591
-        """
+        """Monolithic kernels need to express router support."""
         # NOTE(dbari): TopK routing could also be enabled, but need to validate models
         # NOTE(dbari): Default is not implemented and should not be enabled until it is
 
@@ -295,6 +309,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
             return routing_method in [
                 RoutingMethodType.DeepSeekV3,
                 RoutingMethodType.Simulated,
+                RoutingMethodType.Renormalize,
+                RoutingMethodType.RenormalizeNaive,
             ]
         elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
             # NOTE(dbari): as above, potentially allow others here.
@@ -302,6 +318,8 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
                 RoutingMethodType.DeepSeekV3,
                 RoutingMethodType.Llama4,
                 RoutingMethodType.Simulated,
+                RoutingMethodType.Renormalize,
+                RoutingMethodType.RenormalizeNaive,
             ]
         else:
             raise ValueError("Unsupported quantization scheme.")
@@ -324,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
@@ -344,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(
@@ -369,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/layernorm.py

@@ -241,8 +241,12 @@ class RMSNorm(CustomOp):
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """PyTorch-native implementation equivalent to forward()."""
         if residual is None:
+            # TODO(luka): address the weight=None passing issue more generally
             return ir.ops.rms_norm(
-                x, self.weight.data, self.variance_epsilon, self.variance_size_override
+                x,
+                self.weight.data if self.has_weight else None,
+                self.variance_epsilon,
+                self.variance_size_override,
             )
 
         return self.forward_static(

vllm/model_executor/layers/linear.py

@@ -60,7 +60,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [
     "ModelOptFp8PbWoLinearMethod",
     "QuarkLinearMethod",
     "ModelOptNvFp4LinearMethod",
-    "PetitNvFp4LinearMethod",
 ]
 
 

vllm/model_executor/layers/quantization/__init__.py

@@ -31,8 +31,14 @@ QuantizationMethods = Literal[
     "inc",
     "mxfp4",
     "mxfp8",
-    "petit_nvfp4",
     "cpu_awq",
+    "online",
+    # Below are values of the OnlineQuantScheme enum, specified as strings to
+    # avoid circular import issues. This is here to provide a shortcut where
+    # the user can specify "LLM(..., quantization='fp8_per_tensor')" as
+    # shorthand for creating a more complicated online quant config object
+    "fp8_per_tensor",
+    "fp8_per_block",
 ]
 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
 
@@ -41,7 +47,6 @@ DEPRECATED_QUANTIZATION_METHODS = [
     "fbgemm_fp8",
     "fp_quant",
     "experts_int8",
-    "petit_nvfp4",
 ]
 
 # The customized quantization methods which will be added to this dict.
@@ -103,6 +108,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         raise ValueError(f"Invalid quantization method: {quantization}")
 
     # lazy import to avoid triggering `torch.compile` too early
+    from vllm.config.quantization import OnlineQuantScheme
     from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
 
     from .awq import AWQConfig
@@ -129,7 +135,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
     from .moe_wna16 import MoeWNA16Config
     from .mxfp4 import Mxfp4Config
     from .mxfp8 import Mxfp8Config
-    from .petit import PetitNvFp4Config
+    from .online.base import OnlineQuantizationConfig
     from .torchao import TorchAOConfig
 
     method_to_config: dict[str, type[QuantizationConfig]] = {
@@ -155,9 +161,21 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
         "inc": INCConfig,
         "mxfp4": Mxfp4Config,
         "mxfp8": Mxfp8Config,
-        "petit_nvfp4": PetitNvFp4Config,
         "cpu_awq": CPUAWQConfig,
+        "online": OnlineQuantizationConfig,
     }
+
+    # Below are values of the OnlineQuantScheme enum. This is here to provide
+    # a shortcut where the user can specify
+    # "LLM(..., quantization='fp8_per_tensor')" as shorthand for creating a
+    # more complicated online quant config object
+    for scheme in OnlineQuantScheme:
+        assert scheme.value not in method_to_config, (
+            f"Online quant scheme {scheme.value!r} conflicts with an "
+            f"existing quantization method"
+        )
+        method_to_config[scheme.value] = OnlineQuantizationConfig
+
     # Update the `method_to_config` with customized quantization methods.
     method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
 

vllm/model_executor/layers/quantization/fp8.py

@@ -497,6 +497,8 @@ class Fp8LinearMethod(LinearMethodBase):
         return self.fp8_linear.apply_weights(layer, x, bias)
 
 
+# TODO(future PR): remove this class in favor of
+# online/fp8.py::Fp8PerTensorOnlineLinearMethod
 class Fp8OnlineLinearMethod(Fp8LinearMethod):
     """Online version of Fp8LinearMethod which loads a full precision checkpoint
     and quantizes weights during loading."""
@@ -919,6 +921,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
         )
 
 
+# TODO(future PR): remove this class in favor of
+# online/fp8.py::Fp8PerTensorOnlineMoEMethod
 class Fp8OnlineMoEMethod(Fp8MoEMethod):
     """MoE method for online FP8 quantization.
     Supports loading quantized FP16/BF16 model checkpoints with dynamic

vllm/model_executor/layers/quantization/online/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/model_executor/layers/quantization/online/base.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+
+from vllm.config.quantization import (
+    OnlineQuantizationConfigArgs,
+    OnlineQuantScheme,
+)
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoE,
+)
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
+    UnquantizedFusedMoEMethod,
+)
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization import QuantizationMethods
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
+    should_ignore_layer,
+)
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PerBlockOnlineLinearMethod,
+    Fp8PerBlockOnlineMoEMethod,
+    Fp8PerTensorOnlineLinearMethod,
+    Fp8PerTensorOnlineMoEMethod,
+)
+
+
+class OnlineQuantizationConfig(QuantizationConfig):
+    """Model-level config class for online quantization (quantize fp16/bf16 weights
+    during model loading, without requiring a pre-quantized checkpoint)."""
+
+    def __init__(
+        self,
+        args: OnlineQuantizationConfigArgs,
+    ) -> None:
+        super().__init__()
+        if (
+            args.global_scheme is None
+            and args.linear_scheme_override is None
+            and args.moe_scheme_override is None
+        ):
+            raise ValueError(
+                "OnlineQuantizationConfig requires at least one of "
+                "global_scheme, linear_scheme_override, or "
+                "moe_scheme_override to be set."
+            )
+        self.args = args
+        self.quant_scheme = args.global_scheme
+        self.ignored_layers: list[str] = args.ignore
+
+    @classmethod
+    def get_name(cls) -> QuantizationMethods:
+        return "online"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # Note: as more online quant schemes will be added, this
+        # value will become the minimum across all supported schemes.
+        return 75
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "OnlineQuantizationConfig":
+        raise NotImplementedError(
+            "OnlineQuantizationConfig does not support loading from a "
+            "checkpoint config. Use quantization_config or "
+            "quantization='fp8_per_tensor'/'fp8_per_block' instead."
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> "QuantizeMethodBase | None":
+        if isinstance(layer, LinearBase):
+            if should_ignore_layer(
+                prefix,
+                ignore=self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+
+            linear_scheme = self.args.linear_scheme_override or self.args.global_scheme
+            if linear_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+                return Fp8PerBlockOnlineLinearMethod()
+            else:
+                return Fp8PerTensorOnlineLinearMethod()
+        elif isinstance(layer, FusedMoE):
+            if should_ignore_layer(
+                prefix,
+                ignore=self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedFusedMoEMethod(layer.moe_config)
+
+            moe_scheme = self.args.moe_scheme_override or self.args.global_scheme
+            if moe_scheme == OnlineQuantScheme.FP8_PER_BLOCK:
+                return Fp8PerBlockOnlineMoEMethod(layer=layer)
+            else:
+                return Fp8PerTensorOnlineMoEMethod(layer=layer)
+        return None

vllm/model_executor/layers/quantization/online/fp8.py

@@ -0,0 +1,632 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch.nn import Module
+
+if TYPE_CHECKING:
+    import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+    from vllm.model_executor.layers.fused_moe import FusedMoE
+    from vllm.model_executor.layers.fused_moe.config import (
+        FusedMoEConfig,
+        FusedMoEQuantConfig,
+    )
+    from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
+from vllm.model_executor.layers.fused_moe import (
+    FusedMoEMethodBase,
+)
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+    select_fp8_moe_backend,
+)
+from vllm.model_executor.layers.linear import (
+    LinearMethodBase,
+)
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    W8A8BlockFp8LinearOp,
+    maybe_post_process_fp8_weight_block,
+    process_fp8_weight_block_strategy,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    kFp8Dynamic128Sym,
+    kFp8DynamicTensorSym,
+    kFp8DynamicTokenSym,
+    kFp8Static128BlockSym,
+    kFp8StaticTensorSym,
+)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    cutlass_block_fp8_supported,
+    cutlass_fp8_supported,
+)
+from vllm.model_executor.model_loader.reload.layerwise import (
+    initialize_online_processing,
+)
+from vllm.model_executor.parameter import ModelWeightParameter
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import is_deep_gemm_supported, per_block_cast_to_fp8
+
+# ---------------------------------------------------------------------------
+# Online FP8 Linear Methods
+# ---------------------------------------------------------------------------
+
+
+class _Fp8OnlineLinearBase(LinearMethodBase):
+    """Shared base for online FP8 linear methods. Loads fp16/bf16 checkpoint
+    weights onto meta device and materializes them just-in-time."""
+
+    uses_meta_device: bool = True
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        layer.logical_widths = output_partition_sizes
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.orig_dtype = params_dtype
+        layer.weight_block_size = None
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                device="meta",  # materialized and processed during loading
+                dtype=params_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        initialize_online_processing(layer)
+
+
+class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
+    """Online tensorwise FP8 linear quantization.
+    Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
+
+    def __init__(self):
+        self.out_dtype = torch.get_default_dtype()
+
+        # Use per-token quantization for better perf if dynamic and cutlass
+        if cutlass_fp8_supported():
+            activation_quant_key = kFp8DynamicTokenSym
+        else:
+            activation_quant_key = kFp8DynamicTensorSym
+
+        self.fp8_linear = init_fp8_linear_kernel(
+            activation_quant_key=activation_quant_key,
+            weight_quant_key=kFp8StaticTensorSym,
+            out_dtype=torch.get_default_dtype(),
+            module_name=self.__class__.__name__,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        layer.input_scale = None
+        qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
+
+        # Update layer with new values.
+        replace_parameter(layer, "weight", qweight.t().data)
+        replace_parameter(layer, "weight_scale", weight_scale.data)
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # if batch invariant mode is enabled, use BF16 dequant
+        if envs.VLLM_BATCH_INVARIANT:
+            weight_fp8 = layer.weight.to(torch.bfloat16)
+            weight_scale = layer.weight_scale.to(torch.bfloat16)
+            if weight_scale.numel() == 1:
+                # Per-tensor: simple scalar multiplication
+                weight_bf16 = weight_fp8 * weight_scale
+            else:
+                # Multiple scales (fused modules like QKV)
+                if (
+                    weight_scale.dim() == 1
+                    and weight_scale.shape[0] == weight_fp8.shape[0]
+                ):
+                    # Per-row scaling
+                    weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
+                else:
+                    # Fallback
+                    weight_bf16 = weight_fp8 * weight_scale
+            return torch.nn.functional.linear(x, weight_bf16.t(), bias)
+
+        return self.fp8_linear.apply_weights(layer, x, bias)
+
+
+class Fp8PerBlockOnlineLinearMethod(_Fp8OnlineLinearBase):
+    """Online blockwise FP8 linear quantization.
+    Loads fp16/bf16 weights and quantizes them per-block during loading."""
+
+    def __init__(self):
+        self.out_dtype = torch.get_default_dtype()
+        self.weight_block_size = [128, 128]
+
+        self.use_deep_gemm = is_deep_gemm_supported()
+        self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
+        self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
+
+        self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
+            weight_group_shape=GroupShape(*self.weight_block_size),
+            act_quant_group_shape=GroupShape(1, self.weight_block_size[0]),
+            cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
+            use_aiter_and_is_supported=self.use_aiter_and_is_supported,
+            use_deep_gemm=self.use_deep_gemm,
+        )
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        output_size: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        super().create_weights(
+            layer,
+            input_size_per_partition,
+            output_partition_sizes,
+            input_size,
+            output_size,
+            params_dtype,
+            **extra_weight_attrs,
+        )
+        layer.weight_block_size = self.weight_block_size
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        layer.input_scale = None
+        block_size = self.weight_block_size
+
+        qweight, weight_scale_inv = per_block_cast_to_fp8(
+            layer.weight, block_size=block_size, use_ue8m0=False
+        )
+
+        qweight, weight_scale_inv = process_fp8_weight_block_strategy(
+            qweight, weight_scale_inv
+        )
+
+        replace_parameter(layer, "weight", qweight.data)
+        replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
+
+        maybe_post_process_fp8_weight_block(layer)
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        assert self.weight_block_size is not None
+
+        # Note: batch invariance already handled in the function below
+        return self.w8a8_block_fp8_linear.apply(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale_inv,
+            input_scale=layer.input_scale,
+            bias=bias,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Online FP8 MoE Methods
+# ---------------------------------------------------------------------------
+
+
+class _Fp8OnlineMoEBase(FusedMoEMethodBase):
+    """Shared base for online FP8 MoE methods. Loads fp16/bf16 checkpoint
+    weights onto meta device and materializes them just-in-time."""
+
+    uses_meta_device: bool = True
+
+    # Declared here for mypy; actual values are set in __init__.
+    fp8_backend: "Fp8MoeBackend"
+    experts_cls: "type[mk.FusedMoEExperts] | None"
+    weight_scale_name: str
+    weight_block_size: list[int] | None
+    moe: "FusedMoEConfig"
+    is_monolithic: bool
+    moe_quant_config: "FusedMoEQuantConfig | None"
+    moe_kernel: "mk.FusedMoEKernel | None"
+
+    def __init__(
+        self,
+        *,
+        weight_block_size: list[int] | None,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(layer.moe_config)
+        self.weight_block_size = weight_block_size
+        self.block_quant: bool = self.weight_block_size is not None
+        self.weight_scale_name = (
+            "weight_scale_inv" if self.block_quant else "weight_scale"
+        )
+
+        # Set weight key and activation key for kernel compatibility
+        if self.block_quant:
+            weight_key = kFp8Static128BlockSym
+            activation_key = kFp8Dynamic128Sym
+        else:
+            weight_key = kFp8StaticTensorSym
+            activation_key = kFp8DynamicTensorSym
+
+        # Select Fp8 MoE backend
+        self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
+            config=self.moe,
+            weight_key=weight_key,
+            activation_key=activation_key,
+            allow_vllm_cutlass=False,
+        )
+
+    def create_weights(
+        self,
+        layer: 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.orig_dtype = params_dtype
+        layer.weight_block_size = None
+
+        # WEIGHTS
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size,
+                device="meta",
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition,
+                device="meta",  # materialized and processed during loading
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # BIASES (for models like GPT-OSS that have biased MoE)
+        if self.moe.has_bias:
+            w13_bias = torch.nn.Parameter(
+                torch.zeros(
+                    num_experts,
+                    2 * intermediate_size_per_partition,
+                    device="meta",  # materialized and processed during loading
+                    dtype=layer.orig_dtype,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w13_bias", w13_bias)
+            set_weight_attrs(w13_bias, extra_weight_attrs)
+
+            w2_bias = torch.nn.Parameter(
+                torch.zeros(
+                    num_experts,
+                    hidden_size,
+                    device="meta",  # materialized and processed during loading
+                    dtype=layer.orig_dtype,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w2_bias", w2_bias)
+            set_weight_attrs(w2_bias, extra_weight_attrs)
+
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
+
+        initialize_online_processing(layer)
+
+    def _setup_kernel(
+        self,
+        layer: "FusedMoE",
+        w13: torch.Tensor,
+        w2: torch.Tensor,
+        w13_scale: torch.Tensor,
+        w2_scale: torch.Tensor,
+        w13_input_scale: torch.Tensor | None,
+        w2_input_scale: torch.Tensor | None,
+    ) -> None:
+        from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+            convert_to_fp8_moe_kernel_format,
+            make_fp8_moe_kernel,
+        )
+
+        # Shuffle weights to runtime format.
+        w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
+            fp8_backend=self.fp8_backend,
+            layer=layer,
+            w13=w13,
+            w2=w2,
+            w13_scale=w13_scale,
+            w2_scale=w2_scale,
+            w13_input_scale=w13_input_scale,
+            w2_input_scale=w2_input_scale,
+        )
+
+        # Replace parameters with updated versions. Note that this helper
+        # function ensures the replacement is compatible with RL weight reloads.
+        replace_parameter(layer, "w13_weight", w13)
+        replace_parameter(layer, "w2_weight", w2)
+        replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
+        replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
+
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        if self.moe_quant_config:
+            assert self.experts_cls is not None
+            self.moe_kernel = make_fp8_moe_kernel(
+                moe_quant_config=self.moe_quant_config,
+                moe_config=self.moe,
+                fp8_backend=self.fp8_backend,
+                experts_cls=self.experts_cls,
+                routing_tables=layer._maybe_init_expert_routing_tables(),
+                shared_experts=layer.shared_experts,
+            )
+
+    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":
+        from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+            make_fp8_moe_quant_config,
+        )
+
+        w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
+        w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
+        a1_scale = layer.w13_input_scale
+        a2_scale = layer.w2_input_scale
+
+        quant_config = make_fp8_moe_quant_config(
+            fp8_backend=self.fp8_backend,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=self.weight_block_size,
+        )
+
+        # Inject biases into the quant config if the model has them
+        # (e.g. GPT-OSS biased MoE)
+        if quant_config is not None and self.moe.has_bias:
+            w13_bias = getattr(layer, "w13_bias", None)
+            w2_bias = getattr(layer, "w2_bias", None)
+            if w13_bias is not None:
+                quant_config._w1.bias = w13_bias
+            if w2_bias is not None:
+                quant_config._w2.bias = w2_bias
+
+        return quant_config
+
+    @property
+    def supports_eplb(self) -> bool:
+        return True
+
+    def apply_monolithic(
+        self,
+        layer: "FusedMoE",
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+    ) -> torch.Tensor | tuple[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 | tuple[torch.Tensor, torch.Tensor]:
+        assert not self.is_monolithic
+        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,
+        )
+
+
+class Fp8PerTensorOnlineMoEMethod(_Fp8OnlineMoEBase):
+    """Online tensorwise FP8 MoE quantization.
+    Loads fp16/bf16 weights and quantizes them per-tensor during loading."""
+
+    def __init__(
+        self,
+        *,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(
+            weight_block_size=None,
+            layer=layer,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        # TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        # If checkpoint is fp16, quantize in place.
+        fp8_dtype = current_platform.fp8_dtype()
+        w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
+        w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
+        w13_scale = torch.ones(
+            layer.num_experts, device=w13.device, dtype=torch.float32
+        )
+        w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)
+        layer.w13_input_scale = None
+        layer.w2_input_scale = None
+
+        for expert in range(layer.local_num_experts):
+            w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
+                layer.w13_weight[expert, :, :]
+            )
+            w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
+                layer.w2_weight[expert, :, :]
+            )
+
+        # Shuffle weights to runtime format and setup kernel.
+        self._setup_kernel(
+            layer,
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            w13_input_scale=layer.w13_input_scale,
+            w2_input_scale=layer.w2_input_scale,
+        )
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True
+
+
+class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):
+    """Online blockwise FP8 MoE quantization.
+    Loads fp16/bf16 weights and quantizes them per-block during loading."""
+
+    def __init__(
+        self,
+        *,
+        layer: torch.nn.Module,
+    ):
+        super().__init__(
+            weight_block_size=[128, 128],
+            layer=layer,
+        )
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        fp8_dtype = current_platform.fp8_dtype()
+        w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
+        w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
+
+        block_size = self.weight_block_size
+        assert block_size is not None
+        block_n, block_k = block_size
+
+        # Create block-shaped scales (computed here rather than in
+        # create_weights because online quant doesn't need them until now).
+        num_experts = layer.local_num_experts
+        _, w13_out, w13_in = layer.w13_weight.shape
+        _, w2_out, w2_in = layer.w2_weight.shape
+
+        w13_scale = torch.ones(
+            num_experts,
+            (w13_out + block_n - 1) // block_n,
+            (w13_in + block_k - 1) // block_k,
+            dtype=torch.float32,
+            device=w13.device,
+        )
+        w2_scale = torch.ones(
+            num_experts,
+            (w2_out + block_n - 1) // block_n,
+            (w2_in + block_k - 1) // block_k,
+            dtype=torch.float32,
+            device=w2.device,
+        )
+
+        for expert in range(num_experts):
+            w13[expert], w13_scale[expert] = per_block_cast_to_fp8(
+                layer.w13_weight[expert],
+                block_size=block_size,
+                use_ue8m0=False,
+            )
+            w2[expert], w2_scale[expert] = per_block_cast_to_fp8(
+                layer.w2_weight[expert],
+                block_size=block_size,
+                use_ue8m0=False,
+            )
+
+        layer.weight_block_size = block_size
+
+        # Shuffle weights to runtime format and setup kernel.
+        self._setup_kernel(
+            layer,
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            layer.w13_input_scale,
+            layer.w2_input_scale,
+        )
+
+        # Prevent duplicate processing (e.g., during weight reload)
+        layer._already_called_process_weights_after_loading = True

vllm/model_executor/layers/quantization/petit.py

@@ -1,319 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py
-
-from typing import Any
-
-import regex as re
-import torch
-from torch.nn.parameter import Parameter
-
-from vllm.logger import init_logger
-from vllm.model_executor.layers.attention import Attention
-from vllm.model_executor.layers.linear import (
-    LinearBase,
-    LinearMethodBase,
-    UnquantizedLinearMethod,
-)
-from vllm.model_executor.layers.quantization import QuantizationMethods
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-)
-from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
-from vllm.model_executor.layers.quantization.utils.petit_utils import (
-    apply_petit_nvfp4_linear,
-    prepare_nvfp4_layer_for_petit,
-    verify_petit_nvfp4_supported,
-)
-from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
-from vllm.model_executor.parameter import ModelWeightParameter, PerTensorScaleParameter
-from vllm.platforms import current_platform
-
-# Initialize logger for the module
-logger = init_logger(__name__)
-
-
-# Configuration class to support the NVFP4 quantized model
-# generated by the ModelOpt quantization tool
-class PetitNvFp4Config(QuantizationConfig):
-    """Config class for Petit FP4."""
-
-    def __init__(
-        self,
-        is_checkpoint_nvfp4_serialized: bool = False,
-        kv_cache_quant_algo: str | None = None,
-        group_size: int | None = None,
-        exclude_modules: list[str] | None = None,
-    ) -> None:
-        self._check_hardware_support()
-        self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
-        if is_checkpoint_nvfp4_serialized:
-            logger.warning(
-                "Detected nvfp4 checkpoint. Please note that the "
-                "format is experimental and subject to change."
-            )
-        self.group_size = group_size
-        self.kv_cache_quant_algo = kv_cache_quant_algo
-        self.exclude_modules = exclude_modules
-
-    def _check_hardware_support(self) -> None:
-        """
-        Verifies that the current hardware is supported by the Petit backend.
-        This backend is specifically designed for AMD GPUs and is not
-        supported on the CUDA platform.
-        """
-        # This check ensures the code is NOT running on an NVIDIA GPU.
-        if current_platform.is_cuda():
-            raise ValueError(
-                "The 'petit' quantization backend is designed for AMD GPUs "
-                "and is not supported on the CUDA platform. For NVIDIA GPUs, "
-                "please use a different quantization method such as FP8, AWQ, "
-                "or GPTQ."
-            )
-
-    @classmethod
-    def get_name(cls) -> QuantizationMethods:
-        return "petit_nvfp4"
-
-    @classmethod
-    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
-        return [torch.bfloat16, torch.half]
-
-    @classmethod
-    def get_min_capability(cls) -> int:
-        # Petit supports the gfx90a and gfx942 GPUs
-        return 90
-
-    @classmethod
-    def get_config_filenames(cls) -> list[str]:
-        return ["hf_quant_config.json"]
-
-    @classmethod
-    def from_config(cls, config: dict[str, Any]) -> "PetitNvFp4Config":
-        qc = cls.get_from_keys(config, ["quantization"])
-
-        quant_method_raw = qc.get("quant_algo")
-        if not isinstance(quant_method_raw, str) or not quant_method_raw:
-            raise ValueError("Missing or invalid 'quant_algo' in quantization config.")
-        quant_method = quant_method_raw.upper()
-
-        group_size_raw = qc.get("group_size")
-        if not isinstance(group_size_raw, int):
-            raise ValueError(
-                "Missing or invalid 'group_size' (int) in hf_quant_config.json."
-            )
-        group_size = group_size_raw
-
-        verify_petit_nvfp4_supported(quant_method, group_size)
-
-        kv_cache_quant_algo_raw = qc.get("kv_cache_quant_algo") or "auto"
-        if not isinstance(kv_cache_quant_algo_raw, str):
-            raise ValueError("'kv_cache_quant_algo' must be a string if provided.")
-        kv_cache_quant_algo = kv_cache_quant_algo_raw
-
-        exclude_raw = qc.get("exclude_modules", [])
-        if exclude_raw is None:
-            exclude_modules: list[str] = []
-        elif isinstance(exclude_raw, list) and all(
-            isinstance(x, str) for x in exclude_raw
-        ):
-            exclude_modules = exclude_raw
-        else:
-            raise ValueError("'exclude_modules' must be a list[str] (or omitted).")
-
-        is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method
-
-        return cls(
-            is_checkpoint_nvfp4_serialized=is_checkpoint_nvfp4_serialized,
-            kv_cache_quant_algo=kv_cache_quant_algo,
-            group_size=group_size,
-            exclude_modules=exclude_modules,
-        )
-
-    @classmethod
-    def override_quantization_method(
-        cls, hf_quant_cfg, user_quant
-    ) -> QuantizationMethods | None:
-        if not current_platform.is_rocm():
-            return None
-
-        qc = hf_quant_cfg.get("quantization", hf_quant_cfg)
-        algo = (qc.get("quant_algo") or qc.get("quant_method") or "").upper()
-        if algo in ("NVFP4", "MODELOPT_FP4", "MODELOPT"):
-            return cls.get_name()  # "petit_nvfp4"
-        return None
-
-    @classmethod
-    def is_petit_nvfp4_compatible(cls, quant_config: dict[str, Any]) -> bool:
-        qc = quant_config.get("quantization", quant_config)
-        algo = (qc.get("quant_algo") or qc.get("quant_method") or "").upper()
-        return algo == "NVFP4"
-
-    def is_layer_excluded(self, prefix: str, exclude_modules: list[str]) -> bool:
-        for pattern in exclude_modules:
-            regex_str = pattern.replace(".", r"\.").replace("*", r".*")
-            if re.fullmatch(regex_str, prefix):
-                return True
-        return False
-
-    def get_quant_method(
-        self, layer: torch.nn.Module, prefix: str
-    ) -> "QuantizeMethodBase | None":
-        exclude = self.require_exclude_modules()
-
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped(prefix, exclude) or self.is_layer_excluded(
-                prefix, exclude
-            ):
-                return UnquantizedLinearMethod()
-            return PetitNvFp4LinearMethod(self)
-        elif isinstance(layer, Attention):
-            return PetitFp8KVCacheMethod(self)
-        return None
-
-    def get_scaled_act_names(self) -> list[str]:
-        return []
-
-    def require_group_size(self) -> int:
-        if self.group_size is None:
-            logger.warning("group_size not set; defaulting to 16 for NVFP4.")
-            return 16
-        return self.group_size
-
-    def require_kv_cache_quant_algo(self) -> str:
-        return self.kv_cache_quant_algo or "auto"
-
-    def require_exclude_modules(self) -> list[str]:
-        return list(self.exclude_modules or [])
-
-
-class PetitFp8KVCacheMethod(BaseKVCacheMethod):
-    """
-    Supports loading kv-cache scaling factors from FP8 checkpoints.
-    """
-
-    def __init__(self, quant_config: PetitNvFp4Config):
-        super().__init__(quant_config)
-
-
-class PetitNvFp4LinearMethod(LinearMethodBase):
-    """Linear method for NVFP4.
-    Supports loading NVFP4 checkpoints with the following structure:
-
-    |Tensor Name           | datatype      |  shape      |
-    |----------------------------------------------------|
-    |input_scale           | torch.float32 | scalar      |
-    |weight                | NVFP4(SE2M1)  | [1, X, y/2] |
-    |weight_scale          | FP8-E4M3      | [X, Y]      |
-    |weight_scale_2        | torch.float32 | scalar      |
-
-    The weights are quantized per block of 16 elements.
-    Args: quant_config: The ModelOpt quantization config.
-    """
-
-    def __init__(self, quant_config: PetitNvFp4Config):
-        self.quant_config = quant_config
-
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        input_size_per_partition: int,
-        output_partition_sizes: list[int],
-        input_size: int,
-        output_size: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        del input_size, output_size
-        if not self.quant_config.is_checkpoint_nvfp4_serialized:
-            raise ValueError(
-                "NVFP4 quantization was selected, "
-                " dynamic quantization is not supported."
-            )
-
-        output_size_per_partition = sum(output_partition_sizes)
-        weight_loader = extra_weight_attrs.get("weight_loader")
-
-        layer.logical_widths = output_partition_sizes
-
-        layer.input_size_per_partition = input_size_per_partition
-        layer.output_size_per_partition = output_size_per_partition
-        if input_size_per_partition % 16 != 0:
-            raise ValueError(
-                "Unsupported model when in features size is not multiple of 16"
-            )
-
-        weight_dtype = (
-            torch.float8_e4m3fn
-            if self.quant_config.is_checkpoint_nvfp4_serialized
-            else params_dtype
-        )
-
-        weight = ModelWeightParameter(
-            data=torch.empty(
-                # 2 fp4 data is packed in one uint8 in the input dimension
-                output_size_per_partition,
-                input_size_per_partition // 2,
-                dtype=torch.uint8,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight", weight)
-
-        input_scale = PerTensorScaleParameter(
-            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
-            weight_loader=weight_loader,
-        )
-
-        layer.register_parameter("input_scale", input_scale)
-
-        weight_scale_2 = PerTensorScaleParameter(
-            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight_scale_2", weight_scale_2)
-
-        group_size = self.quant_config.require_group_size()
-        weight_scale = ModelWeightParameter(
-            data=torch.empty(
-                output_size_per_partition,
-                input_size_per_partition // group_size,
-                dtype=weight_dtype,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-
-        layer.register_parameter("weight_scale", weight_scale)
-
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        input_scale_2 = layer.input_scale.max().to(torch.float32)
-        weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
-        layer.input_scale = Parameter(input_scale_2, requires_grad=False)
-        layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False)
-        layer.alpha = Parameter(
-            layer.input_scale * layer.weight_scale_2, requires_grad=False
-        )
-
-        prepare_nvfp4_layer_for_petit(layer)
-        del layer.input_scale
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        return apply_petit_nvfp4_linear(
-            input=x,
-            weight=layer.weight,
-            weight_scale=layer.weight_scale,
-            weight_scale_2=layer.weight_scale_2,
-            size_n=layer.output_size_per_partition,
-            size_k=layer.input_size_per_partition,
-            bias=bias,
-        )

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/layers/quantization/utils/petit_utils.py

@@ -1,120 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from typing import TYPE_CHECKING
-
-import torch
-
-# TYPE_CHECKING is used for static type analysis to prevent circular imports.
-if TYPE_CHECKING:
-    from types import ModuleType
-
-# 1. Create a global variable as a placeholder for the module
-_petit_kernel: "ModuleType | None" = None
-
-_PETIT_INSTALL_MSG = (
-    "Petit is not installed. Please install it with `pip install petit-kernel`."
-)
-
-
-def _import_petit_kernel() -> "ModuleType":
-    """
-    A helper function to handle the lazy import.
-    The first time this function is called, it will import the petit_kernel
-    library and store it in the global _petit_kernel variable.
-    Subsequent calls will return the already-loaded module directly.
-    """
-    global _petit_kernel
-    if _petit_kernel is not None:
-        return _petit_kernel
-
-    try:
-        import petit_kernel
-
-        _petit_kernel = petit_kernel
-        return _petit_kernel
-    except ImportError:
-        # The 'from None' syntax prevents chaining the original ImportError,
-        # making the traceback cleaner.
-        raise ImportError(_PETIT_INSTALL_MSG) from None
-
-
-def _check_petit_nvfp4_supported(
-    quant_method: str, group_size: int | None
-) -> tuple[bool, str | None]:
-    if quant_method != "NVFP4":
-        return (
-            False,
-            (
-                "Petit currently only supports: NVFP4 quantizations in sglang. "
-                "Please check the `hf_quant_config.json` file for your model's "
-                "quant configuration."
-            ),
-        )
-    if group_size is not None and group_size != 16:
-        return (
-            False,
-            "Petit currently only supports: group_size=16 quantizations.",
-        )
-    return (True, None)
-
-
-def verify_petit_nvfp4_supported(quant_method: str, group_size: int | None) -> None:
-    supported, error_msg = _check_petit_nvfp4_supported(quant_method, group_size)
-    if not supported:
-        assert error_msg is not None
-        raise ValueError(error_msg)
-
-
-def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None:
-    # 2. Call _import_petit_kernel() to trigger (or get) the import.
-    petit_kernel = _import_petit_kernel()
-
-    # Repack weights to petit format
-    part_size_n = layer.output_size_per_partition
-    part_size_k = layer.input_size_per_partition
-    qweight = layer.weight.view(torch.int32).contiguous()
-
-    # 3. Call functions through the imported module variable.
-    petit_qweight = petit_kernel.repack_nvfp4(
-        qweight, size_n=part_size_n, size_k=part_size_k
-    )
-    layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False)
-
-    # Permute scales
-    weight_scale = petit_kernel.process_nvfp4_scales(
-        scales=layer.weight_scale, size_k=part_size_k, size_n=part_size_n
-    )
-    layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
-
-
-def apply_petit_nvfp4_linear(
-    input: torch.Tensor,
-    weight: torch.Tensor,
-    weight_scale: torch.Tensor,
-    weight_scale_2: torch.Tensor,
-    size_n: int,
-    size_k: int,
-    bias: torch.Tensor | None = None,
-) -> torch.Tensor:
-    # Trigger (or get) the import here as well.
-    petit_kernel = _import_petit_kernel()
-
-    reshaped_x = input.reshape(-1, input.shape[-1])
-    out_shape = input.shape[:-1] + (size_n,)
-
-    # TODO: Use auto-tuning to find the performant solution_id
-    # Call the function via the module variable.
-    output = petit_kernel.mul_nvfp4_a16(
-        a=reshaped_x,
-        b=weight,
-        s=weight_scale,
-        global_scale=weight_scale_2,
-        size_m=reshaped_x.size(0),
-        size_n=size_n,
-        size_k=size_k,
-        solution_id=-1,
-    )
-    if bias is not None:
-        output.add_(bias)  # In-place add
-
-    return output.reshape(out_shape)

vllm/model_executor/model_loader/weight_utils.py

@@ -296,6 +296,13 @@ def get_quant_config(
         )
 
     if hf_quant_config is not None:
+        if model_config.quantization_config is not None:
+            raise ValueError(
+                "Setting `quantization_config` for online "
+                "quantization when the model checkpoint already "
+                "has a `quantization_config` is not supported"
+            )
+
         # For modelopt_mixed, config.json's quantization_config may or may
         # not contain the per-layer quantized_layers map.  Newer checkpoints
         # embed it directly; older ones keep it only in hf_quant_config.json.
@@ -319,6 +326,12 @@ def get_quant_config(
     quantization_config_file = hf_overrides.get("quantization_config_file", None)
     if quantization_config_file is not None:
         if hasattr(quant_cls, "from_config_file"):
+            if model_config.quantization_config is not None:
+                raise ValueError(
+                    "Setting `quantization_config` for online "
+                    "quantization when the model checkpoint already "
+                    "has a `quantization_config` is not supported"
+                )
             return quant_cls.from_config_file(quantization_config_file)
         else:
             raise NotImplementedError(
@@ -329,6 +342,12 @@ def get_quant_config(
     quantization_config_json = hf_overrides.get("quantization_config_dict_json", None)
     if quantization_config_json is not None:
         if hasattr(quant_cls, "from_config_dict_json"):
+            if model_config.quantization_config is not None:
+                raise ValueError(
+                    "Setting `quantization_config` for online "
+                    "quantization when the model checkpoint already "
+                    "has a `quantization_config` is not supported"
+                )
             return quant_cls.from_config_dict_json(quantization_config_json)
         else:
             raise NotImplementedError(
@@ -337,6 +356,19 @@ def get_quant_config(
                 f"{quant_cls}"
             )
 
+    # Online quantization doesn't read from checkpoint configs — it quantizes
+    # fp16/bf16 weights on the fly during loading.
+    if model_config.quantization_config is not None:
+        from vllm.config.quantization import OnlineQuantizationConfigArgs
+        from vllm.model_executor.layers.quantization.online.base import (
+            OnlineQuantizationConfig,
+        )
+
+        assert isinstance(
+            model_config.quantization_config, OnlineQuantizationConfigArgs
+        )
+        return OnlineQuantizationConfig(args=model_config.quantization_config)
+
     # Inflight BNB quantization
     if model_config.quantization == "bitsandbytes":
         return quant_cls.from_config({})

vllm/model_executor/models/arctic.py

@@ -16,7 +16,6 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
-from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
@@ -42,6 +41,7 @@ from vllm.transformers_utils.configs.arctic import ArcticConfig
 
 from .interfaces import SupportsPP, SupportsQuant
 from .utils import (
+    AutoWeightsLoader,
     extract_layer_index,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -49,8 +49,6 @@ from .utils import (
     maybe_prefix,
 )
 
-logger = init_logger(__name__)
-
 
 class ArcticMLP(nn.Module):
     def __init__(
@@ -384,6 +382,7 @@ class ArcticModel(nn.Module):
         cache_config = vllm_config.cache_config
         quant_config = vllm_config.quant_config
 
+        self.config = config
         self.vocab_size = config.vocab_size
         self.embed_tokens = VocabParallelEmbedding(
             self.vocab_size, config.hidden_size, org_num_embeddings=self.vocab_size
@@ -426,6 +425,116 @@ class ArcticModel(nn.Module):
         hidden_states = self.norm(hidden_states)
         return hidden_states
 
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+
+        mlp_params_mapping: list[tuple[str, str, int]] = []
+        expert_params_mapping: list[tuple[str, str, int]] = []
+
+        for layer in range(self.config.num_hidden_layers):
+            is_moe_layer = (layer + 1) % self.config.moe_layer_frequency == 0
+            if is_moe_layer and self.config.use_residual:
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.residual_mlp.w13.weight",
+                        f"layers.{layer}.residual_mlp.w1.weight",
+                        0,
+                    )
+                )
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.residual_mlp.w13.weight",
+                        f"layers.{layer}.residual_mlp.w3.weight",
+                        1,
+                    )
+                )
+
+            if is_moe_layer:
+                for expert_id in range(self.config.num_local_experts):
+                    expert_params_mapping.append(
+                        ("ws", f"experts.{expert_id}.w1.weight", expert_id)
+                    )
+                    expert_params_mapping.append(
+                        ("w2s", f"experts.{expert_id}.w2.weight", expert_id)
+                    )
+                    expert_params_mapping.append(
+                        ("ws", f"experts.{expert_id}.w3.weight", expert_id)
+                    )
+            else:
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
+                        f"layers.{layer}.block_sparse_moe.mlp.w1.weight",
+                        0,
+                    )
+                )
+                mlp_params_mapping.append(
+                    (
+                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
+                        f"layers.{layer}.block_sparse_moe.mlp.w3.weight",
+                        1,
+                    )
+                )
+
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                for param_name, weight_name, shard_id in mlp_params_mapping:
+                    if weight_name not in name:
+                        continue
+                    name = name.replace(weight_name, param_name)
+                    if is_pp_missing_parameter(name, self):
+                        continue
+                    param = params_dict[name]
+                    weight_loader = param.weight_loader
+                    weight_loader(param, loaded_weight, shard_id)
+                    break
+                else:
+                    for param_name, weight_name, shard_id in expert_params_mapping:
+                        if weight_name not in name:
+                            continue
+                        name = name.replace(weight_name, param_name)
+                        if is_pp_missing_parameter(name, self):
+                            continue
+                        param = params_dict[name]
+                        weight_loader = param.weight_loader
+                        weight_loader(
+                            param, loaded_weight, weight_name, expert_id=shard_id
+                        )
+                        break
+                    else:
+                        if name.endswith(".bias") and 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
+                        )
+                        weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
 
 class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
     packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
@@ -478,117 +587,8 @@ class ArcticForCausalLM(nn.Module, SupportsPP, SupportsQuant):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("qkv_proj", "q_proj", "q"),
-            ("qkv_proj", "k_proj", "k"),
-            ("qkv_proj", "v_proj", "v"),
-        ]
-
-        mlp_params_mapping: list[tuple[str, str, int]] = []
-        expert_params_mapping: list[tuple[str, str, int]] = []
-        num_layers = self.config.num_hidden_layers
-
-        for layer in range(num_layers):
-            mlp_params_mapping.append(
-                (
-                    f"layers.{layer}.residual_mlp.w13.weight",
-                    f"layers.{layer}.residual_mlp.w1.weight",
-                    0,
-                )
-            )
-            mlp_params_mapping.append(
-                (
-                    f"layers.{layer}.residual_mlp.w13.weight",
-                    f"layers.{layer}.residual_mlp.w3.weight",
-                    1,
-                )
-            )
-            if layer % 2 == 0:
-                # MLP layers
-                mlp_params_mapping.append(
-                    (
-                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
-                        f"layers.{layer}.block_sparse_moe.mlp.w1.weight",
-                        0,
-                    )
-                )
-                mlp_params_mapping.append(
-                    (
-                        f"layers.{layer}.block_sparse_moe.mlp.w13.weight",
-                        f"layers.{layer}.block_sparse_moe.mlp.w3.weight",
-                        1,
-                    )
-                )
-            else:
-                # MoE layers
-                for expert_id in range(self.config.num_local_experts):
-                    expert_params_mapping.append(
-                        ("ws", f"experts.{expert_id}.w1.weight", expert_id)
-                    )
-                    expert_params_mapping.append(
-                        ("w2s", f"experts.{expert_id}.w2.weight", expert_id)
-                    )
-                    expert_params_mapping.append(
-                        ("ws", f"experts.{expert_id}.w3.weight", expert_id)
-                    )
-
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-
-        logger.info(
-            "It will take ~10 minutes loading from the 16-bit weights. "
-            "Alternatively, use the prequantized 8-bit weights of arctic "
-            "and set load-format to `sharded_state` will accelerate loading."
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),
         )
-        for name, loaded_weight in weights:
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                for param_name, weight_name, shard_id in mlp_params_mapping:
-                    if weight_name not in name:
-                        continue
-                    name = name.replace(weight_name, param_name)
-                    if is_pp_missing_parameter(name, self):
-                        continue
-                    param = params_dict[name]
-                    weight_loader = param.weight_loader
-                    weight_loader(param, loaded_weight, shard_id)
-                    break
-                else:
-                    for param_name, weight_name, shard_id in expert_params_mapping:
-                        if weight_name not in name:
-                            continue
-                        name = name.replace(weight_name, param_name)
-                        if is_pp_missing_parameter(name, self):
-                            continue
-                        param = params_dict[name]
-                        weight_loader = param.weight_loader
-                        weight_loader(
-                            param, loaded_weight, weight_name, expert_id=shard_id
-                        )
-                        break
-                    else:
-                        if name.endswith(".bias") and 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
-                        )
-                        weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        return loader.load_weights(weights)

vllm/model_executor/models/deepseek_mtp.py

@@ -184,11 +184,16 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.config = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
         self.model = DeepSeekMultiTokenPredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
         # Set MoE hyperparameters
         self.set_moe_parameters()
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
 
     def set_moe_parameters(self):
         self.expert_weights = []
@@ -241,11 +246,16 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             ("gate_up_proj", "up_proj", 1),
             ("fused_qkv_a_proj", "q_a_proj", 0),
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
-            # Fused indexer wk + weights_proj
-            ("wk_weights_proj", "wk", 0),
-            ("wk_weights_proj", "weights_proj", 1),
         ]
 
+        if self.is_fp4_ckpt:
+            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+            indexer_fused_mapping = [
+                ("wk_weights_proj", "wk", 0),
+                ("wk_weights_proj", "weights_proj", 1),
+            ]
+            stacked_params_mapping.extend(indexer_fused_mapping)
+
         expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
             self,
             ckpt_gate_proj_name="gate_proj",

vllm/model_executor/models/deepseek_v2.py

@@ -625,6 +625,11 @@ class Indexer(nn.Module):
         super().__init__()
         self.vllm_config = vllm_config
         self.config = config
+        self.quant_config = quant_config
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
         # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
         self.topk_tokens = config.index_topk
         self.n_head = config.index_n_heads  # 64
@@ -639,18 +644,36 @@ class Indexer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.wq_b",
         )
-        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
-        # weights_proj does not get quantized, so we run both with quant_config=None
-        # wk may be upcasted from the default quant; experiments show fusion is always
-        # faster unless WK proj is in FP4, which is not the case for all known quants.
-        self.wk_weights_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [self.head_dim, self.n_head],
-            bias=False,
-            quant_config=None,
-            disable_tp=True,
-            prefix=f"{prefix}.wk_weights_proj",
-        )
+        if self.is_fp4_ckpt:
+            # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
+            # weights_proj does not get quantized,
+            # so we run both with quant_config=None
+            # wk may be upcasted from the default quant;
+            # experiments show fusion is always faster unless WK proj is in FP4,
+            # which is not the case for all known quants.
+            self.wk_weights_proj = MergedColumnParallelLinear(
+                hidden_size,
+                [self.head_dim, self.n_head],
+                bias=False,
+                quant_config=None,
+                disable_tp=True,
+                prefix=f"{prefix}.wk_weights_proj",
+            )
+        else:
+            self.wk = ReplicatedLinear(
+                hidden_size,
+                self.head_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.wk",
+            )
+            self.weights_proj = ReplicatedLinear(
+                hidden_size,
+                self.n_head,
+                bias=False,
+                quant_config=None,
+                prefix=f"{prefix}.weights_proj",
+            )
         self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
         self.softmax_scale = self.head_dim**-0.5
 
@@ -691,11 +714,14 @@ class Indexer(nn.Module):
         q_pe, q_nope = torch.split(
             q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
         )
-
-        # Fused wk + weights_proj: one GEMM, then split
-        kw, _ = self.wk_weights_proj(hidden_states)
-        k = kw[:, : self.head_dim]
-        weights_raw = kw[:, self.head_dim :]
+        if self.is_fp4_ckpt:
+            # Fused wk + weights_proj: one GEMM, then split
+            kw, _ = self.wk_weights_proj(hidden_states)
+            k = kw[:, : self.head_dim]
+            weights = kw[:, self.head_dim :]
+        else:
+            k, _ = self.wk(hidden_states)
+            weights, _ = self.weights_proj(hidden_states)
 
         k = self.k_norm(k)
         k_pe, k_nope = torch.split(
@@ -726,7 +752,7 @@ class Indexer(nn.Module):
         q_scale = q_scale.view(-1, self.n_head, 1)
 
         weights = (
-            weights_raw.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
+            weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
         )
         weights = weights.squeeze(-1)
 
@@ -1314,6 +1340,10 @@ class DeepseekV2ForCausalLM(
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
+        self.is_fp4_ckpt = (
+            self.quant_config is not None
+            and self.quant_config.get_name() == "modelopt_fp4"
+        )
 
         qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
         qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
@@ -1439,12 +1469,13 @@ class DeepseekV2ForCausalLM(
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
         ]
-        # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
-        indexer_fused_mapping = [
-            ("wk_weights_proj", "wk", 0),
-            ("wk_weights_proj", "weights_proj", 1),
-        ]
-        stacked_params_mapping.extend(indexer_fused_mapping)
+        if self.is_fp4_ckpt:
+            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+            indexer_fused_mapping = [
+                ("wk_weights_proj", "wk", 0),
+                ("wk_weights_proj", "weights_proj", 1),
+            ]
+            stacked_params_mapping.extend(indexer_fused_mapping)
 
         if self.use_mha:
             stacked_params_mapping.extend(mha_params_mapping)

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_5.py

@@ -620,6 +620,7 @@ class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid)
             self.packed_modules_mapping = {k: list(v) for k, v in base.items()}
             self.packed_modules_mapping.pop("in_proj_qkvz", None)
             self.packed_modules_mapping["in_proj_qkv"] = ["in_proj_qkv"]
+            self.packed_modules_mapping["in_proj_z"] = ["in_proj_z"]
 
     def embed_input_ids(
         self,

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

@@ -381,7 +381,12 @@ class CpuPlatform(Platform):
     def get_global_cpu_mask(cls) -> set[int]:
         # get global cpu mask
         if cls.global_cpu_mask is None:
-            cls.global_cpu_mask = os.sched_getaffinity(0)
+            if hasattr(os, "sched_getaffinity"):
+                cls.global_cpu_mask = os.sched_getaffinity(0)
+            else:
+                # macOS does not support sched_getaffinity
+                cpu_count = os.cpu_count() or 1
+                cls.global_cpu_mask = set(range(cpu_count))
         return cls.global_cpu_mask
 
     @classmethod

vllm/platforms/interface.py

@@ -207,6 +207,15 @@ class Platform:
         """
         return cls.simple_compile_backend
 
+    @classmethod
+    def import_ir_kernels(cls) -> None:
+        """
+        The default implementation imports ``vllm.kernels``, which registers
+        the built-in IR op implementations. Out-of-tree (OOT) platforms should
+        override this method to import their own kernel modules.
+        """
+        import vllm.kernels  # noqa: F401
+
     @classmethod
     def device_id_to_physical_device_id(cls, device_id: int):
         # Treat empty device control env var as unset. This is a valid

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
 
@@ -402,7 +407,6 @@ class RocmPlatform(Platform):
         "gguf",
         "quark",
         "mxfp4",
-        "petit_nvfp4",
         "torchao",
         "bitsandbytes",
     ]

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/utils/ompmultiprocessing.py

@@ -7,6 +7,7 @@ Copyright (c) 2026 Cambridge Greys Ltd
 
 import json
 import os
+import platform
 import subprocess
 
 
@@ -38,10 +39,37 @@ def parse_mask(mask):
     return set(result)
 
 
+def _get_default_affinity() -> set[int]:
+    """Get the set of CPUs the process is allowed to run on."""
+    if hasattr(os, "sched_getaffinity"):
+        return os.sched_getaffinity(0)
+    # macOS does not support sched_getaffinity; fall back to cpu_count
+    cpu_count = os.cpu_count() or 1
+    return set(range(cpu_count))
+
+
+def _get_cpu_topology_json() -> bytes:
+    """Get CPU topology as JSON.
+
+    On Linux this uses ``lscpu -Je``.  On other platforms (e.g. macOS) we
+    synthesize a simple topology where every logical CPU is its own core
+    on NUMA node 0, which is sufficient for the OMP place-list builder.
+    """
+    if platform.system() == "Linux":
+        return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout
+
+    # Fallback for non-Linux (macOS, etc.)
+    cpu_count = os.cpu_count() or 1
+    cpus = []
+    for i in range(cpu_count):
+        cpus.append({"cpu": str(i), "core": str(i), "node": "0"})
+    return json.dumps({"cpus": cpus}).encode()
+
+
 def enumerate_resources(resource_map, mask=None, allowed=None):
     """Enumerate system resources"""
     if allowed is None:
-        allowed = os.sched_getaffinity(0)
+        allowed = _get_default_affinity()
     if mask is not None:
         allowed = allowed & mask
 
@@ -140,9 +168,7 @@ class OMPProcessManager:
             else:
                 masks = [None]
             if mock is None:
-                data = subprocess.run(
-                    ["lscpu", "-Je"], check=True, capture_output=True
-                ).stdout
+                data = _get_cpu_topology_json()
             else:
                 with open(mock, mode="rb") as jf:
                     data = jf.read()

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__)
 
@@ -4192,6 +4192,7 @@ class GPUModelRunner(
         spec_config = self.speculative_config
         propose_drafts_after_bookkeeping = False
         if spec_config is not None:
+            # Decide whether to run the drafter or zero out draft tokens.
             input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
                 spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
                 <= self.effective_drafter_max_model_len
@@ -4227,10 +4228,6 @@ class GPUModelRunner(
                     self._copy_valid_sampled_token_count(
                         next_token_ids, valid_sampled_tokens_count
                     )
-                    self._draft_token_ids = torch.zeros(
-                        1, device=self.device, dtype=torch.int32
-                    ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
-                    self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
             elif (
                 spec_config.use_ngram_gpu()
                 and not spec_config.disable_padded_drafter_batch
@@ -4253,15 +4250,20 @@ class GPUModelRunner(
                     self._copy_valid_sampled_token_count(
                         next_token_ids, valid_sampled_tokens_count
                     )
-                    # Since we couldn't run the drafter,
-                    # just use zeros for the draft tokens.
-                    self._draft_token_ids = torch.zeros(
-                        1, device=self.device, dtype=torch.int32
-                    ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
-                    self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
             else:
                 propose_drafts_after_bookkeeping = input_fits_in_drafter
 
+            if not input_fits_in_drafter:
+                # Zero out draft tokens so the scheduler doesn't schedule
+                # stale drafts from the previous step.
+                # For Nemotron-H: it is necessary to zero out the draft tokens,
+                # otherwise the stale tokens will corrupt Mamba recurrent
+                # state and logprobs for sequences near max_model_len.
+                self._draft_token_ids = torch.zeros(
+                    1, device=self.device, dtype=torch.int32
+                ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
+                self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
+
         with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
             (
                 num_nans_in_logits,
@@ -5986,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,
             )
 
@@ -6844,7 +6846,7 @@ class GPUModelRunner(
             # group
             self.drafter.validate_same_kv_cache_group(kv_cache_config)
 
-        if has_kv_transfer_group():
+        if has_kv_transfer_group() and not is_profiling:
             kv_transfer_group = get_kv_transfer_group()
             if self.cross_layers_kv_cache is not None:
                 assert self.cross_layers_attn_backend is not None

vllm/v1/worker/workspace.py

@@ -31,7 +31,7 @@ _manager: "WorkspaceManager | None" = None
 class WorkspaceManager:
     """Manager for workspace allocation.
 
-    Manages workspace buffers for DBO (Dual Batch Overlap) execution.
+    Manages one workspace buffer per active ubatch slot.
     Can be locked to prevent further growth during execution.
     """
 
@@ -39,7 +39,9 @@ class WorkspaceManager:
         self._device = device
         # Cache num ubatches at init based on configuration (default to 1)
         self._num_ubatches = num_ubatches if num_ubatches is not None else 1
-        self._current_workspaces: list[torch.Tensor | None] = [None, None]
+        self._current_workspaces: list[torch.Tensor | None] = [
+            None
+        ] * self._num_ubatches
         self._locked: bool = False
 
     @staticmethod
@@ -224,7 +226,7 @@ def init_workspace_manager(
 
     Args:
         device: The device to allocate workspace on.
-        num_ubatches: Number of micro-batches. Defaults to 1.
+        num_ubatches: Number of workspace ubatch slots. Defaults to 1.
     """
     global _manager
     if _manager is not None: