Merge branch 'main' into codex/rocm-artifact-tensorizer

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

.buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml

@@ -0,0 +1,68 @@
+group: Intel
+steps:
+  - label: ":docker: Build XPU image"
+    soft_fail: true
+    optional: true
+    depends_on: []
+    key: image-build-xpu
+    commands:
+      - bash -lc '.buildkite/image_build/image_build_xpu.sh "public.ecr.aws/q9t5s3a7" "vllm-ci-test-repo" "$BUILDKITE_COMMIT"'
+    env:
+      DOCKER_BUILDKIT: "1"
+    retry:
+      automatic:
+        - exit_status: -1  # Agent was lost
+          limit: 2
+        - exit_status: -10  # Agent was lost
+          limit: 2
+  - label: "XPU example Test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    optional: true
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+      - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh example'
+  - label: "XPU V1 test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    optional: true
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+      - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh v1'
+  - label: "XPU server test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    optional: true
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - .buildkite/hardware_tests/intel_xpu_ci/test-intel.yaml
+      - .buildkite/scripts/hardware_ci/run-intel-ci-test.sh
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'bash .buildkite/scripts/hardware_ci/run-intel-ci-test.sh server'

.buildkite/intel_jobs/misc_intel.yaml

@@ -57,13 +57,16 @@ steps:
   commands:
     - >-
       bash .buildkite/scripts/hardware_ci/run-intel-test.sh
-      'export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
+      'pip install lm_eval[api]>=0.4.12 &&
+      export VLLM_WORKER_MULTIPROC_METHOD=spawn &&
       cd tests &&
       pytest -v -s v1/logits_processors --ignore=v1/logits_processors/test_custom_online.py --ignore=v1/logits_processors/test_custom_offline.py &&
       pytest -v -s v1/test_oracle.py &&
       pytest -v -s v1/test_request.py &&
       pytest -v -s v1/test_outputs.py &&
-      pytest -v -s v1/sample/test_topk_topp_sampler.py'
+      pytest -v -s v1/sample/test_topk_topp_sampler.py &&
+      pytest -v -s v1/sample/test_logprobs.py &&
+      pytest -v -s v1/sample/test_logprobs_e2e.py'
 
 - label: XPU CPU Offload
   timeout_in_minutes: 60

.buildkite/intel_jobs/model_runner_v2_intel.yaml

@@ -0,0 +1,54 @@
+group: Model Runner V2 Intel
+depends_on:
+  - image-build-xpu
+steps:
+- label: Model Runner V2 Core Tests (Intel)
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/v1/worker/gpu/
+    - vllm/v1/worker/gpu_worker.py
+    - vllm/v1/core/sched/
+    - vllm/v1/attention/
+    - tests/v1/engine/test_llm_engine.py
+    - tests/v1/e2e/
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
+      cd tests &&
+      pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics" &&
+      ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram" &&
+      pytest -v -s v1/e2e/general/test_min_tokens.py'
+
+- label: Model Runner V2 Examples (Intel)
+  timeout_in_minutes: 45
+  device: intel_gpu
+  no_plugin: true
+  working_dir: "."
+  env:
+    REGISTRY: "public.ecr.aws/q9t5s3a7"
+    REPO: "vllm-ci-test-repo"
+    VLLM_TEST_DEVICE: "xpu"
+  source_file_dependencies:
+    - vllm/v1/worker/gpu/
+    - vllm/v1/core/sched/
+    - vllm/v1/worker/gpu_worker.py
+    - examples/basic/offline_inference/
+    - examples/generate/multimodal/
+    - examples/features/
+  commands:
+    - >-
+      bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+      'export VLLM_USE_V2_MODEL_RUNNER=1 &&
+      cd examples &&
+      python3 basic/offline_inference/chat.py &&
+      python3 basic/offline_inference/generate.py --model facebook/opt-125m &&
+      python3 generate/multimodal/vision_language_offline.py --seed 0 &&
+      python3 features/automatic_prefix_caching/prefix_caching_offline.py'

.buildkite/intel_jobs/test-intel.yaml

@@ -60,6 +60,7 @@ steps:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
+        bash v1/kv_connector/nixl_integration/run_xpu_disagg_accuracy_test.sh &&
         pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py &&
         pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py &&
         pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp" &&
@@ -87,3 +88,20 @@ steps:
         cd tests &&
         pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py &&
         pytest -v -s benchmarks/test_serve_cli.py'
+  - label: "XPU quantization test"
+    depends_on:
+      - image-build-xpu
+    timeout_in_minutes: 30
+    device: intel_gpu
+    no_plugin: true
+    env:
+      REGISTRY: "public.ecr.aws/q9t5s3a7"
+      REPO: "vllm-ci-test-repo"
+    source_file_dependencies:
+      - vllm/
+      - .buildkite/intel_jobs/test-intel.yaml 
+    commands:
+      - >-
+        bash .buildkite/scripts/hardware_ci/run-intel-test.sh
+        'cd tests &&
+        pytest -v -s quantization/test_auto_round.py'

.buildkite/scripts/hardware_ci/run-intel-ci-test.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+
+set -euo pipefail
+
+test_suite="${1:-}"
+
+if [[ -z "${test_suite}" ]]; then
+  echo "Usage: $0 <example|v1|server>" >&2
+  exit 1
+fi
+
+case "${test_suite}" in
+  example)
+    pip install tblib==3.1.0
+
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 -O3 -cc.cudagraph_mode=NONE
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --quantization fp8
+    python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --kv-cache-dtype fp8
+    python3 examples/basic/offline_inference/generate.py --model nvidia/Llama-3.1-8B-Instruct-FP8 --block-size 64 --enforce-eager --quantization modelopt --kv-cache-dtype fp8 --attention-backend TRITON_ATTN --max-model-len 4096
+    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --block-size 64 --enforce-eager --max-model-len 8192
+    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
+    python3 examples/basic/offline_inference/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
+    python3 examples/basic/offline_inference/generate.py --model superjob/Qwen3-4B-Instruct-2507-GPTQ-Int4 --max-model-len 8192
+    ;;
+  v1)
+    cd tests
+
+    pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py --ignore=v1/core/test_scheduler_e2e.py
+    pytest -v -s v1/engine --ignore=v1/engine/test_output_processor.py
+    pytest -v -s v1/sample --ignore=v1/sample/test_logprobs.py --ignore=v1/sample/test_logprobs_e2e.py -k "not test_topk_only and not test_topp_only and not test_topk_and_topp"
+    pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
+    pytest -v -s v1/structured_output
+    pytest -v -s v1/test_serial_utils.py
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
+    pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
+    ;;
+  server)
+    pip install av
+    cd tests
+
+    pytest -v -s entrypoints/multimodal/openai/chat_completion/test_audio_in_video.py
+    pytest -v -s benchmarks/test_serve_cli.py
+    ;;
+  *)
+    echo "Unknown Intel test suite: ${test_suite}" >&2
+    exit 1
+    ;;
+esac

.buildkite/scripts/hardware_ci/run-intel-test.sh

@@ -243,8 +243,10 @@ container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head
 
 # ---- Command source selection ----
 commands=""
+commands_source=""
 if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
   commands="${VLLM_TEST_COMMANDS}"
+  commands_source="env"
   echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
 elif [[ $# -gt 0 ]]; then
   all_yaml=true
@@ -303,8 +305,12 @@ if [[ -z "$commands" ]]; then
 fi
 
 echo "Raw commands: $commands"
-commands=$(re_quote_pytest_markers "$commands")
-echo "After re-quoting: $commands"
+if [[ "$commands_source" != "env" ]]; then
+  commands=$(re_quote_pytest_markers "$commands")
+  echo "After re-quoting: $commands"
+else
+  echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
+fi
 commands=$(apply_intel_test_overrides "$commands")
 echo "Final commands: $commands"
 

.buildkite/test-amd.yaml

@@ -415,22 +415,6 @@ steps:
   commands:
   - pytest -v -s kernels/mamba
 
-#-----------------------------------------------------------  mi250 · lora  ------------------------------------------------------------#
-
-- label: LoRA %N # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  parallelism: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/lora
-  - tests/lora
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
-
 #------------------------------------------------------  mi250 · models / basic  -------------------------------------------------------#
 
 - label: Basic Models Test (Other CPU) # TBD
@@ -608,6 +592,11 @@ steps:
   - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
   - pip uninstall bge_m3_sparse_plugin -y
   # END: `bge_m3_sparse io_processor` test
+  # BEGIN: `colbert_query io_processor` test
+  - pip install -e ./plugins/colbert_query_plugin
+  - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
+  - pip uninstall colbert_query_plugin -y
+  # END: `colbert_query io_processor` test
   # BEGIN: `stat_logger` plugins test
   - pip install -e ./plugins/vllm_add_dummy_stat_logger
   - pytest -v -s plugins_tests/test_stats_logger_plugins.py
@@ -1694,6 +1683,20 @@ steps:
 
 #-----------------------------------------------------------  mi300 · lora  ------------------------------------------------------------#
 
+- label: LoRA %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  parallelism: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/lora
+  - tests/lora
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
+
 - label: LoRA TP (Distributed) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1778,10 +1781,9 @@ steps:
   - tests/models/multimodal/generation
   - tests/models/multimodal/test_mapping.py
   commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
-  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-  - pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
+  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
+  - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py
+  - pytest -v -s models/multimodal/test_mapping.py
 
 - label: Multi-Modal Models (Extended Generation 2) # TBD
   timeout_in_minutes: 180
@@ -1793,9 +1795,8 @@ steps:
   - vllm/
   - tests/models/multimodal/generation
   commands:
-  - uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@rocm-7.0-v2.3.0'
-  - uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
-  - pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
+  - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
+  - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
 
 
 - label: Multi-Modal Models (Extended Generation 3) # TBD
@@ -2754,7 +2755,7 @@ steps:
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   commands:
-  - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
+  - python3 benchmarks/attention_benchmarks/benchmark.py --backends ROCM_ATTN ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN --batch-specs "8q1s1k"
 
 #--------------------------------------------------------  mi355 · distributed  --------------------------------------------------------#
 

.buildkite/test_areas/benchmarks.yaml

@@ -23,4 +23,4 @@ steps:
   - benchmarks/attention_benchmarks/
   - vllm/v1/attention/
   commands:
-  - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k" --repeats 1 --warmup-iters 1
+  - python3 benchmarks/attention_benchmarks/benchmark.py --backends flash flashinfer --batch-specs "8q1s1k"

.buildkite/test_areas/lora.yaml

@@ -12,6 +12,17 @@ steps:
   commands:
     - pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py 
   parallelism: 4
+  mirror:
+    amd:
+      device: mi325_1
+      working_dir: "/vllm-workspace/tests"
+      timeout_in_minutes: 60
+      source_file_dependencies:
+      - vllm/lora
+      - tests/lora
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd
 
 
 - label: LoRA TP (Distributed)

.buildkite/test_areas/misc.yaml

@@ -21,6 +21,12 @@ steps:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     # TODO: create another `optional` test group for slow tests
     - pytest -v -s -m 'not slow_test' v1/spec_decode
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 65
+      depends_on:
+      - image-build-amd
 
 - label: V1 Sample + Logits
   key: v1-sample-logits

.buildkite/test_areas/models_multimodal.yaml

@@ -30,7 +30,6 @@ steps:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
     - pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
     - pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
-    - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
   mirror:
     amd:
       device: mi325_1
@@ -63,9 +62,16 @@ steps:
   - tests/models/multimodal
   commands:
     - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
-    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
+    - pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py  --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
+    - pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
     - pytest models/multimodal/generation/test_memory_leak.py -m core_model
     - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
+  mirror:
+    amd:
+      soft_fail: true
+      device: mi325_1
+      depends_on:
+      - image-build-amd
 
 - label: Multi-Modal Processor (CPU)
   key: multi-modal-processor-cpu

.buildkite/test_areas/plugins.yaml

@@ -27,6 +27,10 @@ steps:
   - pip install -e ./plugins/bge_m3_sparse_plugin
   - pytest -v -s plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
   - pip uninstall bge_m3_sparse_plugin -y
+  # test colbert_query io_processor plugin
+  - pip install -e ./plugins/colbert_query_plugin
+  - pytest -v -s plugins_tests/test_colbert_query_io_processor_plugins.py
+  - pip uninstall colbert_query_plugin -y
   # end io_processor plugins test
   # begin stat_logger plugins test
   - pip install -e ./plugins/vllm_add_dummy_stat_logger

.buildkite/test_areas/rust_frontend.yaml

@@ -99,9 +99,13 @@ steps:
   - vllm/v1/engine/
   - vllm/v1/worker/
   - tests/utils.py
+  - tests/v1/distributed/test_external_lb_dp.py
+  - tests/v1/distributed/test_hybrid_lb_dp.py
   - tests/v1/distributed/test_internal_lb_dp.py
   commands:
   - export VLLM_USE_RUST_FRONTEND=1
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - export NCCL_CUMEM_HOST_ENABLE=0
   - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
+  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"

.github/CODEOWNERS

@@ -120,16 +120,6 @@
 /vllm/model_executor/models/transformers @hmellor
 /tests/models/test_transformers.py @hmellor
 
-# Observability
-/vllm/config/observability.py @markmc
-/vllm/v1/metrics @markmc
-/tests/v1/metrics @markmc
-/vllm/tracing.py @markmc
-/tests/v1/tracing/test_tracing.py @markmc
-/vllm/config/kv_events.py @markmc
-/vllm/distributed/kv_events.py @markmc
-/tests/distributed/test_events.py @markmc
-
 # Docs
 /docs/mkdocs @hmellor
 /docs/**/*.yml @hmellor

.github/actionlint.yaml

@@ -0,0 +1,5 @@
+# Custom self-hosted runner labels (e.g. the autoscaling vllm-runners pool) so
+# actionlint doesn't flag them as unknown in `runs-on`.
+self-hosted-runner:
+  labels:
+    - vllm-runners

.github/mergify.yml

@@ -388,9 +388,13 @@ pull_request_rules:
     - or:
       - files~=^tests/tool_use/
       - files~=^tests/tool_parsers/
+      - files~=^tests/parser/
+      - files~=^tests/reasoning/
       - files~=^tests/entrypoints/openai/.*tool.*
       - files~=^tests/entrypoints/anthropic/.*tool.*
       - files~=^vllm/tool_parsers/
+      - files~=^vllm/parser/
+      - files~=^vllm/reasoning/
       - files=docs/features/tool_calling.md
       - files~=^examples/tool_calling/
   actions:

.github/workflows/pre-commit.yml

@@ -46,12 +46,16 @@ jobs:
   pre-commit:
     needs: pre-run-check
     if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
-    runs-on: ubuntu-latest
+    runs-on: [self-hosted, linux, x64, vllm-runners]
     steps:
     - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
     - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548 # v6.1.0
       with:
         python-version: "3.12"
+    # Provide shellcheck on PATH so tools/pre_commit/shellcheck.sh skips its
+    # wget + tar -xJ self-download, which the self-hosted runner image lacks
+    # (no wget/xz). Pinned to shellcheck 0.10.0 to match the script's "stable".
+    - run: python -m pip install shellcheck-py==0.10.0.1
     - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
     - run: echo "::add-matcher::.github/workflows/matchers/markdownlint.json"
     - run: echo "::add-matcher::.github/workflows/matchers/mypy.json"

.gitignore

@@ -15,6 +15,9 @@ vllm/third_party/flashmla/flash_mla_interface.py
 # DeepGEMM vendored package built from source
 vllm/third_party/deep_gemm/
 
+# fmha_sm100 vendored package built from source
+vllm/third_party/fmha_sm100/
+
 # triton jit
 .triton
 

CMakeLists.txt

@@ -363,35 +363,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     SRCS "${VLLM_EXT_SRC}"
     CUDA_ARCHS "${CUDA_ARCHS}")
 
-  # Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
-  endif()
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND ES_MXFP8_GROUPED_MM_ARCHS)
-    set(ES_MXFP8_GROUPED_MM_SRCS
-      "csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu"
-      "csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${ES_MXFP8_GROUPED_MM_SRCS}"
-      CUDA_ARCHS "${ES_MXFP8_GROUPED_MM_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${ES_MXFP8_GROUPED_MM_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_ES_MXFP8_GROUPED_MM_SM100=1")
-    message(STATUS "Building ES MXFP8 grouped kernels for archs: ${ES_MXFP8_GROUPED_MM_ARCHS}")
-  else()
-    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8
-        AND ES_MXFP8_GROUPED_MM_ARCHS)
-      message(STATUS "Not building ES MXFP8 grouped kernels as CUDA Compiler version is "
-                     "not >= 12.8.")
-    else()
-      message(STATUS "Not building ES MXFP8 grouped kernels as no compatible archs found "
-                     "in CUDA target architectures.")
-    endif()
-  endif()
-
-
-
 # if CUDA endif
 endif()
 
@@ -440,6 +411,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"
     "csrc/libtorch_stable/pos_encoding_kernels.cu"
     "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
+    "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
     "csrc/libtorch_stable/layernorm_kernels.cu"
     "csrc/libtorch_stable/layernorm_quant_kernels.cu"
     "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
@@ -1398,6 +1370,7 @@ endif()
 # For CUDA we also build and ship some external projects.
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
+    include(cmake/external_projects/fmha_sm100.cmake)
     include(cmake/external_projects/flashmla.cmake)
     include(cmake/external_projects/qutlass.cmake)
 

benchmarks/attention_benchmarks/benchmark.py

@@ -53,6 +53,16 @@ from common import (
 from vllm.v1.worker.workspace import init_workspace_manager
 
 
+def _str2bool(v) -> bool:
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ("true", "1", "yes", "t"):
+        return True
+    if v.lower() in ("false", "0", "no", "f"):
+        return False
+    raise argparse.ArgumentTypeError(f"expected a boolean, got {v!r}")
+
+
 def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
     """Run standard attention benchmark (Flash/Triton/FlashInfer)."""
     from runner import run_attention_benchmark
@@ -485,6 +495,20 @@ def main():
         help="Prefill backends to compare (fa2, fa3, fa4). "
         "Uses the first decode backend for impl construction.",
     )
+    parser.add_argument(
+        "--fp8-output-scale",
+        type=float,
+        help="Static per-tensor scale enabling the MLA prefill FP8-output "
+        "comparison on FA4 (fused write vs standalone post-quant).",
+    )
+    parser.add_argument(
+        "--fuse-quant-op",
+        nargs="+",
+        type=_str2bool,
+        help="FP8-output write path(s) to run: false = bf16 attention + "
+        "standalone static-FP8 quant, true = FA4 writes FP8 directly. "
+        "Default: both.",
+    )
 
     # Batch specifications
     parser.add_argument(
@@ -618,6 +642,12 @@ def main():
         # Prefill backends (e.g., ["fa3", "fa4"])
         args.prefill_backends = yaml_config.get("prefill_backends", None)
 
+        # FP8 output benchmark knobs; CLI wins.
+        if args.fp8_output_scale is None:
+            args.fp8_output_scale = yaml_config.get("fp8_output_scale", None)
+        if args.fuse_quant_op is None:
+            args.fuse_quant_op = yaml_config.get("fuse_quant_op", None)
+
         # Check for special modes
         args.mode = yaml_config.get("mode", None)
 
@@ -787,8 +817,59 @@ def main():
             "skipped (timings are placeholder zeros).[/]"
         )
 
+    # FA4 fused FP8 output vs standalone post-quant, on the same fa4 kernel:
+    # the delta is the post-quant kernel the fused path removes.
+    fp8_output_scale = getattr(args, "fp8_output_scale", None)
+    if fp8_output_scale is not None:
+        decode_backend = backends[0]
+        fuse_variants = args.fuse_quant_op or [False, True]
+        label_of = {False: "post_quant", True: "fused"}
+        console.print(
+            f"[yellow]FP8 output comparison @ scale={fp8_output_scale} "
+            f"(prefill=fa4, decode impl={decode_backend})[/]"
+        )
+        fp8_results = []
+        total = len(fuse_variants) * len(args.batch_specs)
+        with tqdm(total=total, desc="FP8 output benchmarking") as pbar:
+            for spec in args.batch_specs:
+                for fuse in fuse_variants:
+                    config = BenchmarkConfig(
+                        backend=decode_backend,
+                        batch_spec=spec,
+                        num_layers=args.num_layers,
+                        head_dim=args.head_dim,
+                        num_q_heads=args.num_q_heads,
+                        num_kv_heads=args.num_kv_heads,
+                        block_size=args.block_size,
+                        device=args.device,
+                        repeats=args.repeats,
+                        warmup_iters=args.warmup_iters,
+                        profile_memory=args.profile_memory,
+                        kv_cache_dtype=args.kv_cache_dtype,
+                        use_cuda_graphs=args.cuda_graphs,
+                        prefill_backend="fa4",
+                    )
+                    result = run_benchmark(
+                        config, output_scale=fp8_output_scale, fuse_quant_op=fuse
+                    )
+                    label = label_of[fuse]
+                    labeled_config = replace(result.config, backend=label)
+                    result = replace(result, config=labeled_config)
+                    fp8_results.append(result)
+
+                    if not result.success:
+                        console.print(f"[red]Error {label} {spec}: {result.error}[/]")
+
+                    pbar.update(1)
+
+        console.print("\n[bold green]FP8 Output Results:[/]")
+        formatter = ResultsFormatter(console)
+        labels = [label_of[f] for f in fuse_variants]
+        formatter.print_table(fp8_results, labels, compare_to_fastest=True)
+        all_results = fp8_results
+
     # Handle special mode: decode_vs_prefill comparison
-    if hasattr(args, "mode") and args.mode == "decode_vs_prefill":
+    elif hasattr(args, "mode") and args.mode == "decode_vs_prefill":
         console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
         console.print(
             "[dim]For each query length, testing both decode and prefill pipelines[/]"

benchmarks/attention_benchmarks/configs/mla_fa4_fp8_output.yaml

@@ -0,0 +1,44 @@
+# MLA prefill FP8-output microbenchmark (FA4).
+# Compares the fused FP8 write against bf16 attention + a standalone static-FP8
+# quant; the delta is the post-quant kernel the fused path removes.
+# DeepSeek-Coder-V2-Lite dims; FA4 needs SM100/110.
+#
+# Usage:
+#   python benchmark.py --config configs/mla_fa4_fp8_output.yaml
+
+description: "MLA prefill FA4 fused-FP8 output vs post-quant"
+
+model:
+  name: "deepseek-v2-lite"
+  num_layers: 27
+  num_q_heads: 16
+  num_kv_heads: 1
+  head_dim: 576
+  kv_lora_rank: 512
+  qk_nope_head_dim: 128
+  qk_rope_head_dim: 64
+  v_head_dim: 128
+  block_size: 128
+
+# Pure prefill (q_len == kv_len) so every token goes through forward_mha.
+batch_specs:
+  - "q512"
+  - "q1k"
+  - "q2k"
+  - "q4k"
+  - "q8k"
+  - "2q4k"
+  - "4q4k"
+  - "8q4k"
+
+# Only used to construct the MLA impl; the pure-prefill specs skip decode.
+decode_backends:
+  - CUTLASS_MLA
+
+# Sweep the two FP8 write paths (prefill backend is fixed to fa4).
+fp8_output_scale: 0.1
+fuse_quant_op: [false, true]
+
+device: "cuda:0"
+repeats: 50
+warmup_iters: 10

benchmarks/attention_benchmarks/mla_runner.py

@@ -708,6 +708,8 @@ def _run_single_benchmark(
     device: torch.device,
     indexer=None,
     kv_cache_dtype: str | None = None,
+    output_scale: float | None = None,
+    fuse_quant_op: bool = False,
 ) -> BenchmarkResult:
     """
     Run a single benchmark iteration.
@@ -721,6 +723,11 @@ def _run_single_benchmark(
         mla_dims: MLA dimension configuration
         device: Target device
         indexer: Optional MockIndexer for sparse backends
+        output_scale: Static per-tensor FP8 scale for prefill output. None
+            keeps the plain bf16 output (no quantization).
+        fuse_quant_op: With output_scale set, True lets the prefill kernel write
+            FP8 directly; False runs bf16 attention then a standalone static-FP8
+            quant. The delta isolates the saved post-quant kernel.
 
     Returns:
         BenchmarkResult with timing statistics
@@ -824,23 +831,55 @@ def _run_single_benchmark(
             num_prefill, mla_dims, query_fmt, device, torch.bfloat16
         )
 
+    # Prefill FP8 output: fused (kernel writes e4m3) vs separate post-quant.
+    prefill_fp8_output = None
+    prefill_output_scale = None
+    prefill_quant_op = None
+    if has_prefill and output_scale is not None:
+        from vllm.platforms import current_platform
+
+        prefill_output_scale = torch.tensor(
+            [output_scale], device=device, dtype=torch.float32
+        )
+        if fuse_quant_op:
+            prefill_fp8_output = torch.empty_like(
+                prefill_inputs["output"], dtype=current_platform.fp8_dtype()
+            )
+        else:
+            from vllm.model_executor.layers.quantization.input_quant_fp8 import (
+                QuantFP8,
+            )
+            from vllm.model_executor.layers.quantization.utils.quant_utils import (
+                GroupShape,
+            )
+
+            prefill_quant_op = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
+
+    fused_output = output_scale is not None and fuse_quant_op
+
     # Build forward function (runs a single decode/prefill pass)
     def forward_fn():
         results = []
         if has_decode:
             results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer))
         if has_prefill:
-            results.append(
-                impl.forward_mha(
-                    prefill_inputs["q"],
-                    prefill_inputs["k_c_normed"],
-                    prefill_inputs["k_pe"],
-                    kv_cache,
-                    metadata,
-                    prefill_inputs["k_scale"],
-                    prefill_inputs["output"],
-                )
+            out = impl.forward_mha(
+                prefill_inputs["q"],
+                prefill_inputs["k_c_normed"],
+                prefill_inputs["k_pe"],
+                kv_cache,
+                metadata,
+                prefill_inputs["k_scale"],
+                prefill_fp8_output if fused_output else prefill_inputs["output"],
+                prefill_output_scale if fused_output else None,
             )
+            if fused_output:
+                out = prefill_fp8_output
+            elif prefill_quant_op is not None:
+                out, _ = prefill_quant_op(
+                    prefill_inputs["output"], prefill_output_scale
+                )
+            results.append(out)
         return results[0] if len(results) == 1 else tuple(results)
 
     def benchmark_fn():
@@ -881,6 +920,8 @@ def _run_mla_benchmark_batched(
     configs_with_params: list[tuple],  # [(config, threshold, num_splits), ...]
     index_topk: int = 2048,
     prefill_backend: str | None = None,
+    output_scale: float | None = None,
+    fuse_quant_op: bool = False,
 ) -> list[BenchmarkResult]:
     """
     Unified batched MLA benchmark runner for all backends.
@@ -1020,6 +1061,8 @@ def _run_mla_benchmark_batched(
                     device,
                     indexer=indexer,
                     kv_cache_dtype=kv_cache_dtype,
+                    output_scale=output_scale,
+                    fuse_quant_op=fuse_quant_op,
                 )
                 results.append(result)
 
@@ -1047,6 +1090,8 @@ def run_mla_benchmark(
     num_kv_splits: int | None = None,
     index_topk: int = 2048,
     prefill_backend: str | None = None,
+    output_scale: float | None = None,
+    fuse_quant_op: bool = False,
 ) -> BenchmarkResult | list[BenchmarkResult]:
     """
     Unified MLA benchmark runner for all backends.
@@ -1066,6 +1111,9 @@ def run_mla_benchmark(
         index_topk: Topk value for sparse MLA backends (default 2048)
         prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
             When set, forces the specified FlashAttention version for prefill.
+        output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).
+        fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill
+            kernel vs a standalone post-quant kernel. See _run_single_benchmark.
 
     Returns:
         BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
@@ -1090,7 +1138,12 @@ def run_mla_benchmark(
 
     # Use unified batched execution
     results = _run_mla_benchmark_batched(
-        backend, configs_with_params, index_topk, prefill_backend=prefill_backend
+        backend,
+        configs_with_params,
+        index_topk,
+        prefill_backend=prefill_backend,
+        output_scale=output_scale,
+        fuse_quant_op=fuse_quant_op,
     )
 
     # Return single result or list based on input

benchmarks/benchmark_pin_memory.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark and regression-test pinned (page-locked) CPU memory for vLLM.
+
+Verifies that enabling pinned memory does not regress throughput or latency
+compared to unpinned memory.  Each condition runs in an isolated ``spawn``
+subprocess so both start from a cold CUDA context, giving an unbiased
+comparison.
+
+Usage
+-----
+Run all tests with the default model::
+
+    python benchmarks/benchmark_pin_memory.py -v
+
+Override the model and optional max-model-len::
+
+    python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B -v
+    python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B \
+        --max-model-len 8192 -v
+
+Run only throughput or latency tests::
+
+    python benchmarks/benchmark_pin_memory.py -v -k test_throughput
+    python benchmarks/benchmark_pin_memory.py -v -k test_latency
+
+Run only the v1 or v2 runner variant::
+
+    python benchmarks/benchmark_pin_memory.py -v -k v1
+    python benchmarks/benchmark_pin_memory.py -v -k v2
+
+Note: on WSL2, v1 runner tests are skipped because pin memory is not available
+for the v1 runner without cpu_offload_gb.  Run on other platforms to exercise v1.
+"""
+
+import argparse
+import json
+import multiprocessing
+import sys
+import tempfile
+
+import pytest
+
+# Allow up to 2% degradation.  Both benchmark runs start from an identical
+# cold CUDA context (separate spawn subprocesses), so the measured difference
+# reflects the genuine pin_memory overhead rather than cold/warm ordering bias.
+_THROUGHPUT_TOLERANCE = 0.98
+_THROUGHPUT_NUM_REQUESTS = 200
+_THROUGHPUT_INPUT_LEN = 128
+_THROUGHPUT_OUTPUT_LEN = 512
+_THROUGHPUT_MAX_NUM_SEQS = 128
+
+# Latency benchmark constants — match latency.py defaults.
+_LATENCY_TOLERANCE = 1.02  # Allow up to 2% latency regression.
+_LATENCY_BATCH_SIZE = 64
+_LATENCY_INPUT_LEN = 32
+_LATENCY_OUTPUT_LEN = 128
+_LATENCY_WARMUP_ITERS = 5
+_LATENCY_BENCH_ITERS = 15
+
+_DEFAULT_MODEL = "unsloth/Qwen3-1.7B"
+_DEFAULT_MAX_MODEL_LEN = 16384
+
+
+def _benchmark_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--model", default=_DEFAULT_MODEL)
+    parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
+    args, _ = parser.parse_known_args()
+    return args
+
+
+@pytest.fixture
+def model() -> str:
+    return _benchmark_args().model
+
+
+@pytest.fixture
+def max_model_len() -> int:
+    return _benchmark_args().max_model_len
+
+
+def _skip_if_pin_memory_not_available(engine_args_kwargs: dict) -> None:
+    """Skip the current pytest test if pin_memory is unavailable for this config."""
+    import vllm.utils.platform_utils as pu
+    from vllm.config import set_current_vllm_config
+    from vllm.engine.arg_utils import EngineArgs
+
+    vllm_config = EngineArgs(**engine_args_kwargs).create_engine_config()
+    with set_current_vllm_config(vllm_config):
+        pu.is_pin_memory_available.cache_clear()
+        if not pu.is_pin_memory_available():
+            import os
+
+            runner = "v2" if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") == "1" else "v1"
+            model = engine_args_kwargs.get("model", "unknown")
+            print(
+                f"\033[33mSKIP: pin_memory not available for "
+                f"{runner} runner, model={model}\033[0m"
+            )
+            pytest.skip("pin_memory not available for this configuration")
+
+
+def _throughput_worker(
+    pin: bool,
+    engine_args_kwargs: dict,
+    q: "multiprocessing.Queue[float]",
+    v2_mode: bool = False,
+) -> None:
+    """Run throughput benchmark in a fresh spawn subprocess.
+
+    Delegates to vllm/benchmarks/throughput.py main() using the random dataset,
+    so the methodology matches the official benchmark.  Results are written to a
+    temp JSON file and forwarded through the queue as tokens/s.
+
+    v2_mode: when True, monkeypatches is_uva_available() to always return True
+    so the v2 model runner's UVA buffers remain functional even when pin=False.
+    This isolates the non-UVA pin_memory paths in v2.
+    """
+    import vllm.utils.platform_utils as pu
+    from vllm.platforms import current_platform
+
+    pu.is_pin_memory_available.cache_clear()
+    pu.is_uva_available.cache_clear()
+    type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
+    if v2_mode:
+        pu.is_uva_available = lambda: True
+
+    from vllm.benchmarks.throughput import add_cli_args
+    from vllm.benchmarks.throughput import main as throughput_main
+
+    parser = argparse.ArgumentParser()
+    add_cli_args(parser)
+    args = parser.parse_args([])
+
+    for key, val in engine_args_kwargs.items():
+        setattr(args, key, val)
+    args.max_num_seqs = _THROUGHPUT_MAX_NUM_SEQS
+    args.dataset_name = "random"
+    args.input_len = _THROUGHPUT_INPUT_LEN
+    args.output_len = _THROUGHPUT_OUTPUT_LEN
+    # Nullify defaults that conflict with explicit input/output_len.
+    args.random_input_len = None
+    args.random_output_len = None
+    args.random_prefix_len = None
+    args.num_prompts = _THROUGHPUT_NUM_REQUESTS
+    args.seed = 0
+    args.disable_detokenize = True
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        tmp_path = f.name
+    args.output_json = tmp_path
+
+    throughput_main(args)
+
+    with open(tmp_path) as f:
+        results = json.load(f)
+    q.put(results["tokens_per_second"])
+
+
+def _run_throughput_benchmark(
+    pin: bool,
+    engine_args_kwargs: dict,
+    v2_mode: bool = False,
+) -> float:
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    p = ctx.Process(
+        target=_throughput_worker,
+        args=(pin, engine_args_kwargs, q, v2_mode),
+    )
+    p.start()
+    p.join()
+    if p.exitcode != 0:
+        raise RuntimeError(
+            f"Throughput benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
+        )
+    return q.get()
+
+
+def _latency_worker(
+    pin: bool,
+    engine_args_kwargs: dict,
+    q: "multiprocessing.Queue[dict]",
+    v2_mode: bool = False,
+) -> None:
+    """Run latency benchmark in a fresh spawn subprocess.
+
+    Follows latency.py methodology: fixed batch of dummy token IDs, warmup
+    iterations to reach steady state, then timed iterations reduced to avg
+    and percentiles.  Results are written to a temp JSON file by latency_main
+    and forwarded through the queue.
+    """
+    import vllm.utils.platform_utils as pu
+    from vllm.platforms import current_platform
+
+    pu.is_pin_memory_available.cache_clear()
+    pu.is_uva_available.cache_clear()
+    type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
+    if v2_mode:
+        pu.is_uva_available = lambda: True
+
+    from vllm.benchmarks.latency import add_cli_args
+    from vllm.benchmarks.latency import main as latency_main
+
+    parser = argparse.ArgumentParser()
+    add_cli_args(parser)
+    args = parser.parse_args([])
+
+    for key, val in engine_args_kwargs.items():
+        setattr(args, key, val)
+    args.input_len = _LATENCY_INPUT_LEN
+    args.output_len = _LATENCY_OUTPUT_LEN
+    args.batch_size = _LATENCY_BATCH_SIZE
+    args.num_iters_warmup = _LATENCY_WARMUP_ITERS
+    args.num_iters = _LATENCY_BENCH_ITERS
+    args.profile = False
+    args.disable_detokenize = True
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        tmp_path = f.name
+    args.output_json = tmp_path
+
+    latency_main(args)
+
+    with open(tmp_path) as f:
+        results = json.load(f)
+    q.put(results)
+
+
+def _run_latency_benchmark(
+    pin: bool,
+    engine_args_kwargs: dict,
+    v2_mode: bool = False,
+) -> dict:
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    p = ctx.Process(
+        target=_latency_worker,
+        args=(pin, engine_args_kwargs, q, v2_mode),
+    )
+    p.start()
+    p.join()
+    if p.exitcode != 0:
+        raise RuntimeError(
+            f"Latency benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
+        )
+    return q.get()
+
+
+@pytest.mark.parametrize(
+    "test_v2_runner",
+    [
+        pytest.param(False, id="v1"),
+        pytest.param(True, id="v2"),
+    ],
+)
+class TestPinnedMemory:
+    """Verify pinned memory yields >= throughput vs unpinned via real vLLM inference."""
+
+    def test_throughput(self, monkeypatch, test_v2_runner, model, max_model_len):
+        """Benchmark throughput with pin_memory forced on then off.
+
+        Delegates to vllm/benchmarks/throughput.py main() with the random
+        dataset.  Each condition runs in an isolated spawn subprocess so both
+        start from a cold CUDA context, giving an unbiased comparison.
+        """
+        monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
+
+        engine_args_kwargs = dict(
+            model=model,
+            gpu_memory_utilization=0.88,
+            max_model_len=max_model_len,
+            enable_prefix_caching=False,
+        )
+
+        _skip_if_pin_memory_not_available(engine_args_kwargs)
+
+        unpinned_tps = _run_throughput_benchmark(
+            False, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+        pinned_tps = _run_throughput_benchmark(
+            True, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+
+        pct_diff = (pinned_tps - unpinned_tps) / unpinned_tps * 100
+        runner = "v2" if test_v2_runner else "v1"
+        print(
+            f"\n=== Throughput results ({runner} runner, {model}) ==="
+            f"\npin_memory=True:  {pinned_tps:.1f} tok/s"
+            f"\npin_memory=False: {unpinned_tps:.1f} tok/s"
+            f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
+        )
+
+        assert pinned_tps >= unpinned_tps * _THROUGHPUT_TOLERANCE, (
+            f"Pinned throughput ({pinned_tps:.1f} tok/s) fell more than "
+            f"{(1.0 - _THROUGHPUT_TOLERANCE) * 100:.1f}% below "
+            f"unpinned ({unpinned_tps:.1f} tok/s)."
+        )
+
+    def test_latency(self, monkeypatch, test_v2_runner, model, max_model_len):
+        """Benchmark per-batch latency with pin_memory forced on then off.
+
+        Follows vllm/benchmarks/latency.py: fixed dummy-token batch, warmup
+        iterations to reach steady state, then timed iterations reduced to avg
+        and percentiles.  Subprocesses run serially so each gets a cold CUDA
+        context without GPU memory pressure from the other run.
+        """
+        monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
+
+        engine_args_kwargs = dict(
+            model=model,
+            gpu_memory_utilization=0.88,
+            max_model_len=max_model_len,
+            enable_prefix_caching=False,
+        )
+
+        _skip_if_pin_memory_not_available(engine_args_kwargs)
+
+        unpinned = _run_latency_benchmark(
+            False, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+        pinned = _run_latency_benchmark(
+            True, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+
+        pct_diff = (
+            (pinned["avg_latency"] - unpinned["avg_latency"])
+            / unpinned["avg_latency"]
+            * 100
+        )
+        runner = "v2" if test_v2_runner else "v1"
+        print(
+            f"\n=== Latency results ({runner} runner, {model}) ==="
+            f"\npin_memory=True:  avg={pinned['avg_latency']:.3f}s"
+            f"  p50={pinned['percentiles']['50']:.3f}s"
+            f"  p99={pinned['percentiles']['99']:.3f}s"
+            f"\npin_memory=False: avg={unpinned['avg_latency']:.3f}s"
+            f"  p50={unpinned['percentiles']['50']:.3f}s"
+            f"  p99={unpinned['percentiles']['99']:.3f}s"
+            f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
+        )
+
+        assert pinned["avg_latency"] <= unpinned["avg_latency"] * _LATENCY_TOLERANCE, (
+            f"Pinned avg latency ({pinned['avg_latency']:.3f}s) exceeded "
+            f"unpinned ({unpinned['avg_latency']:.3f}s) by more than "
+            f"{(_LATENCY_TOLERANCE - 1.0) * 100:.1f}%."
+        )
+
+
+if __name__ == "__main__":
+    _parser = argparse.ArgumentParser(add_help=False)
+    _parser.add_argument("--model", default=_DEFAULT_MODEL)
+    _parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
+    _, _remaining = _parser.parse_known_args()
+    sys.exit(pytest.main([__file__] + _remaining))

cmake/external_projects/fmha_sm100.cmake

@@ -0,0 +1,48 @@
+include(FetchContent)
+
+# If FMHA_SM100_SRC_DIR is set, fmha_sm100 is installed from that directory
+# instead of downloading. This is useful for local MSA development.
+if(DEFINED ENV{FMHA_SM100_SRC_DIR})
+  set(FMHA_SM100_SRC_DIR $ENV{FMHA_SM100_SRC_DIR})
+endif()
+
+if(FMHA_SM100_SRC_DIR)
+  FetchContent_Declare(
+    fmha_sm100
+    SOURCE_DIR ${FMHA_SM100_SRC_DIR}
+    CONFIGURE_COMMAND ""
+    BUILD_COMMAND ""
+  )
+else()
+  FetchContent_Declare(
+    fmha_sm100
+    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
+    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
+    GIT_PROGRESS TRUE
+    CONFIGURE_COMMAND ""
+    BUILD_COMMAND ""
+  )
+endif()
+
+FetchContent_GetProperties(fmha_sm100)
+if(NOT fmha_sm100_POPULATED)
+  FetchContent_Populate(fmha_sm100)
+endif()
+message(STATUS "fmha_sm100 is available at ${fmha_sm100_SOURCE_DIR}")
+
+add_custom_target(fmha_sm100)
+
+set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")
+
+install(FILES
+  "${FMHA_SM100_PY_ROOT}/__init__.py"
+  "${FMHA_SM100_PY_ROOT}/sparse.py"
+  DESTINATION vllm/third_party/fmha_sm100
+  COMPONENT fmha_sm100)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
+  DESTINATION vllm/third_party/fmha_sm100/cute
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)

csrc/cpu/cpu_attn.cpp

@@ -11,13 +11,25 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
   return cpu_attention::Fp8KVCacheDataType::kAuto;
 }
 
+bool cpu_attn_has_isa(const std::string& isa) {
+  if (isa == "rvv") {
+#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
+    return true;
+#else
+    return false;
+#endif
+  }
+  return false;
+}
+
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
     const int64_t num_heads_kv, const int64_t head_dim,
     const torch::Tensor& seq_lens, at::ScalarType dtype,
-    const torch::Tensor& query_start_loc, const bool casual,
+    const torch::Tensor& query_start_loc, const bool causal,
     const int64_t window_size, const std::string& isa_hint,
-    const bool enable_kv_split) {
+    const bool enable_kv_split,
+    const std::optional<torch::Tensor>& dynamic_causal) {
   cpu_attention::ISA isa;
   if (isa_hint == "amx") {
     isa = cpu_attention::ISA::AMX;
@@ -44,24 +56,13 @@ torch::Tensor get_scheduler_metadata(
   input.head_dim = head_dim;
   input.query_start_loc = query_start_loc.data_ptr<int32_t>();
   input.seq_lens = seq_lens.data_ptr<int32_t>();
-  if (window_size != -1) {
-    input.left_sliding_window_size = window_size - 1;
-    if (casual) {
-      input.right_sliding_window_size = 0;
-    } else {
-      input.right_sliding_window_size = window_size - 1;
-    }
-  } else {
-    input.left_sliding_window_size = -1;
-    if (casual) {
-      input.right_sliding_window_size = 0;
-    } else {
-      input.right_sliding_window_size = -1;
-    }
-  }
-  input.casual = casual;
+
+  input.sliding_window_size = window_size;
+  input.causal = causal;
   input.isa = isa;
   input.enable_kv_split = enable_kv_split;
+  input.dynamic_causal =
+      dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
 
   VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
     CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
@@ -175,10 +176,11 @@ void cpu_attention_with_kv_cache(
     const torch::Tensor& seq_lens,         // [num_tokens]
     const double scale, const bool causal,
     const std::optional<torch::Tensor>& alibi_slopes,  // [num_heads]
-    const int64_t sliding_window_left, const int64_t sliding_window_right,
+    const int64_t sliding_window,
     const torch::Tensor& block_table,  // [num_tokens, max_block_num]
     const double softcap, const torch::Tensor& scheduler_metadata,
-    const std::optional<torch::Tensor>& s_aux,  // [num_heads]
+    const std::optional<torch::Tensor>& s_aux,           // [num_heads]
+    const std::optional<torch::Tensor>& dynamic_causal,  // [num_reqs]
     const double k_scale = 1.0, const double v_scale = 1.0,
     const std::string& kv_cache_dtype = "auto") {
   TORCH_CHECK_EQ(query.dim(), 3);
@@ -220,13 +222,11 @@ void cpu_attention_with_kv_cache(
   input.alibi_slopes =
       alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
   input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
+  input.dynamic_causal =
+      dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
   input.scale = scale;
   input.causal = causal;
-  input.sliding_window_left = sliding_window_left;
-  input.sliding_window_right = sliding_window_right;
-  if (input.causal) {
-    input.sliding_window_right = 0;
-  }
+  input.sliding_window_size = sliding_window;
   input.softcap = static_cast<float>(softcap);
 
   if (is_fp8) {

csrc/cpu/cpu_attn_impl.hpp

@@ -388,13 +388,13 @@ class AttentionScheduler {
     int32_t head_dim;
     int32_t* query_start_loc;
     int32_t* seq_lens;
-    int32_t left_sliding_window_size;
-    int32_t right_sliding_window_size;
-    bool casual;
+    int32_t sliding_window_size;
+    bool causal;
     cpu_attention::ISA isa;
     int32_t max_num_q_per_iter;  // max Q head num can be hold in registers
     int32_t kv_block_alignment;  // context length alignment requirement
     bool enable_kv_split;
+    bool* dynamic_causal;
   };
 
   static constexpr int32_t MaxQTileIterNum = 128;
@@ -403,7 +403,8 @@ class AttentionScheduler {
       : available_cache_size_(cpu_utils::get_available_l2_size()) {}
 
   torch::Tensor schedule(const ScheduleInput& input) const {
-    const bool casual = input.casual;
+    const bool causal = input.causal;
+    const bool is_dynamic_causal = input.dynamic_causal != nullptr;
     const int32_t thread_num = omp_get_max_threads();
     const int64_t cache_size = cpu_utils::get_available_l2_size();
     const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
@@ -434,8 +435,7 @@ class AttentionScheduler {
     const int32_t default_tile_token_num = default_tile_size / q_head_per_kv;
     const int32_t split_kv_q_token_num_threshold =
         input.enable_kv_split ? 1 : 0;
-    const int32_t left_sliding_window_size = input.left_sliding_window_size;
-    const int32_t right_sliding_window_size = input.right_sliding_window_size;
+    const int32_t sliding_window_size = input.sliding_window_size;
     TORCH_CHECK_LE(split_kv_q_token_num_threshold * q_head_per_kv, 16);
 
     // get total kv len
@@ -444,7 +444,9 @@ class AttentionScheduler {
       const int32_t seq_len = input.seq_lens[req_id];
       const int32_t q_token_num =
           input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
-      const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
+      const bool req_causal =
+          is_dynamic_causal ? input.dynamic_causal[req_id] : causal;
+      const int32_t q_start_pos = seq_len - q_token_num;
       const int32_t kv_start_pos = 0;
       const int32_t kv_end_pos = seq_len;
 
@@ -456,7 +458,7 @@ class AttentionScheduler {
         const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
         const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
             kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
-            left_sliding_window_size, right_sliding_window_size);
+            sliding_window_size, req_causal);
         const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
             align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
                               kv_len_alignment);
@@ -484,7 +486,9 @@ class AttentionScheduler {
       const int32_t seq_len = input.seq_lens[req_id];
       const int32_t q_token_num =
           input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
-      const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
+      const bool req_causal =
+          is_dynamic_causal ? input.dynamic_causal[req_id] : causal;
+      const int32_t q_start_pos = seq_len - q_token_num;
       const int32_t kv_start_pos = 0;
       const int32_t kv_end_pos = seq_len;
       int32_t local_split_id = 0;
@@ -498,7 +502,7 @@ class AttentionScheduler {
         const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
         const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
             kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
-            left_sliding_window_size, right_sliding_window_size);
+            sliding_window_size, req_causal);
         const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
             align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
                               kv_len_alignment);
@@ -708,15 +712,41 @@ class AttentionScheduler {
     return metadata_tensor;
   }
 
+  FORCE_INLINE static std::pair<int32_t, int32_t> calcu_sliding_window_size(
+      int32_t window_size, bool causal) {
+    int32_t left_sliding_window_size, right_sliding_window_size;
+    if (window_size != -1) {
+      left_sliding_window_size = window_size - 1;
+      if (causal) {
+        right_sliding_window_size = 0;
+      } else {
+        right_sliding_window_size = window_size - 1;
+      }
+    } else {
+      left_sliding_window_size = -1;
+      if (causal) {
+        right_sliding_window_size = 0;
+      } else {
+        right_sliding_window_size = -1;
+      }
+    }
+
+    return {left_sliding_window_size, right_sliding_window_size};
+  }
+
   FORCE_INLINE static std::pair<int32_t, int32_t> calcu_kv_tile_pos(
       int32_t kv_left_pos, int32_t kv_right_pos, int32_t q_left_pos,
-      int32_t q_right_pos, int32_t sliding_window_left,
-      int32_t sliding_window_right) {
-    if (sliding_window_left != -1) {
-      kv_left_pos = std::max(kv_left_pos, q_left_pos - sliding_window_left);
+      int32_t q_right_pos, int32_t window_size, bool causal) {
+    auto [left_sliding_window_size, right_sliding_window_size] =
+        calcu_sliding_window_size(window_size, causal);
+
+    if (left_sliding_window_size != -1) {
+      kv_left_pos =
+          std::max(kv_left_pos, q_left_pos - left_sliding_window_size);
     }
-    if (sliding_window_right != -1) {
-      kv_right_pos = std::min(kv_right_pos, q_right_pos + sliding_window_right);
+    if (right_sliding_window_size != -1) {
+      kv_right_pos =
+          std::min(kv_right_pos, q_right_pos + right_sliding_window_size);
     }
     return {kv_left_pos, kv_right_pos};
   }
@@ -805,10 +835,10 @@ struct AttentionInput {
   int32_t* block_table;
   float* alibi_slopes;
   c10::BFloat16* s_aux;
+  bool* dynamic_causal;
   float scale;
   bool causal;
-  int32_t sliding_window_left;
-  int32_t sliding_window_right;
+  int32_t sliding_window_size;
   float softcap;
   // FP8 KV cache scales (used by FP8 attention implementations)
   float k_scale_fp8 = 1.0f;
@@ -1442,15 +1472,16 @@ class AttentionMainLoop {
       const int64_t q_head_num_stride = input->query_num_heads_stride;
       const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
       const int64_t kv_cache_block_num_stride = input->cache_num_blocks_stride;
-      const int32_t sliding_window_left = input->sliding_window_left;
-      const int32_t sliding_window_right = input->sliding_window_right;
+      const int32_t sliding_window_size = input->sliding_window_size;
       const int32_t block_size = input->block_size;
       const float scale = input->scale;
       const float softcap_scale = input->softcap;
       const float* alibi_slopes = input->alibi_slopes;
       const c10::BFloat16* s_aux = input->s_aux;
+      const bool* dynamic_causal = input->dynamic_causal;
+      const bool is_dynamic_causal = dynamic_causal != nullptr;
 
-      const bool casual = input->causal;
+      const bool causal = input->causal;
       int32_t* const block_table = input->block_table;
       const int64_t block_table_stride = input->blt_num_tokens_stride;
 
@@ -1533,6 +1564,11 @@ class AttentionMainLoop {
                 &curr_workitem_groups[workitem_group_idx];
 
             const int32_t current_group_idx = current_workitem_group->req_id;
+            const int32_t current_group_causal =
+                is_dynamic_causal ? dynamic_causal[current_group_idx] : causal;
+            auto [sliding_window_left, sliding_window_right] =
+                AttentionScheduler::calcu_sliding_window_size(
+                    sliding_window_size, current_group_causal);
             const int32_t kv_start_pos =
                 current_workitem_group->kv_split_pos_start;
             const int32_t kv_end_pos = current_workitem_group->kv_split_pos_end;
@@ -1560,8 +1596,7 @@ class AttentionMainLoop {
             const int32_t q_end = input->query_start_loc[current_group_idx + 1];
             const int32_t q_start = input->query_start_loc[current_group_idx];
             const int32_t seq_len = input->seq_lens[current_group_idx];
-            const int32_t q_start_pos =
-                (casual ? seq_len - (q_end - q_start) : 0);
+            const int32_t q_start_pos = seq_len - (q_end - q_start);
             const int32_t block_num = (seq_len + block_size - 1) / block_size;
             // Only apply sink for the first KV split
             bool use_sink = (s_aux != nullptr &&
@@ -1611,8 +1646,8 @@ class AttentionMainLoop {
               const auto [kv_tile_start_pos, kv_tile_end_pos] =
                   AttentionScheduler::calcu_kv_tile_pos(
                       kv_start_pos, kv_end_pos, q_tile_start_pos,
-                      q_tile_end_pos, sliding_window_left,
-                      sliding_window_right);
+                      q_tile_end_pos, sliding_window_size,
+                      current_group_causal);
               const auto [rounded_kv_tile_start_pos, rounded_kv_tile_end_pos] =
                   AttentionScheduler::align_kv_tile_pos(
                       kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
@@ -1725,8 +1760,8 @@ class AttentionMainLoop {
                               actual_kv_tile_pos_right] =
                       AttentionScheduler::calcu_kv_tile_pos(
                           kv_tile_pos_left, kv_tile_pos_right, q_tile_pos_left,
-                          q_tile_pos_right, sliding_window_left,
-                          sliding_window_right);
+                          q_tile_pos_right, sliding_window_size,
+                          current_group_causal);
                   const int32_t q_iter_idx =
                       q_head_tile_token_offset / curr_max_q_token_num_per_iter;
 

csrc/cpu/cpu_fused_moe.cpp

@@ -1,3 +1,5 @@
+#include <sleef.h>
+
 #include "cpu/cpu_types.hpp"
 #include "cpu/utils.hpp"
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
@@ -163,7 +165,6 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
   vec_op::FP32Vec16 w1_vec(0.7978845608028654);
   vec_op::FP32Vec16 w2_vec(0.5);
   vec_op::FP32Vec16 w3_vec(0.044715);
-  alignas(64) float temp[16];
 
   for (int32_t m = 0; m < m_size; ++m) {
     for (int32_t n = 0; n < dim; n += 16) {
@@ -171,12 +172,9 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
       vec_op::FP32Vec16 up_vec(up + n);
       auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
       auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
-
-      inner_vec.save(temp);
-      for (int32_t i = 0; i < 16; ++i) {
-        temp[i] = std::tanh(temp[i]);
-      }
-      vec_op::FP32Vec16 tanh_vec(temp);
+      // Note: can't use fast_exp form because diffusiongemma will generate
+      // wrong results
+      vec_op::FP32Vec16 tanh_vec(Sleef_tanhf16_u10(inner_vec.reg));
       auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
       auto gated_output_fp32 = up_vec * gelu_tanh;
       scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);

csrc/cpu/layernorm.cpp

@@ -4,8 +4,9 @@ namespace {
 template <typename scalar_t>
 void rms_norm_impl(scalar_t* __restrict__ out,
                    const scalar_t* __restrict__ input,
-                   const scalar_t* __restrict__ weight, const float epsilon,
-                   const int num_tokens, const int hidden_size) {
+                   const scalar_t* __restrict__ weight, const bool has_weight,
+                   const float epsilon, const int num_tokens,
+                   const int hidden_size) {
   using scalar_vec_t = vec_op::vec_t<scalar_t>;
   constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
   TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
@@ -27,12 +28,15 @@ void rms_norm_impl(scalar_t* __restrict__ out,
 
     for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
       scalar_vec_t x(input_p + j);
-      scalar_vec_t w(weight + j);
-
       vec_op::FP32Vec8 fp32_x(x);
-      vec_op::FP32Vec8 fp32_w(w);
-
-      vec_op::FP32Vec8 fp32_out = fp32_x * fp32_s_variance * fp32_w;
+      vec_op::FP32Vec8 fp32_out;
+      if (has_weight) {
+        scalar_vec_t w(weight + j);
+        vec_op::FP32Vec8 fp32_w(w);
+        fp32_out = fp32_x * fp32_s_variance * fp32_w;
+      } else {
+        fp32_out = fp32_x * fp32_s_variance;
+      }
 
       scalar_vec_t out(fp32_out);
       out.save(output_p + j);
@@ -44,8 +48,8 @@ template <typename scalar_t>
 void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
                              scalar_t* __restrict__ residual,
                              const scalar_t* __restrict__ weight,
-                             const float epsilon, const int num_tokens,
-                             const int hidden_size) {
+                             const bool has_weight, const float epsilon,
+                             const int num_tokens, const int hidden_size) {
   using scalar_vec_t = vec_op::vec_t<scalar_t>;
   constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
   TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
@@ -72,13 +76,18 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
     vec_op::FP32Vec8 fp32_s_variance(s_variance);
 
     for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
-      scalar_vec_t w(weight + j);
-      scalar_vec_t res(residual_p + j);
-
-      vec_op::FP32Vec8 fp32_w(w);
-      vec_op::FP32Vec8 fp32_res(res);
-
-      vec_op::FP32Vec8 fp32_out = fp32_res * fp32_s_variance * fp32_w;
+      vec_op::FP32Vec8 fp32_out;
+      if (has_weight) {
+        scalar_vec_t w(weight + j);
+        scalar_vec_t res(residual_p + j);
+        vec_op::FP32Vec8 fp32_w(w);
+        vec_op::FP32Vec8 fp32_res(res);
+        fp32_out = fp32_res * fp32_s_variance * fp32_w;
+      } else {
+        scalar_vec_t res(residual_p + j);
+        vec_op::FP32Vec8 fp32_res(res);
+        fp32_out = fp32_res * fp32_s_variance;
+      }
 
       scalar_vec_t out(fp32_out);
       out.save(input_p + j);
@@ -87,31 +96,41 @@ void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
 }
 }  // namespace
 
-void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
-              double epsilon) {
+void rms_norm(torch::Tensor& out, torch::Tensor& input,
+              std::optional<torch::Tensor> weight, double epsilon) {
   int hidden_size = input.size(-1);
   int num_tokens = input.numel() / hidden_size;
+  const bool has_weight = weight.has_value();
+  if (has_weight) {
+    TORCH_CHECK(weight->is_contiguous());
+  }
 
   VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] {
     CPU_KERNEL_GUARD_IN(rms_norm_impl)
     rms_norm_impl(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
-                  weight.data_ptr<scalar_t>(), epsilon, num_tokens,
-                  hidden_size);
+                  has_weight ? weight->data_ptr<scalar_t>() : nullptr,
+                  has_weight, epsilon, num_tokens, hidden_size);
     CPU_KERNEL_GUARD_OUT(rms_norm_impl)
   });
 }
 
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
-                        torch::Tensor& weight, double epsilon) {
+                        std::optional<torch::Tensor> weight, double epsilon) {
   int hidden_size = input.size(-1);
   int num_tokens = input.numel() / hidden_size;
+  const bool has_weight = weight.has_value();
+  if (has_weight) {
+    TORCH_CHECK(weight->scalar_type() == input.scalar_type());
+    TORCH_CHECK(weight->is_contiguous());
+  }
 
   VLLM_DISPATCH_FLOATING_TYPES(
       input.scalar_type(), "fused_add_rms_norm_impl", [&] {
         CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl)
         fused_add_rms_norm_impl(
             input.data_ptr<scalar_t>(), residual.data_ptr<scalar_t>(),
-            weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
+            has_weight ? weight->data_ptr<scalar_t>() : nullptr, has_weight,
+            epsilon, num_tokens, hidden_size);
         CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl)
       });
 }

csrc/cpu/torch_bindings.cpp

@@ -146,13 +146,16 @@ at::Tensor causal_conv1d_update_cpu(
 void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                          const std::string& activation);
 
+bool cpu_attn_has_isa(const std::string& isa);
+
 torch::Tensor get_scheduler_metadata(
     const int64_t num_req, const int64_t num_heads_q,
     const int64_t num_heads_kv, const int64_t head_dim,
     const torch::Tensor& seq_lens, at::ScalarType dtype,
     const torch::Tensor& query_start_loc, const bool casual,
     const int64_t window_size, const std::string& isa_hint,
-    const bool enable_kv_split);
+    const bool enable_kv_split,
+    const std::optional<torch::Tensor>& dynamic_causal);
 
 void cpu_attn_reshape_and_cache(const torch::Tensor& key,
                                 const torch::Tensor& value,
@@ -169,10 +172,10 @@ void cpu_attention_with_kv_cache(
     const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
     const double scale, const bool causal,
     const std::optional<torch::Tensor>& alibi_slopes,
-    const int64_t sliding_window_left, const int64_t sliding_window_right,
-    const torch::Tensor& block_table, const double softcap,
-    const torch::Tensor& scheduler_metadata,
-    const std::optional<torch::Tensor>& s_aux, const double k_scale,
+    const int64_t sliding_window_left, const torch::Tensor& block_table,
+    const double softcap, const torch::Tensor& scheduler_metadata,
+    const std::optional<torch::Tensor>& s_aux,
+    const std::optional<torch::Tensor>& dynamic_causal, const double k_scale,
     const double v_scale, const std::string& kv_cache_dtype);
 
 // Note: just for avoiding importing errors
@@ -309,13 +312,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // Layernorm
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(
-      "rms_norm(Tensor! out, Tensor input, Tensor weight, float epsilon) -> "
+      "rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> "
       "()");
   ops.impl("rms_norm", torch::kCPU, &rms_norm);
 
   // In-place fused Add and RMS Normalization.
   ops.def(
-      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
+      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
       "float epsilon) -> ()");
   ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm);
 
@@ -496,11 +499,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
 
   // CPU attention kernels
+  ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa);
   ops.def(
       "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
       "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
       "query_start_loc, bool casual, int window_size, str isa_hint, bool "
-      "enable_kv_split) -> Tensor",
+      "enable_kv_split, Tensor? dynamic_causal) -> Tensor",
       &get_scheduler_metadata);
   ops.def(
       "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
@@ -512,8 +516,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
       "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
       "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
-      "sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
-      "float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
+      "sliding_window_size, Tensor block_table, "
+      "float softcap, Tensor scheduler_metadata, Tensor? s_aux, Tensor? "
+      "dynamic_causal, "
       "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
       "()",
       &cpu_attention_with_kv_cache);

csrc/libtorch_stable/activation_kernels.cu

@@ -10,11 +10,20 @@
 
 namespace vllm {
 
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+// `alpha` and `beta` are applied to opposite operands:
+//   - alpha lives INSIDE the activation (the activated half): the gated
+//     activation computes act_half * sigmoid(alpha * act_half).
+//   - beta is added to the OTHER (non-activated) half before the multiply.
+// So the result is always ACT(act_half, alpha) * (other_half + beta).
+// Which half is which depends on `act_first` (see below). Defaults
+// alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
           bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
                                             const scalar_t& y,
-                                            const float limit) {
+                                            const float limit,
+                                            const float alpha,
+                                            const float beta) {
   if constexpr (act_first) {
     scalar_t gate = x;
     scalar_t up = y;
@@ -22,7 +31,9 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
       gate = (scalar_t)fminf((float)gate, limit);
       up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
     }
-    return ACT_FN(gate) * up;
+    // act_first: gate is the activated half -> alpha applies to gate;
+    // beta is added to up (the non-activated half).
+    return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta));
   } else {
     scalar_t gate = x;
     scalar_t up = y;
@@ -30,55 +41,68 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
       gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
       up = (scalar_t)fminf((float)up, limit);
     }
-    return gate * ACT_FN(up);
+    // !act_first: up is the activated half -> alpha applies to up;
+    // beta is added to gate (the non-activated half).
+    return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha));
   }
 }
 
-template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
+template <typename packed_t,
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
           bool act_first, bool HAS_CLAMP>
 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
                                                    const packed_t& y,
-                                                   const float limit) {
+                                                   const float limit,
+                                                   const float alpha,
+                                                   const float beta) {
   if constexpr (act_first) {
     packed_t gate = x;
     packed_t up = y;
+    float2 u = cast_to_float2(up);
     if constexpr (HAS_CLAMP) {
       float2 g = cast_to_float2(gate);
-      float2 u = cast_to_float2(up);
       g.x = fminf(g.x, limit);
       g.y = fminf(g.y, limit);
       u.x = fmaxf(fminf(u.x, limit), -limit);
       u.y = fmaxf(fminf(u.y, limit), -limit);
       gate = cast_to_packed<packed_t>(g);
-      up = cast_to_packed<packed_t>(u);
     }
-    return packed_mul(PACKED_ACT_FN(gate), up);
+    // act_first: gate is the activated half -> alpha applies to gate;
+    // beta is added to up (the non-activated half).
+    float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha));
+    activated.x *= u.x + beta;
+    activated.y *= u.y + beta;
+    return cast_to_packed<packed_t>(activated);
   } else {
     packed_t gate = x;
     packed_t up = y;
+    float2 g = cast_to_float2(gate);
     if constexpr (HAS_CLAMP) {
-      float2 g = cast_to_float2(gate);
       float2 u = cast_to_float2(up);
       g.x = fmaxf(fminf(g.x, limit), -limit);
       g.y = fmaxf(fminf(g.y, limit), -limit);
       u.x = fminf(u.x, limit);
       u.y = fminf(u.y, limit);
-      gate = cast_to_packed<packed_t>(g);
       up = cast_to_packed<packed_t>(u);
     }
-    return packed_mul(gate, PACKED_ACT_FN(up));
+    // !act_first: up is the activated half -> alpha applies to up;
+    // beta is added to gate (the non-activated half).
+    float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha));
+    activated.x *= g.x + beta;
+    activated.y *= g.y + beta;
+    return cast_to_packed<packed_t>(activated);
   }
 }
 
 // Activation and gating kernel template.
 template <typename scalar_t, typename packed_t,
-          scalar_t (*ACT_FN)(const scalar_t&),
-          packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
-          bool use_vec, bool HAS_CLAMP, bool use_256b = false>
+          scalar_t (*ACT_FN)(const scalar_t&, const float),
+          packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
+          bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
 __global__ void act_and_mul_kernel(
     scalar_t* __restrict__ out,          // [..., d]
     const scalar_t* __restrict__ input,  // [..., 2, d]
-    const int d, const float limit) {
+    const int d, const float limit, const float alpha, const float beta) {
   const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
   const scalar_t* y_ptr = x_ptr + d;
   scalar_t* out_ptr = out + blockIdx.x * d;
@@ -105,7 +129,7 @@ __global__ void act_and_mul_kernel(
       for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
         x.elts[j] =
             packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
-                x.elts[j], y.elts[j], limit);
+                x.elts[j], y.elts[j], limit, alpha, beta);
       }
       if constexpr (use_256b) {
         st256(x, &out_vec[i]);
@@ -118,29 +142,34 @@ __global__ void act_and_mul_kernel(
     for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
       const scalar_t x = VLLM_LDG(&x_ptr[idx]);
       const scalar_t y = VLLM_LDG(&y_ptr[idx]);
-      out_ptr[idx] =
-          compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
+      out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
+          x, y, limit, alpha, beta);
     }
   }
 }
 
+// Gated activations take an `alpha` argument that scales the sigmoid input
+// (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
+// is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
+// non-default alpha. Activations that do not use alpha simply ignore it.
 template <typename T>
-__device__ __forceinline__ T silu_kernel(const T& x) {
-  // x * sigmoid(x)
-  return (T)(((float)x) / (1.0f + expf((float)-x)));
+__device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
+  // x * sigmoid(alpha * x)
+  return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
 }
 
 template <typename packed_t>
-__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
-  // x * sigmoid(x)
+__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
+                                                       const float alpha) {
+  // x * sigmoid(alpha * x)
   float2 fval = cast_to_float2(val);
-  fval.x = fval.x / (1.0f + expf(-fval.x));
-  fval.y = fval.y / (1.0f + expf(-fval.y));
+  fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
+  fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
   return cast_to_packed<packed_t>(fval);
 }
 
 template <typename T>
-__device__ __forceinline__ T gelu_kernel(const T& x) {
+__device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'none' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -150,7 +179,8 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
 }
 
 template <typename packed_t>
-__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
+__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
+                                                       const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'none' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
@@ -162,7 +192,8 @@ __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
 }
 
 template <typename T>
-__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
+__device__ __forceinline__ T gelu_tanh_kernel(const T& x,
+                                              const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -176,7 +207,7 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
 
 template <typename packed_t>
 __device__ __forceinline__ packed_t
-packed_gelu_tanh_kernel(const packed_t& val) {
+packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
   // Equivalent to PyTorch GELU with 'tanh' approximation.
   // Refer to:
   // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
@@ -202,7 +233,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
 // clamped (max only) and up input is clamped (both sides) before the
 // activation function is applied.
 #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST,        \
-                                      HAS_CLAMP, LIMIT)                        \
+                                      HAS_CLAMP, LIMIT, ALPHA, BETA)           \
   auto dtype = input.scalar_type();                                            \
   int d = input.size(-1) / 2;                                                  \
   int64_t num_tokens = input.numel() / input.size(-1);                         \
@@ -230,7 +261,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>(     \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     } else {                                                                   \
       VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] {   \
@@ -240,7 +271,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
             PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
             ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>(    \
             out.mutable_data_ptr<scalar_t>(),                                  \
-            input.const_data_ptr<scalar_t>(), d, LIMIT);                       \
+            input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA);          \
       });                                                                      \
     }                                                                          \
   } else {                                                                     \
@@ -252,7 +283,7 @@ packed_gelu_tanh_kernel(const packed_t& val) {
           PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>,   \
           ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>(            \
           out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),  \
-          d, LIMIT);                                                           \
+          d, LIMIT, ALPHA, BETA);                                              \
     });                                                                        \
   }
 
@@ -260,14 +291,18 @@ void silu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void silu_and_mul_clamp(torch::stable::Tensor& out,    // [..., d]
                         torch::stable::Tensor& input,  // [..., 2 * d]
-                        double limit) {
+                        double limit, double alpha, double beta) {
+  // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
+  //       * (up.clamp(+-limit) + beta)
+  // alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                true, true, (float)limit);
+                                true, true, (float)limit, (float)alpha,
+                                (float)beta);
 }
 
 void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
@@ -276,21 +311,22 @@ void mul_and_silu(torch::stable::Tensor& out,    // [..., d]
   // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
   // applies the silu to the latter half of the input.
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
-                                false, false, 0.0f);
+                                false, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_and_mul(torch::stable::Tensor& out,    // [..., d]
                   torch::stable::Tensor& input)  // [..., 2 * d]
 {
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
-                                true, false, 0.0f);
+                                true, false, 0.0f, 1.0f, 0.0f);
 }
 
 void gelu_tanh_and_mul(torch::stable::Tensor& out,    // [..., d]
                        torch::stable::Tensor& input)  // [..., 2 * d]
 {
-  LAUNCH_ACTIVATION_GATE_KERNEL(
-      vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
+                                vllm::packed_gelu_tanh_kernel, true, false,
+                                0.0f, 1.0f, 0.0f);
 }
 
 namespace vllm {

csrc/libtorch_stable/fp32_router_gemm.cu

@@ -175,49 +175,52 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a,
 }
 
 // ---------------------------------------------------------------------------
-// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
+// Explicit instantiations: M=1..32, for both input types, for the supported
+// (E, H) pairs:  (256, 3072) [MiniMax-M2/M2.5]  and  (128, 6144) [MiniMax-M3].
 // ---------------------------------------------------------------------------
 
-#define INSTANTIATE(T, M)                              \
-  template void invokeFp32RouterGemm<T, M, 256, 3072>( \
-      float*, T const*, float const*, cudaStream_t);
+#define INSTANTIATE(T, M, E, H)                                    \
+  template void invokeFp32RouterGemm<T, M, E, H>(float*, T const*, \
+                                                 float const*, cudaStream_t);
 
-#define INSTANTIATE_ALL(T) \
-  INSTANTIATE(T, 1)        \
-  INSTANTIATE(T, 2)        \
-  INSTANTIATE(T, 3)        \
-  INSTANTIATE(T, 4)        \
-  INSTANTIATE(T, 5)        \
-  INSTANTIATE(T, 6)        \
-  INSTANTIATE(T, 7)        \
-  INSTANTIATE(T, 8)        \
-  INSTANTIATE(T, 9)        \
-  INSTANTIATE(T, 10)       \
-  INSTANTIATE(T, 11)       \
-  INSTANTIATE(T, 12)       \
-  INSTANTIATE(T, 13)       \
-  INSTANTIATE(T, 14)       \
-  INSTANTIATE(T, 15)       \
-  INSTANTIATE(T, 16)       \
-  INSTANTIATE(T, 17)       \
-  INSTANTIATE(T, 18)       \
-  INSTANTIATE(T, 19)       \
-  INSTANTIATE(T, 20)       \
-  INSTANTIATE(T, 21)       \
-  INSTANTIATE(T, 22)       \
-  INSTANTIATE(T, 23)       \
-  INSTANTIATE(T, 24)       \
-  INSTANTIATE(T, 25)       \
-  INSTANTIATE(T, 26)       \
-  INSTANTIATE(T, 27)       \
-  INSTANTIATE(T, 28)       \
-  INSTANTIATE(T, 29)       \
-  INSTANTIATE(T, 30)       \
-  INSTANTIATE(T, 31)       \
-  INSTANTIATE(T, 32)
+#define INSTANTIATE_ALL(T, E, H) \
+  INSTANTIATE(T, 1, E, H)        \
+  INSTANTIATE(T, 2, E, H)        \
+  INSTANTIATE(T, 3, E, H)        \
+  INSTANTIATE(T, 4, E, H)        \
+  INSTANTIATE(T, 5, E, H)        \
+  INSTANTIATE(T, 6, E, H)        \
+  INSTANTIATE(T, 7, E, H)        \
+  INSTANTIATE(T, 8, E, H)        \
+  INSTANTIATE(T, 9, E, H)        \
+  INSTANTIATE(T, 10, E, H)       \
+  INSTANTIATE(T, 11, E, H)       \
+  INSTANTIATE(T, 12, E, H)       \
+  INSTANTIATE(T, 13, E, H)       \
+  INSTANTIATE(T, 14, E, H)       \
+  INSTANTIATE(T, 15, E, H)       \
+  INSTANTIATE(T, 16, E, H)       \
+  INSTANTIATE(T, 17, E, H)       \
+  INSTANTIATE(T, 18, E, H)       \
+  INSTANTIATE(T, 19, E, H)       \
+  INSTANTIATE(T, 20, E, H)       \
+  INSTANTIATE(T, 21, E, H)       \
+  INSTANTIATE(T, 22, E, H)       \
+  INSTANTIATE(T, 23, E, H)       \
+  INSTANTIATE(T, 24, E, H)       \
+  INSTANTIATE(T, 25, E, H)       \
+  INSTANTIATE(T, 26, E, H)       \
+  INSTANTIATE(T, 27, E, H)       \
+  INSTANTIATE(T, 28, E, H)       \
+  INSTANTIATE(T, 29, E, H)       \
+  INSTANTIATE(T, 30, E, H)       \
+  INSTANTIATE(T, 31, E, H)       \
+  INSTANTIATE(T, 32, E, H)
 
-INSTANTIATE_ALL(float)
-INSTANTIATE_ALL(__nv_bfloat16)
+INSTANTIATE_ALL(float, 256, 3072)
+INSTANTIATE_ALL(__nv_bfloat16, 256, 3072)
+INSTANTIATE_ALL(float, 128, 6144)
+INSTANTIATE_ALL(__nv_bfloat16, 128, 6144)
 
 #undef INSTANTIATE_ALL
 #undef INSTANTIATE

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -22,36 +22,42 @@ inline int getSMVersion() {
 
 }  // namespace
 
-static constexpr int FP32_NUM_EXPERTS = 256;
-static constexpr int FP32_HIDDEN_DIM = 3072;
 static constexpr int FP32_MAX_TOKENS = 32;
 
+// Supported (hidden_dim, num_experts) pairs (must match the instantiations in
+// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3.
+static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) {
+  return (hidden_dim == 3072 && num_experts == 256) ||
+         (hidden_dim == 6144 && num_experts == 128);
+}
+
 // Forward declarations — 4 template params must match fp32_router_gemm.cu
 template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
 void invokeFp32RouterGemm(float* output, InputT const* mat_a,
                           float const* mat_b, cudaStream_t stream);
 
-// LoopUnroller templated on InputT
-template <typename InputT, int kBegin, int kEnd>
+// LoopUnroller templated on InputT, kNumExperts and kHiddenDim
+template <typename InputT, int kNumExperts, int kHiddenDim, int kBegin,
+          int kEnd>
 struct Fp32LoopUnroller {
   static void unroll(int num_tokens, float* output, InputT const* mat_a,
                      float const* mat_b, cudaStream_t stream) {
     if (num_tokens == kBegin) {
-      invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+      invokeFp32RouterGemm<InputT, kBegin, kNumExperts, kHiddenDim>(
           output, mat_a, mat_b, stream);
     } else {
-      Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
-                                                         mat_a, mat_b, stream);
+      Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kBegin + 1,
+                       kEnd>::unroll(num_tokens, output, mat_a, mat_b, stream);
     }
   }
 };
 
-template <typename InputT, int kEnd>
-struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
+template <typename InputT, int kNumExperts, int kHiddenDim, int kEnd>
+struct Fp32LoopUnroller<InputT, kNumExperts, kHiddenDim, kEnd, kEnd> {
   static void unroll(int num_tokens, float* output, InputT const* mat_a,
                      float const* mat_b, cudaStream_t stream) {
     if (num_tokens == kEnd) {
-      invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
+      invokeFp32RouterGemm<InputT, kEnd, kNumExperts, kHiddenDim>(
           output, mat_a, mat_b, stream);
     } else {
       throw std::invalid_argument(
@@ -60,6 +66,23 @@ struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
   }
 };
 
+// Dispatch over the supported (num_experts, hidden_dim) pairs.
+template <typename InputT>
+void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens,
+                            float* output, InputT const* mat_a,
+                            float const* mat_b, cudaStream_t stream) {
+  if (num_experts == 256 && hidden_dim == 3072) {
+    Fp32LoopUnroller<InputT, 256, 3072, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, output, mat_a, mat_b, stream);
+  } else if (num_experts == 128 && hidden_dim == 6144) {
+    Fp32LoopUnroller<InputT, 128, 6144, 1, FP32_MAX_TOKENS>::unroll(
+        num_tokens, output, mat_a, mat_b, stream);
+  } else {
+    throw std::invalid_argument(
+        "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair");
+  }
+}
+
 void fp32_router_gemm(
     torch::stable::Tensor& output,       // [num_tokens, num_experts]
     torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
@@ -85,10 +108,10 @@ void fp32_router_gemm(
   STD_TORCH_CHECK(
       mat_a.size(1) == mat_b.size(1),
       "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
-  STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
-                  "fp32_router_gemm: expected hidden_dim=3072");
-  STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
-                  "fp32_router_gemm: expected num_experts=256");
+  STD_TORCH_CHECK(
+      fp32_router_gemm_supported(hidden_dim, num_experts),
+      "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are "
+      "(3072, 256) and (6144, 128)");
   STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
                   "fp32_router_gemm: num_tokens must be in [0, 32]");
   STD_TORCH_CHECK(
@@ -113,12 +136,13 @@ void fp32_router_gemm(
   if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
     auto const* mat_a_ptr =
         reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
-    Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
-        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+    dispatchFp32RouterGemm<__nv_bfloat16>(num_experts, hidden_dim, num_tokens,
+                                          out_ptr, mat_a_ptr, mat_b_ptr,
+                                          stream);
   } else {
     auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
-    Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
-        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
+    dispatchFp32RouterGemm<float>(num_experts, hidden_dim, num_tokens, out_ptr,
+                                  mat_a_ptr, mat_b_ptr, stream);
   }
 }
 

csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

@@ -18,7 +18,7 @@
  *   ROPE_DIM  = 64   (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
  *   NOPE_DIM  = 448
  *   QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
- *   FP8_MAX   = 448.0f
+ *   FP8_MAX   = 224.0f on ROCm FNUZ / 448.0f on OCP
  *   is_neox=false (GPT-J interleaved pairs)
  *   cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
  *     second along last dim; each half is rope_dim/2 = 32 values)
@@ -61,10 +61,11 @@
 #ifdef USE_ROCM
 // ROCm-compatible FP8 conversion helpers
 __device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
-  #if defined(HIP_FP8_TYPE_OCP)
-  __hip_fp8_e4m3 fp8_val(val);
-  #else
+  // gfx942 uses FNUZ FP8; other ROCm targets use OCP E4M3.
+  #if defined(__gfx942__)
   __hip_fp8_e4m3_fnuz fp8_val(val);
+  #else
+  __hip_fp8_e4m3 fp8_val(val);
   #endif
   return reinterpret_cast<uint8_t&>(fp8_val);
 }
@@ -90,7 +91,13 @@ constexpr int kQuantBlock = 64;
 constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock;   // 7
 constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1;  // 8 (7 real + 1 pad)
 constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2;  // 448 + 128 = 576
+// FNUZ on gfx942 / OCP elsewhere. FNUZ uses 224.0 (not the dtype's raw
+// 240.0) to match the rest of vLLM's FNUZ pipeline.
+#if defined(USE_ROCM) && defined(__gfx942__)
+constexpr float kFp8Max = 224.0f;
+#else
 constexpr float kFp8Max = 448.0f;
+#endif
 
 #ifndef USE_ROCM
 // When num_tokens is less than this threshold,

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -0,0 +1,675 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Horizontally-fused MiniMax-M3 attention pre-processing kernel.
+ *
+ * Replaces the per-token Python sequence in
+ * ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
+ *
+ *     q  = q_norm(q);  k = k_norm(k);  q, k = rotary_emb(pos, q, k)
+ *     index_q = index_q_norm(index_q);  index_k = index_k_norm(index_k)
+ *     index_q, index_k = rotary_emb(pos, index_q, index_k)
+ *     _insert_kv(k, v, index_k)
+ *
+ * All branches share head_dim=128 and the *same* partial-NeoX RoPE table
+ * (``rotary_dim`` rotated, the trailing dims pass through).  The four norms
+ * are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
+ * independent weights.
+ *
+ * Everything lives in a single fused ``qkv`` tensor.  The sparse layer's
+ * fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
+ *
+ *     [ q | k | v | index_q | index_k ]   (the "5 results")
+ *
+ * while the dense layer emits just ``[ q | k | v ]``.  The kernel reads the
+ * index branch straight out of that packed row -- no separate index tensors.
+ *
+ * One kernel, one grid; each warp owns one (token, head-slot) pair.  Slot
+ * enumeration per token:
+ *     [0, nq)                         Q  heads   -> norm(q_w)  + RoPE, write
+ * qkv [nq, nq+nkv)                    K  heads   -> norm(k_w)  + RoPE, write
+ * qkv
+ *                                                  (+ insert into key cache)
+ *     [nq+nkv, nq+2*nkv)              V  heads   -> insert into value cache
+ *     IQ heads (niq)                             -> norm(iq_w) + RoPE, write iq
+ *     IK       (1)                               -> norm(ik_w) + RoPE
+ *                                                  (+ insert into index cache)
+ *
+ * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
+ * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
+ *
+ * Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
+ * template bools (3 instantiations: dense <false,false>, sparse-profiling
+ * <true,false>, sparse-serving <true,true>), so the index slots, the V slots
+ * and the cache inserts fold away entirely on paths that don't use them. The
+ * dense layer passes no caches/index: norm+RoPE happens in place and the
+ * generic ``Attention`` layer owns the cache write.
+ *
+ * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
+ * ``qkv`` tensor.  Caches (bf16) are scatter-written by slot.
+ */
+
+#include <cmath>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+#include "torch_utils.h"
+
+#include "../cuda_compat.h"
+#include "../type_convert.cuh"
+#include "../attention/dtype_fp8.cuh"
+#include "dispatch_utils.h"
+
+#ifdef USE_ROCM
+  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
+#else
+  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+#endif
+
+#ifndef FINAL_MASK
+  #ifdef USE_ROCM
+    #define FINAL_MASK 0xffffffffffffffffULL
+  #else
+    #define FINAL_MASK 0xffffffffu
+  #endif
+#endif
+
+namespace vllm {
+namespace minimax_m3_fused_ops {
+
+namespace {
+inline int getSMVersion() {
+  auto* props = get_device_prop();
+  return props->major * 10 + props->minor;
+}
+}  // namespace
+
+// ────────────────────────────────────────────────────────────────────────────
+// Constants (hard-coded for MiniMax-M3-preview).
+// ────────────────────────────────────────────────────────────────────────────
+constexpr int kHeadDim = 128;
+constexpr int kNumLanes = 32;
+constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 4
+
+// ────────────────────────────────────────────────────────────────────────────
+// Helpers
+// ────────────────────────────────────────────────────────────────────────────
+__device__ __forceinline__ float warpReduceSum(float val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  }
+  return val;
+}
+
+// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``), rounded
+// back to scalar_t like the materialized unfused norm output, followed by
+// partial NeoX RoPE on the leading ``rotary_dim`` dims. Each lane owns
+// ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
+template <typename scalar_t>
+__device__ __forceinline__ void normAndRope(
+    float (&elems)[kElemsPerLane], int const laneId, float const eps,
+    scalar_t const* __restrict__ weight,  // [kHeadDim] or nullptr (no norm)
+    bool const do_rope, int const rotary_dim,
+    scalar_t const* __restrict__ cos_ptr,  // cos_sin_cache + pos*rotary_dim
+    bool const apply_norm) {
+  // ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
+  if (apply_norm) {
+    float sumsq = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
+    sumsq = warpReduceSum(sumsq);
+    float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      int const dim = laneId * kElemsPerLane + i;
+      float const w = 1.0f + static_cast<float>(weight[dim]);
+      elems[i] = elems[i] * rms_rcp * w;
+    }
+  }
+
+  // ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
+  // half = rotary_dim/2.  Pair (i, i+half) for i in [0, half).  Lane L owns
+  // dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
+  // first half (own=x[i]) or second half (own=x[i+half]); its partner lives
+  // ``half/4`` lanes away (XOR with that distance).
+  if (do_rope) {
+    int const half = rotary_dim / 2;
+    int const dim0 = laneId * kElemsPerLane;
+    bool const in_rope = dim0 < rotary_dim;
+    int const lane_xor = half / kElemsPerLane;  // partner-lane distance
+
+    float partner[kElemsPerLane];
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
+    }
+    if (in_rope) {
+      bool const first_half = dim0 < half;
+      int const i_base = first_half ? dim0 : (dim0 - half);  // cos/sin index
+      scalar_t const* sin_ptr = cos_ptr + half;
+#pragma unroll
+      for (int i = 0; i < kElemsPerLane; i++) {
+        float const c = static_cast<float>(cos_ptr[i_base + i]);
+        float const s = static_cast<float>(sin_ptr[i_base + i]);
+        if (first_half) {
+          elems[i] = elems[i] * c - partner[i] * s;
+        } else {
+          elems[i] = elems[i] * c + partner[i] * s;
+        }
+      }
+    }
+  }
+}
+
+// Load 4 contiguous bf16 -> 4 fp32 registers.
+template <typename scalar_t>
+__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
+                                          float (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  uint2 v = *reinterpret_cast<uint2 const*>(src);
+  auto const* p =
+      reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
+#pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    float2 f2 = Converter::convert(p[i]);
+    elems[2 * i] = f2.x;
+    elems[2 * i + 1] = f2.y;
+  }
+}
+
+// Store 4 fp32 registers -> 4 contiguous bf16.
+template <typename scalar_t>
+__device__ __forceinline__ void storeElems(
+    scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  using Converter = vllm::_typeConvert<scalar_t>;
+  uint2 v;
+  auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
+#pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
+  }
+  *reinterpret_cast<uint2*>(dst) = v;
+}
+
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+__device__ __forceinline__ void storeCacheElems(
+    cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
+    // kAuto means unquantized KV cache here: cache_t == scalar_t, so store the
+    // model dtype directly. FP8 cache dtypes use the conversion path below.
+    storeElems<scalar_t>(reinterpret_cast<scalar_t*>(dst), elems);
+  } else {
+#pragma unroll
+    for (int i = 0; i < kElemsPerLane; i++) {
+      dst[i] = fp8::scaled_convert<cache_t, float, kv_dt>(elems[i], 1.0f);
+    }
+  }
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Kernel
+// ────────────────────────────────────────────────────────────────────────────
+// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
+// Each warp = one (token, slot).
+//
+// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
+// branch decisions that distinguish the dense layer from the sparse layer
+// (index slots, KV/index inserts, V slots) fold away per instantiation.
+// Three instantiations are built: dense <false,false>, sparse-profiling
+// <true,false> and sparse-serving <true,true>.  Slots per token:
+//     Q : nq                            (always — norm+RoPE)
+//     K : nkv                           (always — norm+RoPE; +K-cache insert)
+//     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
+//     IQ: niq  only if kIsSparse        (norm+RoPE)
+//     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
+          bool kIsSparse, bool kInsertKV>
+__global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
+    scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
+    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
+    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
+    scalar_t const* __restrict__ q_norm_w,
+    scalar_t const* __restrict__ k_norm_w,
+    scalar_t const* __restrict__ iq_norm_w,
+    scalar_t const* __restrict__ ik_norm_w,
+    scalar_t const* __restrict__ cos_sin_cache,  // [max_pos, rotary_dim]
+    int64_t const* __restrict__ positions,       // [N] i64
+    int64_t const* __restrict__ slot_mapping,    // main K/V slots or nullptr
+    int64_t const* __restrict__ index_slot_mapping,  // index K slots/nullptr
+    cache_t* __restrict__ kv_cache,      // [nb,2,bs,nkv,128] or nullptr
+    scalar_t* __restrict__ index_cache,  // [nb*bs, 128] or nullptr
+    float const eps, int const rotary_dim, int const num_tokens, int const nq,
+    int const nkv, int const niq, int const block_size,
+    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
+    // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
+    // NHD/HND layout choice is fully captured by these four strides: NHD keeps
+    // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
+    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
+    int64_t const kv_s_head) {
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
+  // runs with bf16/fp16 inputs in practice.  Discard the bf16 body there.
+  if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
+    return;
+  } else {
+#endif
+    int const warpsPerBlock = blockDim.x / 32;
+    int const laneId = threadIdx.x % 32;
+    int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
+
+    // Slot layout (compile-time gated: dense has neither V nor index slots).
+    int const v_slots = kInsertKV ? nkv : 0;
+    int const idx_slots = kIsSparse ? niq + 1 : 0;
+    int const slots_per_token = nq + nkv + v_slots + idx_slots;
+
+    int const tokenIdx = globalWarpIdx / slots_per_token;
+    int const slot = globalWarpIdx % slots_per_token;
+    if (tokenIdx >= num_tokens) return;
+
+    // Slot boundaries.
+    int const k_begin = nq;
+    int const v_begin = nq + nkv;             // valid only when kInsertKV
+    int const iq_begin = nq + nkv + v_slots;  // index block start
+    int const ik_slot = iq_begin + niq;       // valid only when kIsSparse
+
+    bool const isQ = slot < k_begin;
+    bool const isK = slot >= k_begin && slot < v_begin;
+    bool isV = false;
+    if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
+    bool isIQ = false, isIK = false;
+    if constexpr (kIsSparse) {
+      isIQ = slot >= iq_begin && slot < ik_slot;
+      isIK = slot == ik_slot;
+    }
+
+    int const dim_base = laneId * kElemsPerLane;
+    // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
+    // layer additionally packs [index_q (niq heads) | index_k (1 head)].
+    int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
+
+    // ── Resolve source pointer + per-branch parameters. ────────────────────
+    scalar_t* row_ptr = nullptr;       // in-place output location
+    scalar_t const* norm_w = nullptr;  // nullptr -> skip norm (V)
+    bool do_rope = true;
+    int head = 0;  // kv head index for inserts
+
+    if (isQ) {
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = q_norm_w;
+    } else if (isK) {
+      head = slot - k_begin;
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = k_norm_w;
+    } else if (isV) {
+      // qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
+      // correct in-tensor offset.
+      head = slot - v_begin;
+      row_ptr =
+          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
+      norm_w = nullptr;  // V: no norm, no rope
+      do_rope = false;
+    } else if (isIQ) {
+      // index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
+      int const ih = slot - iq_begin;
+      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
+                (nq + 2 * nkv + ih) * kHeadDim;
+      norm_w = iq_norm_w;
+    } else {  // isIK -- single shared index key at (nq+2*nkv+niq)*128.
+      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
+                (nq + 2 * nkv + niq) * kHeadDim;
+      norm_w = ik_norm_w;
+    }
+
+    // Store destination.  Q and index_q are gathered into dedicated contiguous
+    // output buffers (when provided) so the downstream SM100 sparse kernel's
+    // flat TMA descriptor can address them as [tokens*heads, head_dim]; this
+    // folds the de-interleaving into the store the kernel already does, instead
+    // of a separate q.contiguous() copy.  Everything else stays in place.
+    scalar_t* store_ptr = row_ptr;
+    if (isQ && q_out != nullptr) {
+      store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
+                  slot * kHeadDim;
+    } else if (isIQ && index_q_out != nullptr) {
+      store_ptr = index_q_out +
+                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                  (slot - iq_begin) * kHeadDim;
+    }
+
+    // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
+    // produces ``qkv``) to finish before touching any global memory.  No-op
+    // when PDL is not enabled on the launch.  The CUDA runtime wrapper emits
+    // the griddepcontrol.wait PTX with the required memory clobber internally.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaGridDependencySynchronize();
+#endif
+
+    // ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
+    float elems[kElemsPerLane];
+    loadElems<scalar_t>(row_ptr + dim_base, elems);
+
+    if (!isV) {
+      int64_t const pos = positions[tokenIdx];
+      scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
+      normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
+                            cos_ptr, /*apply_norm=*/norm_w != nullptr);
+      storeElems<scalar_t>(store_ptr + dim_base, elems);
+    }
+
+    // ── Cache inserts (sparse serving only). ───────────────────────────────
+    if constexpr (kInsertKV) {
+      // Guard (not early-return) so every thread reaches the PDL trigger below.
+      int64_t const sm = (isK || isV)
+                             ? slot_mapping[tokenIdx]
+                             : (isIK ? index_slot_mapping[tokenIdx] : -1);
+      if (sm >= 0) {  // skip padded / unscheduled tokens
+        if (isIK) {
+          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
+          storeElems<scalar_t>(dst, elems);
+        } else if (isK || isV) {
+          // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
+          // Paging is logical (block = sm/block_size, token = sm%block_size);
+          // the physical NHD/HND layout is honoured via the passed strides.
+          int64_t const b = sm / block_size;
+          int64_t const t = sm % block_size;
+          int const kv = isK ? 0 : 1;
+          int64_t const off =
+              b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
+          storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off + dim_base,
+                                                    elems);
+        }
+      }
+    }
+
+    // PDL: signal that this kernel is done so a dependent successor may launch
+    // early.  No-op when PDL is not enabled on the launch.
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+    cudaTriggerProgrammaticLaunchCompletion();
+#endif
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
+  }
+#endif
+}
+
+// ────────────────────────────────────────────────────────────────────────────
+// Launch wrapper
+// ────────────────────────────────────────────────────────────────────────────
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
+                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
+                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
+                          scalar_t const* cos_sin_cache,
+                          int64_t const* positions, int64_t const* slot_mapping,
+                          int64_t const* index_slot_mapping, cache_t* kv_cache,
+                          scalar_t* index_cache, float const eps,
+                          int const rotary_dim, int const num_tokens,
+                          int const nq, int const nkv, int const niq,
+                          int const block_size, int64_t const kv_s_block,
+                          int64_t const kv_s_kv, int64_t const kv_s_token,
+                          int64_t const kv_s_head, bool const has_index,
+                          bool const insert_kv, cudaStream_t stream) {
+  // Slot count must match the kernel's compile-time gating.
+  int const v_slots = insert_kv ? nkv : 0;
+  int const idx_slots = has_index ? niq + 1 : 0;
+  int const slots_per_token = nq + nkv + v_slots + idx_slots;
+
+  constexpr int kBlockSize = 256;
+  constexpr int kWarpsPerBlock = kBlockSize / 32;
+  int64_t const total_warps =
+      static_cast<int64_t>(num_tokens) * slots_per_token;
+  int const grid =
+      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
+  if (grid == 0) return;
+
+#ifndef USE_ROCM
+  // PDL: enable programmatic stream serialization whenever the hardware
+  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable, so
+  // leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
+  static int const sm_version = getSMVersion();
+  cudaLaunchConfig_t config;
+  config.gridDim = dim3(grid);
+  config.blockDim = dim3(kBlockSize);
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  config.attrs = attrs;
+  config.numAttrs = (sm_version >= 90) ? 1 : 0;
+
+  #define LAUNCH(IS_SPARSE, INSERT)                                           \
+    cudaLaunchKernelEx(                                                       \
+        &config,                                                              \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,       \
+                                              IS_SPARSE, INSERT>,             \
+        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,    \
+        cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
+        index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size,   \
+        kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
+#else
+  // ROCm: standard kernel launch syntax (no PDL/stream serialization).
+  // clang-format off
+  #define LAUNCH(IS_SPARSE, INSERT)                                          \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,          \
+                                          IS_SPARSE, INSERT>                 \
+        <<<grid, kBlockSize, 0, stream>>>(                                   \
+            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
+            ik_norm_w, cos_sin_cache, positions, slot_mapping,               \
+            index_slot_mapping, kv_cache, index_cache, eps, rotary_dim,      \
+            num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv,       \
+            kv_s_token, kv_s_head)
+  // clang-format on
+#endif
+
+  if (has_index) {
+    if (insert_kv) {
+      LAUNCH(true, true);  // sparse serving
+    } else {
+      LAUNCH(true, false);  // sparse profiling
+    }
+  } else {
+    // Dense layer: never has an index branch and never inserts here (the
+    // generic Attention layer owns the KV insert).
+    LAUNCH(false, false);
+  }
+#undef LAUNCH
+}
+
+}  // namespace minimax_m3_fused_ops
+}  // namespace vllm
+
+#define CALL_FUSED_MINIMAX_M3(_RAW_T, CACHE_T, KV_DTYPE)                       \
+  vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>(     \
+      reinterpret_cast<st*>(qkv.data_ptr()),                                   \
+      q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr,  \
+      index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
+                              : nullptr,                                       \
+      reinterpret_cast<st const*>(q_norm_weight.data_ptr()),                   \
+      reinterpret_cast<st const*>(k_norm_weight.data_ptr()),                   \
+      has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
+                : nullptr,                                                     \
+      has_index ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr()) \
+                : nullptr,                                                     \
+      reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),                   \
+      reinterpret_cast<int64_t const*>(positions.data_ptr()),                  \
+      insert_kv ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())   \
+                : nullptr,                                                     \
+      insert_kv ? reinterpret_cast<int64_t const*>(                            \
+                      effective_index_slot_mapping->data_ptr())                \
+                : nullptr,                                                     \
+      insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr,  \
+      (insert_kv && has_index)                                                 \
+          ? reinterpret_cast<st*>(index_cache->data_ptr())                     \
+          : nullptr,                                                           \
+      static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq,   \
+      nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
+      kv_s_head, has_index, insert_kv, stream)
+
+// ────────────────────────────────────────────────────────────────────────────
+// Torch op wrapper
+// ────────────────────────────────────────────────────────────────────────────
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv,  // [N, qkv_row] (packs index if sparse)
+    torch::stable::Tensor const& q_norm_weight,  // [128]
+    torch::stable::Tensor const& k_norm_weight,  // [128]
+    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rotary_dim]
+    torch::stable::Tensor const& positions,      // [N] i64
+    int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,  // [128]
+    std::optional<torch::stable::Tensor> index_k_norm_weight,  // [128]
+    int64_t num_index_heads,                                  // niq; 0 => dense
+    std::optional<torch::stable::Tensor> slot_mapping,        // [N] i64
+    std::optional<torch::stable::Tensor> index_slot_mapping,  // [N] i64
+    std::optional<torch::stable::Tensor> kv_cache,     // [nb,2,bs,nkv,128]
+    std::optional<torch::stable::Tensor> index_cache,  // [nb,bs,128]
+    int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,  // [N, nq*128] contiguous
+    std::optional<torch::stable::Tensor>
+        index_q_out,  // [N, niq*128] contiguous
+    const std::string& kv_cache_dtype) {
+  STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
+                  "qkv must be contiguous CUDA");
+  STD_TORCH_CHECK(
+      qkv.scalar_type() == torch::headeronly::ScalarType::Half ||
+          qkv.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "qkv must be float16 or bfloat16");
+  STD_TORCH_CHECK(
+      positions.is_cuda() &&
+          positions.scalar_type() == torch::headeronly::ScalarType::Long,
+      "positions must be int64 CUDA");
+  STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
+                  "cos_sin_cache must be contiguous CUDA");
+  STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
+                  "cos_sin_cache dtype must match qkv");
+  STD_TORCH_CHECK(
+      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
+      "cos_sin_cache shape [max_pos, rotary_dim]");
+
+  STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
+                      k_norm_weight.scalar_type() == qkv.scalar_type(),
+                  "q/k norm weight dtype must match qkv");
+  STD_TORCH_CHECK(
+      q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
+          k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
+      "q/k norm weight must have 128 elements");
+  STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
+                      rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
+                  "rotary_dim must be a positive multiple of 8 and <= 128");
+
+  int const num_tokens = static_cast<int>(qkv.size(0));
+  int const nq = static_cast<int>(num_heads);
+  int const nkv = static_cast<int>(num_kv_heads);
+  int const niq = static_cast<int>(num_index_heads);
+
+  // The sparse layer packs the index branch ([index_q (niq heads) | index_k
+  // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
+  bool const has_index = niq > 0;
+  bool const insert_kv = kv_cache.has_value();
+  vllm::Fp8KVCacheDataType const kv_dt =
+      vllm::get_fp8_kv_cache_data_type(kv_cache_dtype);
+  int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
+  int const expected_row =
+      (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
+  STD_TORCH_CHECK(qkv.size(1) == expected_row,
+                  "qkv last dim must be (num_heads + 2*num_kv_heads"
+                  " + num_index_heads + 1) * 128 for sparse, "
+                  "(num_heads + 2*num_kv_heads) * 128 for dense");
+
+  // Only the sparse layer inserts here (dense lets the generic Attention layer
+  // own the KV write); there is no dense+insert kernel instantiation.
+  STD_TORCH_CHECK(
+      !insert_kv || has_index,
+      "insert mode (kv_cache) requires the index branch (sparse layer)");
+  if (has_index) {
+    STD_TORCH_CHECK(
+        index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
+        "index branch requires both index norm weights");
+    STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
+                        index_k_norm_weight->scalar_type() == qkv.scalar_type(),
+                    "index norm weights dtype must match qkv");
+    STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
+                        index_k_norm_weight->numel() == kHeadDim,
+                    "index norm weights must have 128 elements");
+  }
+  // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
+  // off the tensor so the kernel honours whatever physical layout the attention
+  // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
+  // op argument is needed -- the strides ride along with the tensor itself.
+  int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
+  torch::stable::Tensor const* effective_index_slot_mapping = nullptr;
+  if (insert_kv) {
+    STD_TORCH_CHECK(
+        slot_mapping.has_value() && slot_mapping->is_cuda() &&
+            slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
+        "insert mode requires int64 CUDA slot_mapping");
+    STD_TORCH_CHECK(
+        !index_slot_mapping.has_value() ||
+            (index_slot_mapping->is_cuda() &&
+             index_slot_mapping->scalar_type() ==
+                 torch::headeronly::ScalarType::Long &&
+             index_slot_mapping->numel() == slot_mapping->numel()),
+        "index_slot_mapping must be int64 CUDA with slot_mapping length");
+    if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
+      STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
+                      "auto kv_cache dtype must match qkv");
+    } else {
+      STD_TORCH_CHECK(
+          kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
+          "fp8 kv_cache must use uint8 storage");
+    }
+    STD_TORCH_CHECK(index_cache.has_value() &&
+                        index_cache->scalar_type() == qkv.scalar_type(),
+                    "insert mode requires matching index_cache");
+    STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
+                    "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
+                    "head_dim (stride(4)==1)");
+    kv_s_block = kv_cache->stride(0);
+    kv_s_kv = kv_cache->stride(1);
+    kv_s_token = kv_cache->stride(2);
+    kv_s_head = kv_cache->stride(3);
+    effective_index_slot_mapping = index_slot_mapping.has_value()
+                                       ? &index_slot_mapping.value()
+                                       : &slot_mapping.value();
+  }
+  // Optional contiguous gather targets: when given, the normed/roped q (and
+  // index_q) are written here instead of in place, so callers avoid a separate
+  // .contiguous() copy.  index_q_out only makes sense on the sparse path.
+  if (q_out.has_value()) {
+    STD_TORCH_CHECK(
+        q_out->is_cuda() && q_out->is_contiguous() &&
+            q_out->scalar_type() == qkv.scalar_type(),
+        "q_out must be a contiguous CUDA tensor matching qkv dtype");
+    STD_TORCH_CHECK(
+        q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
+        "q_out must have num_tokens * num_heads * 128 elements");
+  }
+  if (index_q_out.has_value()) {
+    STD_TORCH_CHECK(
+        has_index,
+        "index_q_out requires the index branch (num_index_heads > 0)");
+    STD_TORCH_CHECK(
+        index_q_out->is_cuda() && index_q_out->is_contiguous() &&
+            index_q_out->scalar_type() == qkv.scalar_type(),
+        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
+    STD_TORCH_CHECK(index_q_out->numel() ==
+                        static_cast<int64_t>(num_tokens) * niq * kHeadDim,
+                    "index_q_out must have num_tokens * num_index_heads * 128 "
+                    "elements");
+  }
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      qkv.get_device_index());
+  auto stream = get_current_cuda_stream(qkv.get_device_index());
+
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
+        using st = scalar_t;
+        DISPATCH_BY_KV_CACHE_DTYPE(qkv.scalar_type(), kv_cache_dtype,
+                                   CALL_FUSED_MINIMAX_M3);
+      });
+}
+
+#undef CALL_FUSED_MINIMAX_M3

csrc/libtorch_stable/layernorm_kernels.cu

@@ -11,7 +11,7 @@
 namespace vllm {
 
 // TODO(woosuk): Further optimize this kernel.
-template <typename scalar_t, int VEC_SIZE, int NUM_DIMS>
+template <typename scalar_t, int VEC_SIZE, int NUM_DIMS, bool HasWeight>
 __global__ void rms_norm_kernel(
     scalar_t* __restrict__ out,           // [..., hidden_size]
     const scalar_t* __restrict__ input,   // [..., hidden_size]
@@ -20,7 +20,7 @@ __global__ void rms_norm_kernel(
     const int64_t input_stride_d4,        // input.stride(-4)
     const int64_t input_shape_d2,         // input.size(-2)
     const int64_t input_shape_d3,         // input.size(-3)
-    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size], null if !HasWeight
     const float epsilon, const int num_tokens, const int hidden_size) {
   __shared__ float s_variance;
   float variance = 0.0f;
@@ -74,11 +74,19 @@ __global__ void rms_norm_kernel(
   for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) {
     vec_n_t<scalar_t, VEC_SIZE> dst;
     vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[i];
-    vec_n_t<scalar_t, VEC_SIZE> src2 = v_w[i];
+    vec_n_t<scalar_t, VEC_SIZE> src2;
+    if constexpr (HasWeight) {
+      src2 = v_w[i];
+    }
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
+      scalar_t normalized = static_cast<scalar_t>(x * s_variance);
+      if constexpr (HasWeight) {
+        dst.val[j] = normalized * src2.val[j];
+      } else {
+        dst.val[j] = normalized;
+      }
     }
     v_out[i] = dst;
   }
@@ -88,13 +96,13 @@ __global__ void rms_norm_kernel(
    Additional optimizations we can make in this case are
    packed and vectorized operations, which help with the
    memory latency bottleneck. */
-template <typename scalar_t, int width>
+template <typename scalar_t, int width, bool HasWeight>
 __global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
 fused_add_rms_norm_kernel(
     scalar_t* __restrict__ input,  // [..., hidden_size]
     const int64_t input_stride,
     scalar_t* __restrict__ residual,      // [..., hidden_size]
-    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size], null if !HasWeight
     const float epsilon, const int num_tokens, const int hidden_size) {
   // Sanity checks on our vector struct and type-punned pointer arithmetic
   static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
@@ -136,13 +144,21 @@ fused_add_rms_norm_kernel(
     int id = blockIdx.x * vec_hidden_size + idx;
     int64_t strided_id = blockIdx.x * vec_input_stride + idx;
     _f16Vec<scalar_t, width> res = residual_v[id];
-    _f16Vec<scalar_t, width> w = weight_v[idx];
     _f16Vec<scalar_t, width> out;
     using Converter = _typeConvert<scalar_t>;
+    if constexpr (HasWeight) {
+      _f16Vec<scalar_t, width> w = weight_v[idx];
 #pragma unroll
-    for (int j = 0; j < width; ++j) {
-      float x = Converter::convert(res.data[j]);
-      out.data[j] = Converter::convert(x * s_variance) * w.data[j];
+      for (int j = 0; j < width; ++j) {
+        float x = Converter::convert(res.data[j]);
+        out.data[j] = Converter::convert(x * s_variance) * w.data[j];
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < width; ++j) {
+        float x = Converter::convert(res.data[j]);
+        out.data[j] = Converter::convert(x * s_variance);
+      }
     }
     input_v[strided_id] = out;
   }
@@ -151,13 +167,13 @@ fused_add_rms_norm_kernel(
 /* Generic fused_add_rms_norm_kernel
    The width field is not used here but necessary for other specializations.
  */
-template <typename scalar_t, int width>
+template <typename scalar_t, int width, bool HasWeight>
 __global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
 fused_add_rms_norm_kernel(
     scalar_t* __restrict__ input,  // [..., hidden_size]
     const int64_t input_stride,
     scalar_t* __restrict__ residual,      // [..., hidden_size]
-    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size], null if !HasWeight
     const float epsilon, const int num_tokens, const int hidden_size) {
   __shared__ float s_variance;
   float variance = 0.0f;
@@ -181,23 +197,29 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    input[blockIdx.x * input_stride + idx] =
-        (scalar_t)(x * s_variance) * weight[idx];
+    if constexpr (HasWeight) {
+      input[blockIdx.x * input_stride + idx] =
+          (scalar_t)(x * s_variance) * weight[idx];
+    } else {
+      input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
+    }
   }
 }
 
 }  // namespace vllm
 
-void rms_norm(torch::stable::Tensor& out,     // [..., hidden_size]
-              torch::stable::Tensor& input,   // [..., hidden_size]
-              torch::stable::Tensor& weight,  // [hidden_size]
+void rms_norm(torch::stable::Tensor& out,    // [..., hidden_size]
+              torch::stable::Tensor& input,  // [..., hidden_size]
+              std::optional<torch::stable::Tensor> weight,  // [hidden_size]
               double epsilon) {
   STD_TORCH_CHECK(out.is_contiguous());
   if (input.stride(-1) != 1) {
     input = torch::stable::contiguous(input);
   }
   STD_TORCH_CHECK(input.stride(-1) == 1);
-  STD_TORCH_CHECK(weight.is_contiguous());
+  if (weight.has_value()) {
+    STD_TORCH_CHECK(weight->is_contiguous());
+  }
 
   int hidden_size = input.size(-1);
 
@@ -215,46 +237,69 @@ void rms_norm(torch::stable::Tensor& out,     // [..., hidden_size]
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
+  const bool has_weight = weight.has_value();
   VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] {
     VLLM_STABLE_DISPATCH_FLOATING_TYPES(
         input.scalar_type(), "rms_norm_kernel", [&] {
+          const scalar_t* weight_ptr =
+              has_weight ? weight->const_data_ptr<scalar_t>() : nullptr;
           const int calculated_vec_size =
               std::gcd(16 / sizeof(scalar_t), hidden_size);
           const int block_size =
               std::min(hidden_size / calculated_vec_size, max_block_size);
           dim3 block(block_size);
           VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] {
-            vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank>
-                <<<grid, block, 0, stream>>>(
-                    out.mutable_data_ptr<scalar_t>(),
-                    input.const_data_ptr<scalar_t>(), input_stride_d2,
-                    input_stride_d3, input_stride_d4, input_shape_d2,
-                    input_shape_d3, weight.const_data_ptr<scalar_t>(), epsilon,
-                    num_tokens, hidden_size);
+            if (has_weight) {
+              vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, true>
+                  <<<grid, block, 0, stream>>>(
+                      out.mutable_data_ptr<scalar_t>(),
+                      input.const_data_ptr<scalar_t>(), input_stride_d2,
+                      input_stride_d3, input_stride_d4, input_shape_d2,
+                      input_shape_d3, weight_ptr, epsilon, num_tokens,
+                      hidden_size);
+            } else {
+              vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, false>
+                  <<<grid, block, 0, stream>>>(
+                      out.mutable_data_ptr<scalar_t>(),
+                      input.const_data_ptr<scalar_t>(), input_stride_d2,
+                      input_stride_d3, input_stride_d4, input_shape_d2,
+                      input_shape_d3, weight_ptr, epsilon, num_tokens,
+                      hidden_size);
+            }
           });
         });
   });
 }
 
-#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                \
-  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                  \
-      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {           \
-        vllm::fused_add_rms_norm_kernel<scalar_t, width>                \
-            <<<grid, block, 0, stream>>>(                               \
-                input.mutable_data_ptr<scalar_t>(), input_stride,       \
-                residual.mutable_data_ptr<scalar_t>(),                  \
-                weight.const_data_ptr<scalar_t>(), epsilon, num_tokens, \
-                hidden_size);                                           \
+#define LAUNCH_FUSED_ADD_RMS_NORM(width, has_weight)                       \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                     \
+      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {              \
+        if (has_weight) {                                                  \
+          vllm::fused_add_rms_norm_kernel<scalar_t, width, true>           \
+              <<<grid, block, 0, stream>>>(                                \
+                  input.mutable_data_ptr<scalar_t>(), input_stride,        \
+                  residual.mutable_data_ptr<scalar_t>(),                   \
+                  weight->const_data_ptr<scalar_t>(), epsilon, num_tokens, \
+                  hidden_size);                                            \
+        } else {                                                           \
+          vllm::fused_add_rms_norm_kernel<scalar_t, width, false>          \
+              <<<grid, block, 0, stream>>>(                                \
+                  input.mutable_data_ptr<scalar_t>(), input_stride,        \
+                  residual.mutable_data_ptr<scalar_t>(), nullptr, epsilon, \
+                  num_tokens, hidden_size);                                \
+        }                                                                  \
       });
 
 void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
                         torch::stable::Tensor& residual,  // [..., hidden_size]
-                        torch::stable::Tensor& weight,    // [hidden_size]
+                        std::optional<torch::stable::Tensor> weight,
                         double epsilon) {
-  STD_TORCH_CHECK(weight.scalar_type() == input.scalar_type());
   STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type());
   STD_TORCH_CHECK(residual.is_contiguous());
-  STD_TORCH_CHECK(weight.is_contiguous());
+  if (weight.has_value()) {
+    STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type());
+    STD_TORCH_CHECK(weight->is_contiguous());
+  }
   int hidden_size = input.size(-1);
   int64_t input_stride = input.stride(-2);
   int num_tokens = input.numel() / hidden_size;
@@ -269,30 +314,33 @@ void fused_add_rms_norm(torch::stable::Tensor& input,     // [..., hidden_size]
   const torch::stable::accelerator::DeviceGuard device_guard(
       input.get_device_index());
   const cudaStream_t stream = get_current_cuda_stream();
-  /*If the tensor types are FP16/BF16, try to use the optimized kernel
-    with packed + vectorized ops.
-    Max optimization is achieved with a width-8 vector of FP16/BF16s
-    since we can load at most 128 bits at once in a global memory op.
-    However, this requires each tensor's data to be aligned to 16
-    bytes.
-   */
+  constexpr int vector_width = 8;
+  constexpr int req_alignment_bytes = vector_width * 2;
   auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
   auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
-  auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
-  constexpr int vector_width = 8;
-  constexpr int req_alignment_bytes =
-      vector_width * 2;  // vector_width * sizeof(bfloat16 or float16) (float32
-                         // falls back to non-vectorized version anyway)
-  bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
-                          res_ptr % req_alignment_bytes == 0 &&
-                          wt_ptr % req_alignment_bytes == 0;
   bool offsets_are_multiple_of_vector_width =
       hidden_size % vector_width == 0 && input_stride % vector_width == 0;
   bool batch_invariant_launch = vllm::vllm_is_batch_invariant();
-  if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
-      !batch_invariant_launch) {
-    LAUNCH_FUSED_ADD_RMS_NORM(8);
+  const bool has_weight = weight.has_value();
+  if (has_weight) {
+    auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr());
+    bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
+                            res_ptr % req_alignment_bytes == 0 &&
+                            wt_ptr % req_alignment_bytes == 0;
+    if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
+        !batch_invariant_launch) {
+      LAUNCH_FUSED_ADD_RMS_NORM(8, true);
+    } else {
+      LAUNCH_FUSED_ADD_RMS_NORM(0, true);
+    }
   } else {
-    LAUNCH_FUSED_ADD_RMS_NORM(0);
+    bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&
+                            res_ptr % req_alignment_bytes == 0;
+    if (ptrs_are_aligned && offsets_are_multiple_of_vector_width &&
+        !batch_invariant_launch) {
+      LAUNCH_FUSED_ADD_RMS_NORM(8, false);
+    } else {
+      LAUNCH_FUSED_ADD_RMS_NORM(0, false);
+    }
   }
 }

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm.cu

@@ -1,69 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled.cu
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include "cutlass_mxfp8_grouped_mm_launcher.cuh"
-
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a,
-                              const torch::stable::Tensor& b,
-                              const torch::stable::Tensor& sfa,
-                              const torch::stable::Tensor& sfb,
-                              torch::stable::Tensor& d,
-                              const torch::stable::Tensor& problem_sizes,
-                              const torch::stable::Tensor& expert_offsets,
-                              const torch::stable::Tensor& blockscale_offsets) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
-                  "problem_sizes must have shape (num_experts, 3)");
-  STD_TORCH_CHECK(
-      problem_sizes.size(0) == expert_offsets.size(0),
-      "Number of experts in problem_sizes must match expert_offsets");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32");
-  STD_TORCH_CHECK(
-      expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "expert_offsets must be int32");
-  STD_TORCH_CHECK(
-      blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "blockscale_offsets must be int32");
-  STD_TORCH_CHECK(a.dim() == 2,
-                  "a must be a 2D tensor of shape (num_tokens, k)");
-  STD_TORCH_CHECK(b.dim() == 3,
-                  "b must be a 3D tensor of shape (num_experts, k, n)");
-  STD_TORCH_CHECK(a.size(1) == b.size(1) && a.size(1) % 128 == 0,
-                  "k should align 128");
-  STD_TORCH_CHECK(b.size(2) % 128 == 0, "n should align 128");
-  STD_TORCH_CHECK(a.stride(1) == 1, "a must be row major");
-  STD_TORCH_CHECK(b.stride(1) == 1, "b must be column major");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      a.get_device_index());
-  auto stream = get_current_cuda_stream(a.get_device_index());
-  if (d.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::bfloat16_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                             blockscale_offsets, stream);
-  } else if (d.scalar_type() == torch::headeronly::ScalarType::Half) {
-    expert_specialization::cutlass_mxfp8_grouped_mm_dispatch_out_dtype<
-        cutlass::half_t>(a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-                         blockscale_offsets, stream);
-  } else {
-    STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  STD_TORCH_CHECK(false,
-                  "No implemented cutlass_mxfp8_grouped_mm for "
-                  "current device");
-#endif
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("cutlass_mxfp8_grouped_mm", TORCH_BOX(&cutlass_mxfp8_grouped_mm));
-}

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_functor.cuh

@@ -1,141 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_functor.cuh
-
-#pragma once
-#include <cuda.h>
-
-#include "cute/tensor.hpp"
-#include "cutlass/util/packed_stride.hpp"
-#include "cutlass_mxfp8_grouped_mm_traits.cuh"
-
-namespace expert_specialization {
-
-using namespace cute;
-
-template <typename GemmTraits>
-struct CutlassMxfp8GroupedMmOffsetFunctor {
-  using Gemm = typename GemmTraits::Gemm;
-  using ElementA = typename Gemm::ElementA;
-  using ElementB = typename Gemm::ElementB;
-  using ElementSF = typename GemmTraits::ElementSF;
-  using ElementD = typename GemmTraits::ElementOutput;
-  // Input
-  int* expert_offsets{nullptr};
-  int* blockscale_offsets{nullptr};
-  // Output
-  ElementA* a_base{nullptr};
-  ElementB* b_base{nullptr};
-  ElementSF* sfa_base{nullptr};
-  ElementSF* sfb_base{nullptr};
-  ElementD* d_base{nullptr};
-  ElementA** a_offsets{nullptr};
-  ElementB** b_offsets{nullptr};
-  ElementSF** sfa_offsets{nullptr};
-  ElementSF** sfb_offsets{nullptr};
-  ElementD** d_offsets{nullptr};
-
-  CutlassMxfp8GroupedMmOffsetFunctor() = default;
-  CutlassMxfp8GroupedMmOffsetFunctor(
-      int* _expert_offsets, int* _blockscale_offsets, ElementA* _a_base,
-      ElementB* _b_base, ElementSF* _sfa_base, ElementSF* _sfb_base,
-      ElementD* _d_base, ElementA** _a_offsets, ElementB** _b_offsets,
-      ElementSF** _sfa_offsets, ElementSF** _sfb_offsets, ElementD** _d_offsets)
-      : expert_offsets{_expert_offsets},
-        blockscale_offsets{_blockscale_offsets},
-        a_base(_a_base),
-        b_base(_b_base),
-        sfa_base(_sfa_base),
-        sfb_base(_sfb_base),
-        d_base(_d_base),
-        a_offsets(_a_offsets),
-        b_offsets(_b_offsets),
-        sfa_offsets(_sfa_offsets),
-        sfb_offsets(_sfb_offsets),
-        d_offsets(_d_offsets) {}
-
-  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
-    int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
-    int64_t blockscale_offset =
-        static_cast<int64_t>(blockscale_offsets[expert_id]);
-    int64_t a_stride = expert_offset * k;
-    int64_t b_stride = expert_id * k * n;
-    int64_t d_stride = expert_offset * n;
-    int64_t sfa_stride = blockscale_offset * (k / 32);
-    int64_t sfb_stride = expert_id * n * (k / 32);
-
-    a_offsets[expert_id] = a_base + a_stride;
-    b_offsets[expert_id] = b_base + b_stride;
-    sfa_offsets[expert_id] = sfa_base + sfa_stride;
-    sfb_offsets[expert_id] = sfb_base + sfb_stride;
-    d_offsets[expert_id] = d_base + d_stride;
-  }
-};
-
-template <typename GemmTraits>
-struct CutlassMxfp8GroupedMmLayoutFunctor {
-  using Sm1xxBlkScaledConfig = typename GemmTraits::Sm1xxBlkScaledConfig;
-  using LayoutSFA = typename GemmTraits::LayoutSFA;
-  using LayoutSFB = typename GemmTraits::LayoutSFB;
-  LayoutSFA* layout_sfa_base{nullptr};
-  LayoutSFB* layout_sfb_base{nullptr};
-
-  CutlassMxfp8GroupedMmLayoutFunctor() = default;
-  CutlassMxfp8GroupedMmLayoutFunctor(LayoutSFA* _layout_sfa_base,
-                                     LayoutSFB* _layout_sfb_base)
-      : layout_sfa_base(_layout_sfa_base), layout_sfb_base(_layout_sfb_base) {}
-
-  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
-    LayoutSFA* layout_sfa_ptr = layout_sfa_base + expert_id;
-    LayoutSFB* layout_sfb_ptr = layout_sfb_base + expert_id;
-    *layout_sfa_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(
-        cute::make_shape(m, n, k, 1));
-    *layout_sfb_ptr = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(
-        cute::make_shape(m, n, k, 1));
-  }
-};
-
-template <typename GemmTraits>
-struct CutlassMxfp8GroupedMmStrideFunctor {
-  using StrideA = typename GemmTraits::StrideA;
-  using StrideB = typename GemmTraits::StrideB;
-  using StrideD = typename GemmTraits::StrideD;
-  StrideA* stride_A_base{nullptr};
-  StrideB* stride_B_base{nullptr};
-  StrideD* stride_D_base{nullptr};
-
-  CutlassMxfp8GroupedMmStrideFunctor() = default;
-  CutlassMxfp8GroupedMmStrideFunctor(StrideA* _stride_A_base,
-                                     StrideB* _stride_B_base,
-                                     StrideD* _stride_D_base)
-      : stride_A_base(_stride_A_base),
-        stride_B_base(_stride_B_base),
-        stride_D_base(_stride_D_base) {}
-
-  void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) {
-    StrideA* stride_A = stride_A_base + expert_id;
-    StrideB* stride_B = stride_B_base + expert_id;
-    StrideD* stride_D = stride_D_base + expert_id;
-    *stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1});
-    *stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1});
-    *stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1});
-  }
-};
-
-template <typename OffsetFunctor, typename LayoutFunctor,
-          typename StrideFunctor>
-__global__ void cutlassMxfp8GroupedMmPreComputeKernel(
-    int* problem_sizes, OffsetFunctor offset_functor,
-    LayoutFunctor layout_functor, StrideFunctor stride_functor) {
-  int64_t expert_id = static_cast<int64_t>(threadIdx.x);
-  int m = problem_sizes[expert_id * 3 + 0];
-  int n = problem_sizes[expert_id * 3 + 1];
-  int k = problem_sizes[expert_id * 3 + 2];
-
-  offset_functor(expert_id, m, n, k);
-  layout_functor(expert_id, m, n, k);
-  stride_functor(expert_id, m, n, k);
-}
-
-}  // namespace expert_specialization

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_launcher.cuh

@@ -1,198 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_launcher.cuh
-
-#pragma once
-
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
-
-#include <cassert>
-#include <iostream>
-#include <string>
-
-#include "cute/tensor.hpp"
-#include "cutlass_mxfp8_grouped_mm_functor.cuh"
-#include "cutlass_mxfp8_grouped_mm_traits.cuh"
-#include "libtorch_stable/torch_utils.h"
-
-namespace expert_specialization {
-
-template <typename GemmTraits>
-void cutlass_mxfp8_grouped_mm_pre_compute(
-    torch::stable::Tensor& a_ptrs, torch::stable::Tensor& b_ptrs,
-    torch::stable::Tensor& sfa_ptrs, torch::stable::Tensor& sfb_ptrs,
-    torch::stable::Tensor& d_ptrs, torch::stable::Tensor& stride_a,
-    torch::stable::Tensor& stride_b, torch::stable::Tensor& stride_d,
-    torch::stable::Tensor& layout_sfa, torch::stable::Tensor& layout_sfb,
-    const torch::stable::Tensor& a, const torch::stable::Tensor& b,
-    const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
-    const torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
-  using OffsetFunctor = CutlassMxfp8GroupedMmOffsetFunctor<GemmTraits>;
-  using ElementA = typename OffsetFunctor::ElementA;
-  using ElementB = typename OffsetFunctor::ElementB;
-  using ElementSF = typename OffsetFunctor::ElementSF;
-  using ElementD = typename OffsetFunctor::ElementD;
-
-  using LayoutFunctor = CutlassMxfp8GroupedMmLayoutFunctor<GemmTraits>;
-  using LayoutSFA = typename LayoutFunctor::LayoutSFA;
-  using LayoutSFB = typename LayoutFunctor::LayoutSFB;
-
-  using StrideFunctor = CutlassMxfp8GroupedMmStrideFunctor<GemmTraits>;
-  using StrideA = typename StrideFunctor::StrideA;
-  using StrideB = typename StrideFunctor::StrideB;
-  using StrideD = typename StrideFunctor::StrideD;
-
-  int num_experts = static_cast<int>(expert_offsets.size(0));
-  STD_TORCH_CHECK(num_experts <= 1024,
-                  "Number of experts cannot exceed 1024, the maximum number of "
-                  "threads per block.");
-
-  OffsetFunctor offset_functor(
-      reinterpret_cast<int*>(expert_offsets.data_ptr()),
-      reinterpret_cast<int*>(blockscale_offsets.data_ptr()),
-      reinterpret_cast<ElementA*>(a.data_ptr()),
-      reinterpret_cast<ElementB*>(b.data_ptr()),
-      reinterpret_cast<ElementSF*>(sfa.data_ptr()),
-      reinterpret_cast<ElementSF*>(sfb.data_ptr()),
-      reinterpret_cast<ElementD*>(d.data_ptr()),
-      reinterpret_cast<ElementA**>(a_ptrs.data_ptr()),
-      reinterpret_cast<ElementB**>(b_ptrs.data_ptr()),
-      reinterpret_cast<ElementSF**>(sfa_ptrs.data_ptr()),
-      reinterpret_cast<ElementSF**>(sfb_ptrs.data_ptr()),
-      reinterpret_cast<ElementD**>(d_ptrs.data_ptr()));
-  LayoutFunctor layout_functor(
-      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()));
-  StrideFunctor stride_functor(reinterpret_cast<StrideA*>(stride_a.data_ptr()),
-                               reinterpret_cast<StrideB*>(stride_b.data_ptr()),
-                               reinterpret_cast<StrideD*>(stride_d.data_ptr()));
-  cutlassMxfp8GroupedMmPreComputeKernel<<<1, num_experts, 0, stream>>>(
-      static_cast<int*>(problem_sizes.data_ptr()), offset_functor,
-      layout_functor, stride_functor);
-}
-
-template <typename GemmTraits>
-void cutlass_mxfp8_grouped_mm(const torch::stable::Tensor& a_ptrs,
-                              const torch::stable::Tensor& b_ptrs,
-                              const torch::stable::Tensor& sfa_ptrs,
-                              const torch::stable::Tensor& sfb_ptrs,
-                              const torch::stable::Tensor& d_ptrs,
-                              const torch::stable::Tensor& stride_a,
-                              const torch::stable::Tensor& stride_b,
-                              const torch::stable::Tensor& stride_d,
-                              const torch::stable::Tensor& layout_sfa,
-                              const torch::stable::Tensor& layout_sfb,
-                              const torch::stable::Tensor& problem_sizes,
-                              cudaStream_t stream) {
-  using Gemm = typename GemmTraits::Gemm;
-  using ElementA = typename Gemm::ElementA;
-  using ElementB = typename Gemm::ElementB;
-  using ElementSF = typename GemmTraits::ElementSF;
-  using ElementD = typename GemmTraits::ElementOutput;
-  using StrideA = typename GemmTraits::StrideA;
-  using StrideB = typename GemmTraits::StrideB;
-  using StrideD = typename GemmTraits::StrideD;
-  using LayoutSFA = typename GemmTraits::LayoutSFA;
-  using LayoutSFB = typename GemmTraits::LayoutSFB;
-  using UnderlyingProblemShape =
-      typename GemmTraits::ProblemShape::UnderlyingProblemShape;
-
-  cutlass::KernelHardwareInfo hw_info;
-  hw_info.device_id = d_ptrs.get_device_index();
-  hw_info.sm_count = get_device_prop()->multiProcessorCount;
-  hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
-  hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
-
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-
-  UnderlyingProblemShape* underlying_problem_shape =
-      reinterpret_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
-
-  typename Gemm::Arguments arguments = {
-      cutlass::gemm::GemmUniversalMode::kGrouped,
-      {num_experts, underlying_problem_shape, nullptr},
-      {reinterpret_cast<const ElementA**>(a_ptrs.data_ptr()),
-       reinterpret_cast<StrideA*>(stride_a.data_ptr()),
-       reinterpret_cast<const ElementB**>(b_ptrs.data_ptr()),
-       reinterpret_cast<StrideB*>(stride_b.data_ptr()),
-       reinterpret_cast<const ElementSF**>(sfa_ptrs.data_ptr()),
-       reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
-       reinterpret_cast<const ElementSF**>(sfb_ptrs.data_ptr()),
-       reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())},
-      {{},
-       nullptr,
-       nullptr,
-       reinterpret_cast<ElementD**>(d_ptrs.data_ptr()),
-       reinterpret_cast<StrideD*>(stride_d.data_ptr())},
-      hw_info,
-      {}  // Scheduler
-  };
-
-  Gemm gemm;
-
-  auto can_implement_status = gemm.can_implement(arguments);
-  STD_TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
-                  "Failed to implement GEMM");
-
-  size_t workspace_size = gemm.get_workspace_size(arguments);
-  torch::stable::Tensor workspace = torch::stable::empty(
-      {static_cast<int64_t>(workspace_size)},
-      torch::headeronly::ScalarType::Byte, std::nullopt, d_ptrs.device());
-
-  auto status = gemm.initialize(arguments, workspace.data_ptr(), stream);
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
-                  "Failed to initialize GEMM");
-
-  status = gemm.run(stream, nullptr, true);  // Enable PDL
-  STD_TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
-}
-
-template <typename OutType>
-void cutlass_mxfp8_grouped_mm_dispatch_out_dtype(
-    const torch::stable::Tensor& a, const torch::stable::Tensor& b,
-    const torch::stable::Tensor& sfa, const torch::stable::Tensor& sfb,
-    torch::stable::Tensor& d, const torch::stable::Tensor& problem_sizes,
-    const torch::stable::Tensor& expert_offsets,
-    const torch::stable::Tensor& blockscale_offsets, cudaStream_t stream) {
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-  auto device = a.device();
-
-  torch::stable::Tensor a_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor b_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor sfa_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor sfb_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor d_ptrs = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-
-  torch::stable::Tensor stride_a = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor stride_b = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor stride_d = torch::stable::empty(
-      num_experts, torch::headeronly::ScalarType::Long, std::nullopt, device);
-  torch::stable::Tensor layout_sfa =
-      torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
-                           std::nullopt, device);
-  torch::stable::Tensor layout_sfb =
-      torch::stable::empty({num_experts, 5}, torch::headeronly::ScalarType::Int,
-                           std::nullopt, device);
-
-  using GemmTraits = CutlassMxfp8GroupedMmGemmTraits<MMA1SMConfig, OutType>;
-  cutlass_mxfp8_grouped_mm_pre_compute<GemmTraits>(
-      a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
-      layout_sfa, layout_sfb, a, b, sfa, sfb, d, problem_sizes, expert_offsets,
-      blockscale_offsets, stream);
-  cutlass_mxfp8_grouped_mm<GemmTraits>(
-      a_ptrs, b_ptrs, sfa_ptrs, sfb_ptrs, d_ptrs, stride_a, stride_b, stride_d,
-      layout_sfa, layout_sfb, problem_sizes, stream);
-}
-
-}  // namespace expert_specialization

csrc/libtorch_stable/moe/mxfp8_moe/cutlass_mxfp8_grouped_mm_traits.cuh

@@ -1,127 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_traits.cuh
-
-#pragma once
-
-// Misc
-#include "cute/tensor.hpp"
-#include "cutlass/arch/arch.h"
-#include "cutlass/arch/mma.h"
-#include "cutlass/cutlass.h"
-#include "cutlass/detail/sm100_blockscaled_layout.hpp"
-#include "cutlass/epilogue/dispatch_policy.hpp"
-#include "cutlass/gemm/dispatch_policy.hpp"
-#include "cutlass/gemm/group_array_problem_shape.hpp"
-#include "cutlass/layout/layout.h"
-#include "cutlass/numeric_conversion.h"
-#include "cutlass/numeric_size.h"
-
-// Collective Builder
-#include "cutlass/epilogue/collective/collective_builder.hpp"
-#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
-#include "cutlass/epilogue/thread/activation.h"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-
-// Integration
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
-
-namespace expert_specialization {
-
-using namespace cute;
-
-// Different configs for 1SM and 2SM MMA kernel
-struct MMA1SMConfig {
-  using MmaTileShape = Shape<_128, _128, _128>;
-  using KernelSchedule =
-      cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf8f6f4Sm100;
-  using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
-  const static dim3 preferred_cluster;
-  const static dim3 fallback_cluster;
-};
-const dim3 MMA1SMConfig::preferred_cluster(1, 4, 1);
-const dim3 MMA1SMConfig::fallback_cluster(1, 2, 1);
-
-template <typename _MMAConfig, typename OutputDtype>
-struct CutlassMxfp8GroupedMmGemmTraits {
-  using MMAConfig = _MMAConfig;
-  using ElementInput = cutlass::float_e4m3_t;
-  using ElementOutput = OutputDtype;
-  using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int, int, int>>;
-
-  // A matrix configuration
-  using ElementA = cutlass::mx_float8_t<ElementInput>;
-  using LayoutA = cutlass::layout::RowMajor;
-  constexpr static int AlignmentA = 32;
-
-  // B matrix configuration
-  using ElementB = cutlass::mx_float8_t<ElementInput>;
-  using LayoutB = cutlass::layout::ColumnMajor;
-  constexpr static int AlignmentB = 32;
-
-  // C/D matrix configuration
-  using ElementC = void;
-  using ElementD = ElementOutput;
-  using LayoutC = cutlass::layout::RowMajor;
-  using LayoutD = cutlass::layout::RowMajor;
-  constexpr static int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
-  constexpr static int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
-  using ElementAccumulator = float;
-
-  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
-  using CustomEVTIdentity =  // acc
-      cutlass::epilogue::fusion::Sm90EVT<
-          cutlass::epilogue::fusion::Sm90Compute<
-              cutlass::epilogue::thread::Identity, ElementD, ElementAccumulator,
-              RoundStyle>,
-          cutlass::epilogue::fusion::Sm90AccFetch>;
-
-  // Core kernel configurations
-  using ArchTag = cutlass::arch::Sm100;
-  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
-  using StageCountType = cutlass::gemm::collective::StageCountAuto;
-
-  // Runtime Cluster Shape
-  using ClusterShape = Shape<int32_t, int32_t, _1>;
-
-  // Define Epilogue
-  using CollectiveEpilogue =
-      typename cutlass::epilogue::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, typename MMAConfig::MmaTileShape,
-          ClusterShape, Shape<_64, _64>, ElementAccumulator, ElementAccumulator,
-          ElementC, LayoutC*, AlignmentC, ElementD, LayoutD*, AlignmentD,
-          typename MMAConfig::EpilogueSchedule,
-          CustomEVTIdentity>::CollectiveOp;
-
-  // Define Mainloop
-  using CollectiveMainloop =
-      typename cutlass::gemm::collective::CollectiveBuilder<
-          ArchTag, OperatorClass, ElementA, LayoutA*, AlignmentA, ElementB,
-          LayoutB*, AlignmentB, ElementAccumulator,
-          typename MMAConfig::MmaTileShape, ClusterShape,
-          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
-              sizeof(typename CollectiveEpilogue::SharedStorage))>,
-          typename MMAConfig::KernelSchedule>::CollectiveOp;
-
-  // Define GemmKernel
-  using GemmKernel =
-      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
-                                           CollectiveEpilogue>;
-  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-
-  using ElementSF = typename Gemm::GemmKernel::ElementSF;
-  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
-  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
-  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
-  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
-  using LayoutSFA =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
-  using LayoutSFB =
-      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
-  using Sm1xxBlkScaledConfig =
-      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
-};
-
-}  // namespace expert_specialization

csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cu

@@ -1,66 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include "libtorch_stable/torch_utils.h"
-
-#include "mxfp8_experts_quant.cuh"
-
-void mxfp8_experts_quant(const torch::stable::Tensor& input,
-                         const torch::stable::Tensor& problem_sizes,
-                         const torch::stable::Tensor& expert_offsets,
-                         const torch::stable::Tensor& blockscale_offsets,
-                         torch::stable::Tensor& quant_output,
-                         torch::stable::Tensor& scale_factor) {
-#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
-  STD_TORCH_CHECK(input.dim() == 2, "input must be 2D tensor");
-  STD_TORCH_CHECK(input.size(1) % 128 == 0, "k must align to 128");
-  STD_TORCH_CHECK(input.stride(1) == 1, "input must be row major");
-  STD_TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be 2D tensor");
-  STD_TORCH_CHECK(
-      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
-      "problem_sizes must be int32");
-  STD_TORCH_CHECK(
-      expert_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "expert_offsets must be int32");
-  STD_TORCH_CHECK(
-      blockscale_offsets.scalar_type() == torch::headeronly::ScalarType::Int,
-      "blockscale_offsets must be int32");
-
-  auto groups = problem_sizes.size(0);
-  STD_TORCH_CHECK(
-      expert_offsets.dim() == 1 && expert_offsets.size(0) == groups,
-      "expert_offsets must be 1D and have size equal to the number of groups");
-  STD_TORCH_CHECK(
-      blockscale_offsets.dim() == 1 && blockscale_offsets.size(0) == groups,
-      "blockscale_offsets must be 1D and have size equal to the number of "
-      "groups");
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      input.get_device_index());
-  if (input.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    expert_specialization::launch_mxfp8_experts_quant<__nv_bfloat16>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else if (input.scalar_type() == torch::headeronly::ScalarType::Half) {
-    expert_specialization::launch_mxfp8_experts_quant<__half>(
-        input, problem_sizes, expert_offsets, blockscale_offsets, quant_output,
-        scale_factor);
-  } else {
-    STD_TORCH_CHECK(false, "dtype must be kFloat16 or kBFloat16");
-  }
-#else
-  STD_TORCH_CHECK(false,
-                  "No implemented mxfp8_experts_quant for "
-                  "current device");
-#endif
-}
-
-// Registered here (not torch_bindings.cpp) because ENABLE_ES_MXFP8_GROUPED_MM
-// is applied only under COMPILE_LANGUAGE:CUDA.
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("mxfp8_experts_quant", TORCH_BOX(&mxfp8_experts_quant));
-}

csrc/libtorch_stable/moe/mxfp8_moe/mxfp8_experts_quant.cuh

@@ -1,416 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Adapted from SGLang:
-// https://github.com/sgl-project/sglang/blob/ded068a76e00878881d52d5bfb791e0f60d7311b/sgl-kernel/csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh
-
-#pragma once
-#include <cuda.h>
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-
-#include <torch/csrc/inductor/aoti_torch/c/shim.h>
-#include <torch/csrc/stable/macros.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/util/Exception.h>
-
-#include <cuda/ptx>
-
-#include "cute/tensor.hpp"
-#include "libtorch_stable/torch_utils.h"
-
-namespace expert_specialization {
-
-using namespace cute;
-
-constexpr uint32_t THREAD_BLOCK_SIZE = 128;
-constexpr uint32_t WARP_SIZE = 32;
-constexpr int BLOCK_M = 128;
-constexpr int BLOCK_K = 128;
-using ThrLayout = Layout<Shape<_16, _8>, Stride<_8, _1>>;
-using ValLayout = Layout<Shape<_1, _16>>;
-using SfR2SThrLayout = Layout<Shape<_16, _4>, Stride<_4, _1>>;
-using SfR2SValLayout = Layout<Shape<_1, _1>>;
-using ScaleFactorTileLayout =
-    Layout<Shape<Shape<_32, _4>, _4>, Stride<Stride<_16, _4>, _1>>;
-
-// Fast reciprocal.
-inline __device__ float reciprocal_approximate_ftz(float a) {
-  float b;
-  asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
-  return b;
-}
-
-// Some code references TRT-LLM:
-// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh
-template <typename FragmentS, typename FragmentD>
-__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s,
-                                                     FragmentD& fragment_d) {
-  using FragmentSLayout = typename FragmentS::layout_type;
-  using FragmentDLayout = typename FragmentD::layout_type;
-  FragmentSLayout fragment_s_layout;
-  FragmentDLayout fragment_d_layout;
-  static_assert(is_static<FragmentSLayout>::value &&
-                size(fragment_s_layout) == 16);
-  static_assert(is_static<FragmentDLayout>::value &&
-                size(fragment_d_layout) == 16);
-
-  constexpr int eles_per_thr = 16;
-  using ValType = typename FragmentS::element_type;
-  using VecType = std::conditional_t<std::is_same_v<ValType, __nv_bfloat16>,
-                                     __nv_bfloat162, __half2>;
-  VecType vec[8];
-  // Assign vals
-  vec[0].x = fragment_s(Int<0>{});
-  vec[0].y = fragment_s(Int<1>{});
-  vec[1].x = fragment_s(Int<2>{});
-  vec[1].y = fragment_s(Int<3>{});
-  vec[2].x = fragment_s(Int<4>{});
-  vec[2].y = fragment_s(Int<5>{});
-  vec[3].x = fragment_s(Int<6>{});
-  vec[3].y = fragment_s(Int<7>{});
-  vec[4].x = fragment_s(Int<8>{});
-  vec[4].y = fragment_s(Int<9>{});
-  vec[5].x = fragment_s(Int<10>{});
-  vec[5].y = fragment_s(Int<11>{});
-  vec[6].x = fragment_s(Int<12>{});
-  vec[6].y = fragment_s(Int<13>{});
-  vec[7].x = fragment_s(Int<14>{});
-  vec[7].y = fragment_s(Int<15>{});
-
-  auto local_max = __habs2(vec[0]);
-  for (int i = 1; i < eles_per_thr / 2; i++) {
-    local_max = __hmax2(__habs2(vec[i]), local_max);
-  }
-  local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max);
-
-  // Get the final absolute maximum values.
-  float block_max(0.0f);
-  if constexpr (std::is_same_v<ValType, __nv_bfloat16>) {
-    block_max = __bfloat162float(__hmax(local_max.x, local_max.y));
-  } else {
-    block_max = __half2float(__hmax(local_max.x, local_max.y));
-  }
-  // Get the SF (max value of the vector / max value of mxfp8).
-  float sf_val = block_max * reciprocal_approximate_ftz(448.0f);
-  // 8 bits representation of the SF.
-  uint8_t fp8_sf_val;
-
-  __nv_fp8_e8m0 tmp_sf_val;
-  tmp_sf_val.__x =
-      __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf);
-  sf_val = static_cast<float>(tmp_sf_val);
-  fp8_sf_val = tmp_sf_val.__x;
-  // Get the output scale (reciprocal of the SFValue).
-  float output_scale =
-      block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f;
-
-  // Convert the input to float.
-  float2 fp2_vals[eles_per_thr / 2];
-
-#pragma unroll
-  for (int i = 0; i < eles_per_thr / 2; i++) {
-    if constexpr (std::is_same_v<ValType, __half>) {
-      fp2_vals[i] = __half22float2(vec[i]);
-    } else {
-      fp2_vals[i] = __bfloat1622float2(vec[i]);
-    }
-    fp2_vals[i].x *= output_scale;
-    fp2_vals[i].y *= output_scale;
-  }
-  union {
-    uint8_t bytes[16];
-    __nv_fp8x2_e4m3 elts[8];
-  } u;
-  u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]);
-  u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]);
-  u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]);
-  u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]);
-  u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]);
-  u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]);
-  u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]);
-  u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]);
-  fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]);
-  fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]);
-  fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]);
-  fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]);
-  fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]);
-  fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]);
-  fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]);
-  fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]);
-  fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]);
-  fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]);
-  fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]);
-  fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]);
-  fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]);
-  fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]);
-  fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]);
-  fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]);
-  return fp8_sf_val;
-}
-
-template <typename TensorS, typename TensorP, typename TensorD,
-          typename TensorSharedSF, typename TensorSF, typename TiledCopyG2R,
-          typename TiledCopyR2G, typename TiledCopyR2S>
-__inline__ __device__ void mxfp8_experts_quant_tile(
-    TensorS& tensor_s, TensorP& tensor_p, TensorD& tensor_d,
-    TensorSharedSF& tensor_shared_sf, TensorSF& tensor_sf, int m,
-    TiledCopyG2R& tiled_copy_g2r, TiledCopyR2G& tiled_copy_r2g,
-    TiledCopyR2S& tiled_copy_r2s) {
-  static_assert(size(get<0>(typename TensorS::layout_type{})) == 128 &&
-                size(get<1>(typename TensorS::layout_type{})) == 128 &&
-                stride(get<1>(typename TensorS::layout_type{})) == 1);
-  static_assert(size(get<0>(typename TensorD::layout_type{})) == 128 &&
-                size(get<1>(typename TensorD::layout_type{})) == 128 &&
-                stride(get<1>(typename TensorD::layout_type{})) == 1);
-  static_assert(size(get<0>(typename TensorP::layout_type{})) == 128 &&
-                size(get<1>(typename TensorP::layout_type{})) == 128);
-  static_assert(size(get<0>(typename TensorSharedSF::layout_type{})) == 128 &&
-                size(get<1>(typename TensorSharedSF::layout_type{})) == 4);
-  static_assert(size(get<0>(typename TensorSF::layout_type{})) == 128 &&
-                size(get<1>(typename TensorSF::layout_type{})) == 4);
-
-  using Tiler_MN = typename TiledCopyG2R::Tiler_MN;
-  auto tiler_mn = Tiler_MN{};
-  static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128);
-
-  auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn);
-  auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn);
-  auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn);
-  static_assert(size<2>(tiled_tensor_s) == 1);
-  static_assert(size<2>(tiled_tensor_p) == 1);
-  static_assert(size<2>(tiled_tensor_d) == 1);
-  auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s);
-  auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p);
-  auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d);
-
-  using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN;
-  auto sf_tiler_mn = SF_Tiler_MN{};
-  static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4);
-
-  auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn);
-  auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn);
-  auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf);
-  auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf);
-
-  constexpr int tile_loop_count = size<1>(tiled_tensor_s);
-  constexpr int rows_in_tile = 16;
-  // We don't need to clear shared memory
-  // clear(squeeze_tiled_tensor_shared_sf);
-#pragma unroll 4
-  for (int t = 0; t < tile_loop_count; t++) {
-    if (t * rows_in_tile >= m) {
-      break;
-    }
-    auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t));
-    auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t));
-    auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t));
-    auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t));
-    auto current_copy_tile_shared_sf =
-        tensor<0>(squeeze_tiled_tensor_shared_sf(_, t));
-
-    // Global to Register copy
-    auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x);
-    auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s);
-    auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p);
-    auto input_fragment = make_fragment_like(thr_tile_g2r_s);
-
-    // Register to Global copy
-    auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x);
-    auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d);
-    auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p);
-    auto output_fragment = make_fragment_like(thr_tile_r2g_d);
-
-    // Register to Shared copy
-    auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2);
-    auto thr_tile_r2s_shared_sf =
-        thr_copy_r2s.partition_D(current_copy_tile_shared_sf);
-    auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf);
-
-    // CopyG2R & convert & CopyR2G
-    copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment);
-    uint8_t fp8_sf_val =
-        cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment);
-    copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d);
-    shared_sf_fragment[0] = fp8_sf_val;
-
-    // Before first copy r2s, clear shared memory and wait previous group
-    if (t == 0 && threadIdx.x == 0) {
-      // Wait for the group to have completed reading from shared memory.
-      cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
-    }
-    __syncthreads();
-
-    if (threadIdx.x % 2 == 0) {
-      copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf);
-    }
-    __syncthreads();
-  }
-
-  // Wait for shared memory writes to be visible to TMA engine.
-  cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);  // b)
-  __syncthreads();
-
-  if (threadIdx.x == 0) {
-    cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared,
-                             squeeze_tiled_tensor_sf.data().get(),
-                             squeeze_tiled_tensor_shared_sf.data().get(), 512);
-    // Wait for TMA transfer to have finished reading shared memory.
-    // Create a "bulk async-group" out of the previous bulk copy operation.
-    cuda::ptx::cp_async_bulk_commit_group();
-  }
-  __syncthreads();
-}
-
-template <typename T_IN, typename TiledCopyG2R, typename TiledCopyR2G,
-          typename TiledCopyR2S>
-__global__ void mxfp8_experts_quant_kernel(
-    const T_IN* input, const int* problem_sizes, const int* expert_offsets,
-    const int* blockscale_offsets, cutlass::float_e4m3_t* quant_output,
-    uint8_t* scale_factor, int groups, TiledCopyG2R tiled_copy_g2r,
-    TiledCopyR2G tiled_copy_r2g, TiledCopyR2S tiled_copy_r2s) {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-  __shared__ __align__(512) uint8_t shared_memory[512];
-  ScaleFactorTileLayout scale_factor_tile_layout{};
-  auto scale_factor_shared =
-      make_tensor(make_smem_ptr(shared_memory),
-                  scale_factor_tile_layout);  // ((_32,_4), _4):((_16,_4), _1)
-  // TODO: Transform Groupwise Schedule into a more efficient Schedule
-  for (int g = 0; g < groups; g++) {
-    int m = problem_sizes[g * 3 + 0];
-    int k = problem_sizes[g * 3 + 2];
-    int64_t expert_offset = static_cast<int64_t>(expert_offsets[g]);
-    int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[g]);
-
-    auto input_tensor = make_tensor(
-        make_gmem_ptr(input + expert_offset * k),
-        make_layout(make_shape(m, k),
-                    LayoutRight{}));  // (M, K):(K, 1) half_t/bfloat16_t
-
-    auto quant_output_tensor = make_tensor(
-        make_gmem_ptr(quant_output + expert_offset * k),
-        make_layout(make_shape(m, k),
-                    LayoutRight{}));  // (M, K):(K, 1) cutlass::float_e4m3_t
-
-    auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32);
-    auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout,
-                                             scale_factor_shape, LayoutRight{});
-    // layout<0>(layout<0>(scale_factor_layout))  (_32,_4):(_16,_4) -- static
-    // layout<1>(layout<0>(scale_factor_layout))  M_align_128 / 128 -- dynamic
-    // shape dynamic stride layout<0>(layout<1>(scale_factor_layout))  _4:_1 --
-    // static layout<1>(layout<1>(scale_factor_layout))  (K / 32) / 4 : _512 --
-    // dynamic shape static stride
-
-    // Reshape to zipped layout for 1D indexing
-    auto zipped_scale_factor_layout = make_layout(
-        make_layout(layout<0>(layout<0>(scale_factor_layout)),
-                    layout<0>(layout<1>(scale_factor_layout))),
-        make_layout(
-            layout<1>(layout<0>(scale_factor_layout)),
-            layout<1>(layout<1>(
-                scale_factor_layout))));  // (((_32,_4),_4),(M_align_128 /
-                                          // 128,(K / 32) /
-                                          // 4)):(((_16,_4),_1),(?,_512))
-
-    auto scale_factor_tensor =
-        make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)),
-                    zipped_scale_factor_layout);
-
-    // Used for cases where M is not divisible by 128 (most scenarios).
-    auto input_shape = shape(input_tensor);  // (M, K):(K, 1)
-    auto identity_tensor = make_identity_tensor(input_shape);
-    auto predict_tensor = cute::lazy::transform(
-        identity_tensor, [&](auto c) { return elem_less(c, input_shape); });
-
-    // (_128, _128)
-    auto tiler = make_shape(Int<BLOCK_M>{}, Int<BLOCK_K>{});
-
-    auto tiled_input_tensor = zipped_divide(
-        input_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
-    auto tiled_quant_output_tensor =
-        zipped_divide(quant_output_tensor,
-                      tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
-    auto tiled_predict_tensor = zipped_divide(
-        predict_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
-
-    auto total_tiles =
-        size<1>(tiled_input_tensor);  // cdiv(M, 128) * cdiv(K, 128)
-    decltype(total_tiles) blk_offset = blockIdx.x;
-    while (blk_offset < total_tiles) {
-      auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
-      auto current_quant_output_tile =
-          tensor<0>(tiled_quant_output_tensor(_, blk_offset));
-      auto current_predict_tile =
-          tensor<0>(tiled_predict_tensor(_, blk_offset));
-      auto current_scale_factor_tile =
-          tensor<0>(scale_factor_tensor(_, blk_offset));
-
-      mxfp8_experts_quant_tile<
-          decltype(current_input_tile), decltype(current_predict_tile),
-          decltype(current_quant_output_tile), decltype(scale_factor_shared),
-          decltype(current_scale_factor_tile), TiledCopyG2R, TiledCopyR2G,
-          TiledCopyR2S>(current_input_tile, current_predict_tile,
-                        current_quant_output_tile, scale_factor_shared,
-                        current_scale_factor_tile, m, tiled_copy_g2r,
-                        tiled_copy_r2g, tiled_copy_r2s);
-      blk_offset += gridDim.x;
-    }
-  }
-#endif
-}
-
-template <typename T_IN>
-void launch_mxfp8_experts_quant(const torch::stable::Tensor& input,
-                                const torch::stable::Tensor& problem_sizes,
-                                const torch::stable::Tensor& expert_offsets,
-                                const torch::stable::Tensor& blockscale_offsets,
-                                torch::stable::Tensor& quant_output,
-                                torch::stable::Tensor& scale_factor) {
-  ThrLayout thr_layout{};
-  ValLayout val_layout{};
-  SfR2SThrLayout r2s_thr_layout{};
-  SfR2SValLayout r2s_val_layout{};
-
-  using CopyOpG2R =
-      UniversalCopy<cutlass::AlignedArray<T_IN, size(val_layout)>>;
-  using CopyAtomG2R = cute::Copy_Atom<CopyOpG2R, T_IN>;
-  auto tiled_copy_g2r = cute::make_tiled_copy(
-      CopyAtomG2R{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
-
-  using CopyOpR2G = UniversalCopy<
-      cutlass::AlignedArray<cutlass::float_e4m3_t, size(val_layout)>>;
-  using CopyAtomR2G = cute::Copy_Atom<CopyOpR2G, cutlass::float_e4m3_t>;
-  auto tiled_copy_r2g = cute::make_tiled_copy(
-      CopyAtomR2G{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
-
-  using CopyOpR2S =
-      UniversalCopy<cutlass::AlignedArray<uint8_t, size(r2s_val_layout)>>;
-  using CopyAtomR2S = cute::Copy_Atom<CopyOpR2S, uint8_t>;
-  auto tiled_copy_r2s = cute::make_tiled_copy(
-      CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)
-
-  int max_active_blocks_per_sm = -1;
-  STD_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(
-      &max_active_blocks_per_sm,
-      mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
-                                 decltype(tiled_copy_r2g),
-                                 decltype(tiled_copy_r2s)>,
-      THREAD_BLOCK_SIZE, 0));
-
-  dim3 grid(get_device_prop()->multiProcessorCount * max_active_blocks_per_sm,
-            1, 1);
-  dim3 block(THREAD_BLOCK_SIZE, 1, 1);
-  int num_experts = static_cast<int>(problem_sizes.size(0));
-  auto stream = get_current_cuda_stream(input.get_device_index());
-  mxfp8_experts_quant_kernel<T_IN, decltype(tiled_copy_g2r),
-                             decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
-      <<<grid, block, 0, stream>>>(
-          reinterpret_cast<const T_IN*>(input.data_ptr()),
-          reinterpret_cast<const int*>(problem_sizes.data_ptr()),
-          reinterpret_cast<const int*>(expert_offsets.data_ptr()),
-          reinterpret_cast<const int*>(blockscale_offsets.data_ptr()),
-          reinterpret_cast<cutlass::float_e4m3_t*>(quant_output.data_ptr()),
-          reinterpret_cast<uint8_t*>(scale_factor.data_ptr()), num_experts,
-          tiled_copy_g2r, tiled_copy_r2g, tiled_copy_r2s);
-}
-
-}  // namespace expert_specialization

csrc/libtorch_stable/ops.h

@@ -3,6 +3,9 @@
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
 
+#include <optional>
+#include <string>
+
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_q,
                                torch::stable::Tensor& output_s,
@@ -185,11 +188,12 @@ torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
 
 // Layernorm kernels (shared CUDA/ROCm)
 void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
-              torch::stable::Tensor& weight, double epsilon);
+              std::optional<torch::stable::Tensor> weight, double epsilon);
 
 void fused_add_rms_norm(torch::stable::Tensor& input,
                         torch::stable::Tensor& residual,
-                        torch::stable::Tensor& weight, double epsilon);
+                        std::optional<torch::stable::Tensor> weight,
+                        double epsilon);
 
 // Layernorm-quant kernels (shared CUDA/ROCm)
 void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
@@ -281,6 +285,25 @@ minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
                          int64_t const nranks, double const eps);
 #endif
 
+// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
+// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
+// the index branch and scatters k/v/index_k into their paged caches.
+void fused_minimax_m3_qknorm_rope_kv_insert(
+    torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
+    torch::stable::Tensor const& k_norm_weight,
+    torch::stable::Tensor const& cos_sin_cache,
+    torch::stable::Tensor const& positions, int64_t num_heads,
+    int64_t num_kv_heads, int64_t rotary_dim, double eps,
+    std::optional<torch::stable::Tensor> index_q_norm_weight,
+    std::optional<torch::stable::Tensor> index_k_norm_weight,
+    int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
+    std::optional<torch::stable::Tensor> index_slot_mapping,
+    std::optional<torch::stable::Tensor> kv_cache,
+    std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
+    std::optional<torch::stable::Tensor> q_out,
+    std::optional<torch::stable::Tensor> index_q_out,
+    const std::string& kv_cache_dtype);
+
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -346,7 +369,8 @@ void free_shared_buffer(int64_t buffer);
 // Activation kernels (shared CUDA/ROCm)
 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
-                        torch::stable::Tensor& input, double limit);
+                        torch::stable::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

@@ -237,21 +237,30 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
   // Get the final absolute maximum values.
   float vecMax = float(__hmax(localMax.x, localMax.y));
 
-  // Get the SF (max value of the vector / max value of e2m1).
-  // maximum value of e2m1 = 6.0.
-  // TODO: use half as compute data type.
-  float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
   // 8 bits representation of the SF.
+  float SFValue;
   uint8_t fp8SFVal;
-  // Write the SF to global memory (STG.8).
+
   if constexpr (UE8M0_SF) {
-    // Extract the 8 exponent bits from float32.
-    // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
-    uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
-    fp8SFVal = tmp & 0xff;
-    // Convert back to fp32.
-    reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
+    // OCP MX spec E8M0 scale computation (MXFP4 path):
+    // scale_exp = biased_exponent(round_up(vecMax)) - 2
+    //   -2 because max E2M1 value is 6.0 ≈ 2^2.58; we use 2^2=4 as the
+    //   safe divisor so that max_val / scale <= 6.0 for values near 2^n.
+    uint32_t max_bits = __float_as_uint(vecMax);
+    // Add rounding bias at mantissa bit 21 (equivalent to bf16 val_to_add=32
+    // at bit 5). Threshold: values with mantissa >= 0.75 (i.e. >= 1.75*2^n)
+    // round up to the next power of 2.
+    uint32_t rounded_bits = (max_bits + (1u << 21)) & 0xFF800000u;
+    uint32_t biased_exp = (rounded_bits >> 23) & 0xFFu;
+    uint32_t scale_exp = (biased_exp > 2u) ? (biased_exp - 2u) : 0u;
+    scale_exp = min(scale_exp, 254u);
+    fp8SFVal = static_cast<uint8_t>(scale_exp);
+    // Reconstruct scale as float32: scale = 2^(scale_exp - 127)
+    uint32_t sf_bits = scale_exp << 23;
+    SFValue = __uint_as_float(sf_bits);
   } else {
+    // NVFP4 path: scale = max / 6.0, stored as E4M3.
+    SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
     // Here SFValue is always positive, so E4M3 is the same as UE4M3.
     __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
     reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
@@ -262,13 +271,21 @@ __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4(
   // Write the SF to global memory (STG.8).
   if (SFout) *SFout = fp8SFVal;
 
-  // Get the output scale.
-  // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
-  //                       reciprocal(SFScaleVal))
-  float outputScale =
-      SFValue != 0.0f ? reciprocal_approximate_ftz(
+  // Get the output scale (= 1 / SFValue for the MXFP4/UE8M0 path where
+  // SFScaleVal=1).  Use exact division for UE8M0 to ensure bit-exact scaling
+  // that matches the reference QDQ implementation (dividing by a power-of-2
+  // scale is exact in IEEE 754).
+  float outputScale;
+  if constexpr (UE8M0_SF) {
+    // SFValue is always a power of 2 for UE8M0, so 1/SFValue is exact.
+    outputScale = SFValue != 0.0f ? (1.0f / SFValue) : 0.0f;
+  } else {
+    // NVFP4 path: use fast approximate reciprocal (original behavior).
+    outputScale = SFValue != 0.0f
+                      ? reciprocal_approximate_ftz(
                             SFValue * reciprocal_approximate_ftz(SFScaleVal))
                       : 0.0f;
+  }
 
   // Convert the input to float.
   float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];

csrc/libtorch_stable/quantization/vectorization_utils.cuh

@@ -24,13 +24,21 @@ __device__ inline void vectorize_with_alignment(
     ScaOp&& scalar_op) {  // InT -> OutT
   static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
                 "VEC_SIZE must be a positive power-of-two");
-  constexpr int WIDTH = VEC_SIZE * sizeof(InT);  // eg: 64 B
+  constexpr int WIDTH = VEC_SIZE * sizeof(InT);       // eg: 16 B
+  constexpr int OUT_WIDTH = VEC_SIZE * sizeof(OutT);  // eg: 16 B
   uintptr_t addr = reinterpret_cast<uintptr_t>(in);
+  uintptr_t out_addr = reinterpret_cast<uintptr_t>(out);
 
-  // fast path when the whole region is already aligned
-  // Note: currently the output is guaranteed to be same as the input, so we
-  // don't check it here, comments here just for future reference.
-  bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
+  // fast path when input and output are both fully aligned. The vector
+  // load/store below go through vec_n_t<T, VEC_SIZE>, declared
+  // __align__(VEC_SIZE * sizeof(T)), so each side must be aligned to its
+  // own vector width. out is NOT generally co-aligned with in: e.g.
+  // reshape_and_cache_flash writes KV-cache rows whose byte offset is a
+  // multiple of head_size, which for head sizes that are not a multiple
+  // of VEC_SIZE puts some rows off the vector-width boundary.
+  bool can_vec = ((addr & (WIDTH - 1)) == 0) &&
+                 ((out_addr & (OUT_WIDTH - 1)) == 0) &&
+                 ((len & (VEC_SIZE - 1)) == 0);
   if (can_vec) {
     int num_vec = len / VEC_SIZE;
 
@@ -55,6 +63,16 @@ __device__ inline void vectorize_with_alignment(
   prefix_elems /= sizeof(InT);
   prefix_elems = min(prefix_elems, len);  // 0 ≤ prefix < 16
 
+  // the prefix below aligns in; if that does not also align out (their
+  // addresses differ modulo the vector width), vectorizing is impossible
+  // and the whole copy must stay scalar.
+  if (((out_addr + prefix_elems * sizeof(OutT)) & (OUT_WIDTH - 1)) != 0) {
+    for (int i = tid; i < len; i += stride) {
+      scalar_op(out[i], in[i]);
+    }
+    return;
+  }
+
   // 1. prefill the when it is unsafe to vectorize
   for (int i = tid; i < prefix_elems; i += stride) {
     scalar_op(out[i], in[i]);

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -304,9 +304,17 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
   const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
   const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
 
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.wait;");
+#endif
+
   if (mn_idx >= tma_aligned_mn) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+    asm volatile("griddepcontrol.launch_dependents;");
+#endif
     return;
   }
+
   const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
 
   // Load 16 input elements (32 B) into registers as two adjacent uint4
@@ -417,6 +425,10 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
       static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
       sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
   *reinterpret_cast<uint4*>(group_output) = packed_out;
+
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
 }
 
 // Public entry point: register-resident packed quant kernel.
@@ -495,23 +507,54 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 
   auto dst_type = output_q.scalar_type();
 
-#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
-  do {                                                                       \
-    dim3 grid(static_cast<unsigned int>(blocks_x),                           \
-              static_cast<unsigned int>(blocks_y));                          \
-    dim3 block(num_threads);                                                 \
-    per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
-                                                      RY>                    \
-        <<<grid, block, 0, stream>>>(                                        \
-            static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),    \
-            reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),     \
-            static_cast<int>(padded_groups_per_row),                         \
-            static_cast<int>(groups_per_row), static_cast<int>(mn),          \
-            static_cast<int>(output_q_mn_extent),                            \
-            static_cast<int>(tma_aligned_mn), num_scale_elems,               \
-            static_cast<float>(eps), static_cast<float>(min_8bit),           \
-            static_cast<float>(max_8bit));                                   \
-  } while (0)
+// PDL (Programmatic Dependent Launch) is NVIDIA-only; ROCm/HIP has no
+// equivalent launch attribute, so fall back to a classic launch there.
+#ifndef USE_ROCM
+  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
+    do {                                                                       \
+      cudaLaunchConfig_t config = {};                                          \
+      config.gridDim = dim3(static_cast<unsigned int>(blocks_x),               \
+                            static_cast<unsigned int>(blocks_y));              \
+      config.blockDim = dim3(num_threads);                                     \
+      config.dynamicSmemBytes = 0;                                             \
+      config.stream = stream;                                                  \
+      cudaLaunchAttribute attrs[1];                                            \
+      attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;        \
+      attrs[0].val.programmaticStreamSerializationAllowed = 1;                 \
+      config.numAttrs = 1;                                                     \
+      config.attrs = attrs;                                                    \
+      cudaLaunchKernelEx(                                                      \
+          &config,                                                             \
+          per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, \
+                                                            KX, RY>,           \
+          static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),        \
+          reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),         \
+          static_cast<int>(padded_groups_per_row),                             \
+          static_cast<int>(groups_per_row), static_cast<int>(mn),              \
+          static_cast<int>(output_q_mn_extent),                                \
+          static_cast<int>(tma_aligned_mn), num_scale_elems,                   \
+          static_cast<float>(eps), static_cast<float>(min_8bit),               \
+          static_cast<float>(max_8bit));                                       \
+    } while (0)
+#else
+  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
+    do {                                                                       \
+      dim3 grid(static_cast<unsigned int>(blocks_x),                           \
+                static_cast<unsigned int>(blocks_y));                          \
+      dim3 block(num_threads);                                                 \
+      per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
+                                                        RY>                    \
+          <<<grid, block, 0, stream>>>(                                        \
+              static_cast<const T*>(input.data_ptr()), output_q.data_ptr(),    \
+              reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()),     \
+              static_cast<int>(padded_groups_per_row),                         \
+              static_cast<int>(groups_per_row), static_cast<int>(mn),          \
+              static_cast<int>(output_q_mn_extent),                            \
+              static_cast<int>(tma_aligned_mn), num_scale_elems,               \
+              static_cast<float>(eps), static_cast<float>(min_8bit),           \
+              static_cast<float>(max_8bit));                                   \
+    } while (0)
+#endif
 
 #define LAUNCH_REG_KERNEL(T, DST_DTYPE)                    \
   do {                                                     \

csrc/libtorch_stable/torch_bindings.cpp

@@ -308,22 +308,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
       "Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
 
-  // Expert-specialization mxfp8 blockscaled grouped quantization (SM100+).
-  ops.def(
-      "mxfp8_experts_quant("
-      " Tensor input, Tensor problem_sizes, Tensor expert_offsets,"
-      " Tensor blockscale_offsets, Tensor! quant_output, Tensor! scale_factor)"
-      " -> ()");
-  // conditionally compiled so impl registration is in source file
-
-  // Expert-specialization mxfp8 blockscaled grouped GEMM (SM100+).
-  ops.def(
-      "cutlass_mxfp8_grouped_mm("
-      " Tensor a, Tensor b, Tensor sfa, Tensor sfb, Tensor! out,"
-      " Tensor problem_sizes, Tensor expert_offsets, Tensor blockscale_offsets)"
-      " -> ()");
-  // conditionally compiled so impl registration is in source file
-
   // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
   // conditionally compiled so impl registration is in source file
   ops.def(
@@ -369,12 +353,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
 
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(
-      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
+      "rms_norm(Tensor! result, Tensor input, Tensor? weight, float epsilon) "
+      "-> "
       "()");
 
   // In-place fused Add and RMS Normalization.
   ops.def(
-      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
+      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
       "float epsilon) -> ()");
 
   // Layernorm-quant
@@ -461,6 +446,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "float eps) -> (Tensor, Tensor)");
 #endif
 
+  // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
+  ops.def(
+      "fused_minimax_m3_qknorm_rope_kv_insert("
+      "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
+      "Tensor cos_sin_cache, Tensor positions, int num_heads, "
+      "int num_kv_heads, int rotary_dim, float eps, "
+      "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
+      "int num_index_heads, "
+      "Tensor? slot_mapping, Tensor? index_slot_mapping, "
+      "Tensor!? kv_cache, Tensor!? index_cache, "
+      "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
+      "str kv_cache_dtype) -> ()");
+
   // Apply repetition penalties to logits in-place.
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -488,9 +486,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
 
   // SwiGLU activation with input clamping.
+  // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
+  // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
   ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
-      "-> ()");
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
+      "float alpha=1.0, float beta=0.0) -> ()");
 
   // Activation function used in GeGLU with `none` approximation.
   ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -679,6 +679,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("minimax_allreduce_rms", TORCH_BOX(&minimax_allreduce_rms));
   ops.impl("minimax_allreduce_rms_qk", TORCH_BOX(&minimax_allreduce_rms_qk));
 #endif
+  ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
+           TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
 
   // Sampler kernels (shared CUDA/ROCm)
   ops.impl("apply_repetition_penalties_",

csrc/ops.h

@@ -34,11 +34,11 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.
-void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
-              double epsilon);
+void rms_norm(torch::Tensor& out, torch::Tensor& input,
+              std::optional<torch::Tensor> weight, double epsilon);
 
 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
-                        torch::Tensor& weight, double epsilon);
+                        std::optional<torch::Tensor> weight, double epsilon);
 
 // rotary_embedding also exist in csrc/libtorch_stable/ops.h (torch::stable
 // ABI for CUDA). It remains here because the CPU build still uses these
@@ -50,7 +50,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
 
 void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
-void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
+void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
+                        double alpha = 1.0, double beta = 0.0);
 
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);

docker/Dockerfile.xpu

@@ -131,8 +131,8 @@ CMD ["/bin/bash"]
 # never included in the final runtime image (mirrors ROCm's build_rixl stage).
 FROM vllm-base AS ucx-nixl-build
 
-ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
-ARG NIXL_VERSION=0.7.0
+ARG UCX_VERSION=v1.21.0-rc2
+ARG NIXL_VERSION=0.10.1
 
 # Build-time only: compiler, autotools, and verbs dev headers
 RUN apt-get update -y && apt-get install -y --no-install-recommends \
@@ -167,8 +167,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 
 FROM vllm-base AS vllm-openai
 
-ARG UCX_VERSION=e5d98879705239d254ede40b4a52891850cb5349
-ARG NIXL_VERSION=0.7.0
+ARG NIXL_VERSION=0.10.1
 
 # Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
 # No compiler or autotools are installed in this stage.
@@ -192,7 +191,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     ibverbs-providers \
     librdmacm1t64 \
     && rm -rf /var/lib/apt/lists/* \
-    && uv pip install --no-deps /tmp/nixl_wheels/nixl-*.whl \
+    && uv pip install --no-deps /tmp/nixl_wheels/nixl*.whl \
+    && uv pip install nixl==${NIXL_VERSION} \
     && rm -rf /tmp/nixl_wheels
 
 RUN --mount=type=cache,target=/root/.cache/uv \

docs/design/attention_backends.md

@@ -170,8 +170,8 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -188,6 +188,18 @@ Priority is **1 = highest** (tried first).
 >
 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
 
+## MiniMax M3 Sparse Attention Backends
+
+Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
+layers. It is wired in directly by the model and is not part of the
+automatic priority lists above. A lightning indexer scores KV blocks, the
+top-k blocks (plus fixed init/local blocks) are selected, and attention
+attends only to those blocks; index keys live in a separate side cache.
+
+| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
+| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
+| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
+
 ## MLA (Multi-head Latent Attention) Backends
 
 MLA uses separate backends for prefill and decode phases.
@@ -202,9 +214,9 @@ hardware and configuration.
 | Backend | Description | Dtypes | Compute Cap. | Notes |
 | ------- | ----------- | ------ | ------------ | ----- |
 | `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
-| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
-| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
-| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
+| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
+| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
+| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
 
 > **‡** Automatic selection tries FlashAttention first. On Blackwell
 > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
@@ -240,6 +252,6 @@ default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
 
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
-| `FLASHINFER_MLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
-| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
+| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |

docs/design/cuda_graphs_multimodal.md

@@ -2,6 +2,8 @@
 
 The [CUDA Graphs](cuda_graphs.md) infrastructure in vLLM primarily targets the **decoder** (language model) forward pass. vLLM also supports capturing the **encoder** (vision transformer) forward pass as CUDA Graphs, independently from the decoder. This is based on <https://github.com/vllm-project/vllm/pull/35963>.
 
+For two-tower vision encoders (e.g., DeepSeek-OCR's SAM + CLIP with dynamic tiling), a **dual-path graph** mode captures two independent sets of CUDA graphs — one for the global image path and one for the local patch path — enabling independent budget selection and partial eager fallback per path. This is based on <https://github.com/vllm-project/vllm/pull/43586>.
+
 !!! note
     Encoder CUDA Graphs are orthogonal to decoder CUDA Graphs — both can be enabled simultaneously. Encoder graphs capture the vision encoder execution (e.g., ViT in Qwen3-VL), while decoder graphs capture the language model execution as described in the [CUDA Graphs design document](cuda_graphs.md).
 
@@ -11,6 +13,8 @@ Vision encoder inference incurs CUDA kernel launch overhead on the host side. Th
 
 Encoder CUDA Graphs eliminate this overhead by pre-capturing the full encoder forward pass at multiple token budget levels during model initialization, then replaying the appropriate graph at runtime.
 
+For two-tower vision encoders such as DeepSeek-OCR (SAM + CLIP with dynamic tiling), the global image path and local patch path have independent token profiles (272 tokens per global image vs. 100 tokens per local patch). Capturing a single monolithic graph for both paths would significantly reduce packing efficiency. The dual-path graph mode captures each path as a separate set of budgets, allowing the manager to pack and replay each path independently.
+
 ## Design
 
 The encoder CUDA Graph system uses a **budget-based capture/replay** strategy, managed by [EncoderCudaGraphManager][vllm.v1.worker.encoder_cudagraph.EncoderCudaGraphManager]. The system contains the following core components:
@@ -37,10 +41,14 @@ class BudgetGraphMetadata:
 
 Budgets are auto-generated as power-of-2 levels from a model-provided range via `get_encoder_cudagraph_budget_range()`, with the maximum budget always included even if it does not fall on a power-of-2 boundary. Budgets can also be explicitly specified by the user via `encoder_cudagraph_token_budgets` in `CompilationConfig`.
 
+When `EncoderCudaGraphConfig.enable_dual_path_graph` is `True`, the manager generates two independent budget lists — `global_token_budgets` (multiples of `global_token_per_image`) and `local_token_budgets` (multiples of `local_token_per_patch`) — and stores captured graphs under `budget_graphs["global"]` and `budget_graphs["local"]` respectively.
+
 ### Greedy bin-packing at runtime
 
 When a batch of images arrives, the manager sorts images by output token count (smallest first) and greedily packs as many images as possible into each sub-batch while staying within the **largest** token budget and the maximum batch size. Once a sub-batch is finalized (the next image would overflow either constraint), the manager finds the **smallest** budget that fits the sub-batch's total tokens and replays the corresponding CUDA Graph. This repeats until the batch is exhausted. Images that exceed all budgets fall back to eager execution.
 
+For dual-path models, the manager routes to `_execute_local_dual_path()`, which constrains both global and local token budgets simultaneously during packing (see [Dual-Path graph capture](#dual-path-graph-capture)).
+
 For each graph replay:
 
 1. Call `prepare_encoder_cudagraph_replay_buffers()` to compute buffer values (including `pixel_values` and precomputed metadata) from actual batch inputs.
@@ -48,6 +56,42 @@ For each graph replay:
 3. Replay the CUDA Graph.
 4. Clone outputs from `output_buffer` (cloning is necessary since the buffer is reused across replays).
 
+### Dual-Path graph capture
+
+For two-tower vision encoders (e.g., DeepSeek-OCR), the `EncoderCudaGraphConfig` sets `enable_dual_path_graph=True` and provides `global_token_per_image` / `local_token_per_patch`. The manager captures two independent sets of CUDA graphs — one for the **global** image path and one for the **local** patch path — stored under `budget_graphs["global"]` and `budget_graphs["local"]` respectively.
+
+**Budget generation.** Two separate budget lists are generated:
+
+* `global_token_budgets` — power-of-2 multiples of `global_token_per_image` (e.g., `[272, 544, 1088, 2176, 4352, 8704, 13824]` for DeepSeek-OCR).
+* `local_token_budgets` — power-of-2 multiples of `local_token_per_patch` (e.g., `[0, 100, 200, 400, 800, 1600, 3200, 6400, 12800]` for DeepSeek-OCR). A budget of `0` is always included to handle images with no local patches (images ≤ 640×640 that produce only global features).
+
+Both lists are capped at the same `max_budget`.
+
+**Dual-path greedy packing.** Each `EncoderItemSpec` provides both `global_output_tokens` (constant per image) and `local_output_tokens` (proportional to the patch count). The dual-path packing algorithm constrains both budgets simultaneously:
+
+* Sort images by total output tokens (global + local), smallest first.
+* Greedily pack images: an image is added to the current sub-batch only if both the accumulated global tokens ≤ `max_global_budget` **and** the accumulated local tokens ≤ `max_local_budget`, with the image count ≤ `max_batch_size`.
+* Once either constraint would overflow, finalize the sub-batch and find the smallest fitting budget **independently** for each path.
+* Repeat until all images are packed.
+
+**Partial graph fallback.** After packing, each sub-batch falls into one of four execution scenarios:
+
+| Global budget | Local budget | Execution |
+| :---: | :---: | --- |
+| Found | Found | Both paths use CUDA graph replay |
+| Found | `None` | Global graph replay + local path skipped (no patches) |
+| `None` | Found | Global eager fallback + local graph replay |
+| `None` | `None` | Both paths fall back to eager execution |
+
+Note that the `0`-budget graph is never actually replayed for local — it signals that local patch processing should be skipped entirely.
+
+**Buffer keys per path.** Global and local paths use different buffer keys. For DeepSeek-OCR, the global path uses `pixel_values` (full images, shape `[B, 3, 1280, 1280]`) while the local path uses `images_crop` (patches, shape `[P, 3, 1024, 1024]`). The manager iterates over each captured graph's own `input_buffers.keys()` rather than a shared `buffer_keys` list, so both paths can use different buffers.
+
+**Post-processing.** The `postprocess_encoder_output` method receives a `local_output` parameter (a tensor or `None`) containing the local-path encoder output. The model is responsible for assembling global and local features into the final per-image embedding. For DeepSeek-OCR, this means reshaping the global output into `[B, 272, n_embed]`, the local output into `[P, 100, n_embed]`, assembling patch grids with newline tokens, and concatenating `[patches_grid, global, view_separator]` for each image.
+
+!!! note
+    The dual-path design enables partial CUDA graph coverage — one path can hit while the other falls back to eager. This avoids wasted compute on zero-padded patch buffers for untiled images and avoids graph invalidation caused by variable `crop_shape` per image.
+
 ### Data-parallel support
 
 When `mm_encoder_tp_mode="data"`, the manager distributes images across TP ranks using load-balanced assignment via `get_load_balance_assignment`, executes locally on each rank, then gathers results back in the original order via `tensor_model_parallel_all_gather`.
@@ -67,29 +111,31 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
 
 * `get_encoder_cudagraph_config()` — returns static configuration (supported modalities, buffer keys, output hidden size, padding logics, max frames per video).
 * `get_encoder_cudagraph_budget_range(vllm_config)` — returns `(min_budget, max_budget)` for auto-inference of token budgets.
-* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size and output token count. Replaces the former three separate methods (`get_num_items`, `get_per_item_output_tokens`, `get_per_item_input_sizes`).
+* `get_encoder_cudagraph_item_specs(mm_kwargs)` — returns `list[EncoderItemSpec]` describing each item with its input size, total output token count (`output_tokens`), and optionally per-path token counts (`global_output_tokens`, `local_output_tokens`) for dual-path models.
 * `select_encoder_cudagraph_items(mm_kwargs, indices)` — extracts a sub-batch of items by index, used during greedy packing and DP sharding.
-* `prepare_encoder_cudagraph_capture_inputs(...)` — creates dummy inputs for graph capture. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph.
-* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch)` — computes buffer values from actual batch inputs. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match `buffer_keys` in the config.
-* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor])` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `pixel_values` tensor is included in `inputs` alongside metadata buffers.
-* `encoder_eager_forward(mm_kwargs)` — fallback eager forward when no graph fits.
-* `postprocess_encoder_output(...)` — post-process encoder output, delegates to `scatter_output_slices` by default.
+* `prepare_encoder_cudagraph_capture_inputs(..., path="default")` — creates dummy inputs for graph capture. The `path` parameter (`"global"` or `"local"`) tells the model which path to generate dummy inputs for. Returns `EncoderCudaGraphCaptureInputs` with a single `values: dict[str, torch.Tensor]` that contains all buffers to be recorded into the graph.
+* `prepare_encoder_cudagraph_replay_buffers(mm_kwargs, max_batch_size, max_frames_per_batch, path="default")` — computes buffer values from actual batch inputs. The `path` parameter selects which modality keys to extract from `mm_kwargs`. Returns `EncoderCudaGraphReplayBuffers` with a `values` dict whose keys match the captured graph's `input_buffers.keys()`.
+* `encoder_cudagraph_forward(inputs: dict[str, torch.Tensor], path="default")` — forward pass accepting only fixed-shaped input tensors (the captured `values` dict). Called during both capture and replay. The `path` parameter dispatches to the correct encoder sub-module (e.g., global vs. local path for DeepSeek-OCR).
+* `encoder_eager_forward(mm_kwargs, path="default")` — fallback eager forward when no graph fits. When `path` is `"global"` or `"local"`, runs only that encoder path without graph capture.
+* `postprocess_encoder_output(..., local_output=None)` — post-process encoder output. The `local_output` parameter receives the local-path encoder output tensor (or `None`), enabling dual-path models to assemble global and local features into the final per-image embedding.
 
 !!! note
     The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
 
 **Supported models:**
 
-| Architecture | Models | CG for Image | CG for Video |
-| ------------ | ------ | ------------ | ------------ |
-| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | - |
-| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ |
-| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ |
-| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
-| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
-| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
-| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |
-| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ |
+| Architecture | Models | CG for Image | CG for Video | Dual-Path Graph |
+| ------------ | ------ | ------------ | ------------ | --------------- |
+| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
+| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
+| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
+| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
+| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
+| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
+| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |
+| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ |
+| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ | ❌︎ |
+| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ❌︎ |
 
 !!! note
     Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
@@ -104,6 +150,8 @@ Three fields in `CompilationConfig` control encoder CUDA Graphs:
 * `encoder_cudagraph_max_vision_items_per_batch` (`int`, default `0`) — maximum number of images/videos per batch during capture. If 0 (default), auto-inferred as `max_budget // min_budget`.
 * `encoder_cudagraph_max_frames_per_batch` (`int`, default `None`) — maximum number of video frames per batch during capture. If `None` (default), auto-inferred as `encoder_cudagraph_max_vision_items_per_batch * max_frames_per_video` (`max_frames_per_video` is a model-specific value from `EncoderCudaGraphConfig`, computed by `get_max_frames_per_video()` on the model). If we limit the video count per prompt to `0`, it will also be set to `0` (i.e., fall back to image-only mode).
 
+Dual-path mode is configured at the model level via `EncoderCudaGraphConfig` fields (`enable_dual_path_graph`, `global_token_per_image`, `local_token_per_patch`) — no additional user configuration is required. The manager automatically generates separate budget lists and routes to dual-path execution when the model opts in.
+
 ## Usage guide
 
 ### Image inference

docs/features/kv_offloading_usage.md

@@ -127,6 +127,29 @@ PYTHONHASHSEED=0 vllm serve ...
 - FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
 - Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
 
+## Per-Request Selective Offload
+
+Individual requests can cap how many of their tokens are eligible for offload by setting `max_offload_tokens` in the request's `kv_transfer_params`. Only the first `max_offload_tokens` tokens of the request are offloaded; blocks beyond that point are skipped on the store path. This is useful when a known prefix (e.g., a system prompt or shared context) is worth caching but later request-specific tokens are not.
+
+| Key | Type | Notes |
+| --- | --- | --- |
+| `max_offload_tokens` | non-negative `int` | Upper bound on tokens to offload for this request. `0` disables offload for the request entirely; omit the key (or set to `None`) for no cap. Non-`int`, negative, or `bool` values are rejected with a warning and treated as no cap. |
+
+!!! note
+    `max_offload_tokens` is experimental and subject to change.
+
+Example (OpenAI-compatible completions request):
+
+```json
+{
+  "model": "<model>",
+  "prompt": "...",
+  "kv_transfer_params": {
+    "max_offload_tokens": 1024
+  }
+}
+```
+
 ## Further Reading
 
 - [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput).

docs/features/mooncake_store_connector_usage.md

@@ -205,6 +205,7 @@ the vLLM JSON config.
 - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
 - `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
+- `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
 
 ## Notes
 

docs/features/quantization/auto_awq.md

@@ -49,7 +49,7 @@ To run an AWQ model with vLLM, you can use [TheBloke/Llama-2-7b-Chat-AWQ](https:
 ```bash
 python examples/deployment/llm_engine_example.py \
     --model TheBloke/Llama-2-7b-Chat-AWQ \
-    --quantization awq
+    --quantization auto_awq
 ```
 
 AWQ models are also supported directly through the LLM entrypoint:
@@ -70,7 +70,7 @@ AWQ models are also supported directly through the LLM entrypoint:
     sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
     # Create an LLM.
-    llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="AWQ")
+    llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="auto_awq")
     # Generate texts from the prompts. The output is a list of RequestOutput objects
     # that contain the prompt, generated text, and other information.
     outputs = llm.generate(prompts, sampling_params)

docs/features/tool_calling.md

@@ -109,18 +109,18 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t
 
 ## Constrained Decoding Behavior
 
-Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode:
+Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode and the per-tool `strict` field:
 
 | `tool_choice` value | Schema-constrained decoding | Behavior |
 | --- | --- | --- |
 | Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. |
 | `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. |
-| `"auto"` | Depends on the parser | Model-specific structural-tag parsers can constrain tool-call arguments with structured outputs. Other parsers generate freely and extract tool calls from raw text. |
+| `"auto"` | Only when `strict: true` is set on at least one tool | Structural-tag parsers constrain tool-call arguments when a tool opts in with `strict: true`. Without it, the model generates freely and tool calls are extracted from raw text. |
 | `"none"` | N/A | No tool calls are produced. |
 
 ### Strict Mode
 
-Strict tool calling makes function-call arguments adhere to the function schema instead of relying only on best-effort parsing. vLLM implements strict tool calling for structural-tag based tool parsers by using the structured outputs backend under the hood.
+For `tool_choice="required"` or named function calling, structural-tag constraints are always applied regardless of the `strict` field. For `tool_choice="auto"`, setting `strict: true` on at least one tool opts in to structural-tag constraints; without it, the model generates freely and tool calls are extracted from raw text. The `strict` field is supported across all three API surfaces: Chat Completion, Responses, and Anthropic Messages.
 
 For best compatibility with strict schema enforcement, define tool parameter schemas in the OpenAI strict-schema style:
 
@@ -128,16 +128,12 @@ For best compatibility with strict schema enforcement, define tool parameter sch
 * Mark all fields in `properties` as required.
 * Represent optional fields by allowing `null`, for example `{"type": ["string", "null"]}`.
 
-vLLM controls structural-tag strict tool calling with the `VLLM_ENFORCE_STRICT_TOOL_CALLING` environment variable. It defaults to `true`.
+vLLM also provides a global toggle via the `VLLM_ENFORCE_STRICT_TOOL_CALLING` environment variable (defaults to `true`). When set to `false`, vLLM does not attach structural tags for tool calling regardless of the per-tool `strict` field. This environment variable only affects structural-tag based tool calling; it does not change schema-derived structured outputs used by named function calling or `tool_choice="required"`.
 
 ```bash
 VLLM_ENFORCE_STRICT_TOOL_CALLING=false vllm serve ...
 ```
 
-When this variable is `true`, structural-tag based tool parsers attach a structural tag to the request, so the structured outputs backend can constrain the model-specific tool-call format and function-call arguments. When it is `false`, vLLM does not attach structural tags for tool calling. In that case, `tool_choice="auto"` falls back to best-effort parser extraction from the raw model output, and no structural-tag constraint is applied.
-
-This environment variable only affects structural-tag based tool calling. It does not change schema-derived structured outputs used by named function calling or `tool_choice="required"`.
-
 ## Automatic Function Calling
 
 To enable this feature, you should set the following flags:
@@ -156,7 +152,7 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso
 If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template!
 
 !!! note
-    With `tool_choice="auto"`, schema-level constraint depends on the selected parser and `VLLM_ENFORCE_STRICT_TOOL_CALLING`. Structural-tag parsers can enforce tool-call constraints when it is `true`; when it is `false`, or when the selected parser has no structural-tag support, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema.
+    With `tool_choice="auto"`, schema-level constraint requires both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`. When these conditions are met and the selected parser supports structural tags, vLLM constrains tool-call arguments. Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema.
 
 ### Hermes Models (`hermes`)
 

docs/getting_started/installation/cpu.md

@@ -142,6 +142,10 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
 === "IBM Z (S390X)"
     --8<-- "docs/getting_started/installation/cpu.s390x.inc.md:build-image-from-source"
 
+## AMD Zen optimizations {#amd-zen-optimizations}
+
+--8<-- "docs/getting_started/installation/cpu.x86.inc.md:amd-zen-optimizations"
+
 ## Related runtime environment variables
 
 - `VLLM_CPU_KVCACHE_SPACE`: specify the KV Cache size (e.g, `VLLM_CPU_KVCACHE_SPACE=40` means 40 GiB space for KV cache), larger setting will allow vLLM to run more requests in parallel. This parameter should be set based on the hardware configuration and memory management pattern of users. Default value is `0`.
@@ -149,12 +153,14 @@ VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cpu VLLM_TARGET_DEVICE=cpu
 - `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`.
 - `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence.
 - `VLLM_CPU_SGL_KERNEL` (x86 only, Experimental): whether to use small-batch optimized kernels for linear layer and MoE layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False).
+- `VLLM_ZENTORCH_WEIGHT_PREPACK` (AMD Zen only): when `ZenCpuPlatform` is active, eagerly prepack linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Default is `1` (enabled). See [AMD Zen optimizations](#amd-zen-optimizations).
 
 ## FAQ
 
 ### Which `dtype` should be used?
 
 - Currently, vLLM CPU uses model default settings as `dtype`. However, due to unstable float16 support in torch CPU, it is recommended to explicitly set `dtype=bfloat16` if there are any performance or accuracy problem.  
+- On AMD Zen CPUs (`ZenCpuPlatform`), `float16` is **not** supported. Only `bfloat16` and `float32` are accepted; models declared with `float16` are auto-downcast to `bfloat16` at model load time. See [AMD Zen optimizations](#amd-zen-optimizations).
 
 ### How to launch a vLLM service on CPU?
 
@@ -227,6 +233,25 @@ By providing MODEL_FILTER and DTYPE_FILTER, only commands for related model ID a
 ON_CPU=1 SERVING_JSON=serving-tests-cpu-text.json DRY_RUN=1 MODEL_FILTER=meta-llama/Llama-3.1-8B-Instruct DTYPE_FILTER=bfloat16  bash .buildkite/performance-benchmarks/scripts/run-performance-benchmarks.sh
 ```
 
+### How do I enable AMD Zen optimizations? {#how-do-i-enable-amd-zen-optimizations}
+
+On an AMD Zen 4 / Zen 5 CPU, install the CPU wheel with the `zen` extra so vLLM pulls the tested `zentorch` version for that release:
+
+```bash
+export VLLM_VERSION=$(curl -s https://api.github.com/repos/vllm-project/vllm/releases/latest | jq -r .tag_name | sed 's/^v//')
+uv pip install "vllm[zen]" --extra-index-url https://wheels.vllm.ai/${VLLM_VERSION}/cpu --index-strategy first-index --torch-backend cpu
+```
+
+vLLM auto-detects the platform and routes linear layers through ZenDNN-optimized kernels - no flag needed. To verify it is engaged, look for the platform-selection line in the server's startup logs:
+
+```bash
+vllm serve Qwen/Qwen3-0.6B 2>&1 | grep "AMD Zen CPU detected with zentorch installed"
+```
+
+For per-backend dispatch details (which kernel each linear layer was bound to), re-run with `VLLM_LOGGING_LEVEL=DEBUG` and grep for `CPU unquantized GEMM dispatch`.
+
+See [AMD Zen optimizations](#amd-zen-optimizations) for detection rules, supported dtypes, and the `VLLM_ZENTORCH_WEIGHT_PREPACK` knob.
+
 ### How to decide `VLLM_CPU_OMP_THREADS_BIND`?
 
 - Default `auto` thread-binding is recommended for most cases. Ideally, each OpenMP thread will be bound to a dedicated physical core respectively, threads of each rank will be bound to the same NUMA node respectively, and 1 CPU per rank will be reserved for other vLLM components when `world_size > 1`. If you have any performance problems or unexpected binding behaviours, please try to bind threads as following.

docs/getting_started/installation/cpu.x86.inc.md

@@ -1,4 +1,4 @@
-<!-- markdownlint-disable MD041 -->
+<!-- markdownlint-disable MD041 MD051 -->
 --8<-- [start:installation]
 
 vLLM supports basic model inferencing and serving on x86 CPU platform, with data types FP32, FP16 and BF16.
@@ -200,7 +200,19 @@ docker build -f docker/Dockerfile.cpu \
         --target vllm-openai .
 ```
 
-#### Launching the OpenAI server
+#### Building with AMD Zen optimizations
+
+For AMD Zen 4 / Zen 5 hosts (`linux/amd64` only), use the `vllm-openai-zen` target. It extends the default `vllm-openai` image and adds `zentorch` via the `vllm[zen]` extra so `ZenCpuPlatform` auto-activates at runtime:
+
+```bash
+docker build -f docker/Dockerfile.cpu \
+        --tag vllm-cpu-zen-env \
+        --target vllm-openai-zen .
+```
+
+The resulting image accepts the same arguments and environment variables as `vllm-openai` (see [Launching the OpenAI server](#launching-the-openai-server) below); no extra flag is needed to engage Zen optimizations. See [AMD Zen optimizations](cpu.md#amd-zen-optimizations) for runtime behavior and the supported-dtype caveats.
+
+#### Launching the OpenAI server {#launching-the-openai-server}
 
 ```bash
 docker run --rm \
@@ -216,5 +228,36 @@ docker run --rm \
 ```
 
 --8<-- [end:build-image-from-source]
+--8<-- [start:amd-zen-optimizations]
+
+On AMD Zen CPUs, vLLM auto-selects `ZenCpuPlatform` (a subclass of `CpuPlatform`) which dispatches linear layers through [`zentorch`](https://github.com/amd/ZenDNN-pytorch-plugin)'s ZenDNN-optimized kernels. See the FAQ entry [How do I enable AMD Zen optimizations?](#how-do-i-enable-amd-zen-optimizations) for the install command.
+
+### Detection rules
+
+`ZenCpuPlatform` is selected when **all** of the following hold:
+
+- vLLM is built for CPU
+- `/proc/cpuinfo` reports `AuthenticAMD` and `avx512`
+- `import zentorch` succeeds
+
+Otherwise, vLLM falls back to the default `CpuPlatform` (oneDNN / sgl-kernel paths).
+
+### Supported dtypes
+
+`float16` is **not** supported on `ZenCpuPlatform`. `ZenCpuPlatform.supported_dtypes` advertises only `bfloat16` and `float32`, so models declared with `torch_dtype=float16` are auto-downcast to `bfloat16` at load time with the standard `"Your device 'cpu' doesn't support torch.float16. Falling back to torch.bfloat16 for compatibility."` warning emitted from `vllm/config/model.py`.
+
+### Environment variables
+
+- `VLLM_ZENTORCH_WEIGHT_PREPACK` (default `1`): eagerly prepacks linear weights into ZenDNN's blocked layout at model load time, eliminating per-inference layout conversion overhead. Set to `0` to disable.
+
+### Docker
+
+The `vllm-openai-zen` Docker target (in `docker/Dockerfile.cpu`) extends the default `vllm-openai` image with `vllm[zen]`. Build it with `docker build -f docker/Dockerfile.cpu --target vllm-openai-zen .` — see [Building with AMD Zen optimizations](#building-with-amd-zen-optimizations) for the full command and run instructions.
+
+### Reference
+
+For the design rationale, see [RFC #35089: In-Tree AMD Zen CPU Backend via zentorch](https://github.com/vllm-project/vllm/issues/35089).
+
+--8<-- [end:amd-zen-optimizations]
 --8<-- [start:extra-information]
 --8<-- [end:extra-information]

docs/getting_started/installation/gpu.rocm.inc.md

@@ -27,6 +27,19 @@ If you need a different ROCm version or want to use an existing PyTorch installa
 --8<-- [end:set-up-using-python]
 --8<-- [start:pre-built-wheels]
 
+!!! warning "Python 3.12 required for ROCm wheels"
+
+    ROCm pre-built wheels are only available for **Python 3.12**. If you are using a different Python version (e.g. 3.11 or 3.13), the installer **will silently fall back** to the CUDA wheel from PyPI, which will fail on AMD GPUs with errors like `libcudart.so: cannot open shared object file`.
+
+    To check your Python version: `python3 --version`
+
+    If you need Python 3.12, you can create an isolated environment with `uv`:
+
+    ```bash
+    uv venv --python 3.12 --seed --managed-python
+    source .venv/bin/activate
+    ```
+
 To install the latest version of vLLM for Python 3.12, ROCm 7.0 and `glibc >= 2.35`.
 
 ```bash

docs/models/hardware_supported_models/cpu.md

@@ -1,5 +1,8 @@
 # CPU - Intel® Xeon®
 
+!!! note "AMD Zen CPUs"
+    On AMD Zen 4 / Zen 5 CPUs, AMD Zen optimizations are auto-enabled when the [`zentorch`](https://github.com/amd/ZenDNN-pytorch-plugin) package is installed. All models supported by vLLM on CPU are supported on AMD Zen as well; model compatibility does not change. This page reflects the current CPU reference validation matrix on Intel systems. See [AMD Zen optimizations](../../getting_started/installation/cpu.md#amd-zen-optimizations) for details.
+
 ## Validated Hardware
 
 | Hardware |

docs/models/pooling_models/README.md

@@ -184,7 +184,7 @@ Our online Server provides endpoints that correspond to the offline APIs:
 - Corresponding to `LLM.classify`:
     - [Classification API](classify.md#online-serving)(`/classify`)
 - Corresponding to `LLM.score`:
-    - [Score API](scoring.md#score-api)(`/score`)
+    - [Score API](scoring.md#score-api) (`/score`, `/v1/score`)
     - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
 - Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all types of pooling models.
 

docs/models/pooling_models/scoring.md

@@ -19,7 +19,7 @@ The score models is designed to compute similarity scores between two input prom
 - Offline APIs:
     - `LLM.score`
 - Online APIs:
-    - [Score API](scoring.md#score-api) (`/score`)
+    - [Score API](scoring.md#score-api) (`/score`, `/v1/score`)
     - [Cohere Rerank API](scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
 
 !!! note
@@ -157,7 +157,7 @@ A code example can be found here: [examples/basic/offline_inference/score.py](..
 
 ### Score API
 
-Our Score API (`/score`) is similar to `LLM.score`, compute similarity scores between two input prompts.
+Our Score API (`/score`, `/v1/score`) is similar to `LLM.score`, compute similarity scores between two input prompts.
 
 #### Parameters
 

docs/models/supported_models.md

@@ -374,7 +374,6 @@ th {
 | `BailingMoeForCausalLM` | Ling | `inclusionAI/Ling-lite-1.5`, `inclusionAI/Ling-plus`, etc. | ✅︎ | ✅︎ |
 | `BailingMoeV2ForCausalLM` | Ling | `inclusionAI/Ling-mini-2.0`, etc. | ✅︎ | ✅︎ |
 | `BailingMoeV2_5ForCausalLM` | Ling | `inclusionAI/Ling-2.5-1T`, `inclusionAI/Ring-2.5-1T` | | ✅︎ |
-| `BambaForCausalLM` | Bamba | `ibm-ai-platform/Bamba-9B-fp8`, `ibm-ai-platform/Bamba-9B` | ✅︎ | ✅︎ |
 | `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
 | `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
 | `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
@@ -386,7 +385,6 @@ th {
 | `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
 | `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | ✅︎ |
-| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
 | `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
 | `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
 | `Ernie4_5_MoeForCausalLM` | Ernie4.5MoE | `baidu/ERNIE-4.5-21B-A3B-PT`, `baidu/ERNIE-4.5-300B-A47B-PT`, etc. | ✅︎ | ✅︎ |
@@ -419,6 +417,7 @@ th {
 | `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ |
 | `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ |
 | `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
+| `HrmTextForCausalLM` | HRM-Text | `sapientinc/HRM-Text-1B`, etc. | | |
 | `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
 | `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
 | `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
@@ -488,7 +487,6 @@ th {
 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
 | `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
-| `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |
 | `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
 | `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | |
 | `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |

docs/serving/online_serving/README.md

@@ -9,12 +9,13 @@ We currently support the following OpenAI APIs:
 - [Completions API](./openai_compatible_server.md#completions-api) (`/v1/completions`)
     - Only applicable to [text generation models](../../models/generative_models.md).
     - *Note: `suffix` parameter is not supported.*
-- [Responses API](./openai_compatible_server.md#responses-api) (`/v1/responses`)
-    - Only applicable to [text generation models](../../models/generative_models.md).
 - [Chat Completions API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions`)
     - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](./openai_compatible_server.md#chat-template).
     - *Note: `user` parameter is ignored.*
     - *Note:* Setting the `parallel_tool_calls` parameter to `false` ensures vLLM only returns zero or one tool call per request. Setting it to `true` (the default) allows returning more than one tool call per request. There is no guarantee more than one tool call will be returned if this is set to `true`, as that behavior is model dependent and not all models are designed to support parallel tool calls.
+- [Chat Completions batch API](./openai_compatible_server.md#chat-api) (`/v1/chat/completions/batch`)
+- [Responses API](./openai_compatible_server.md#responses-api) (`/v1/responses`, `/v1/responses/{response_id}`, `/v1/responses/{response_id}/cancel`)
+    - Only applicable to [text generation models](../../models/generative_models.md).
 - [Embeddings API](../../models/pooling_models/embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`)
     - Only applicable to [embedding models](../../models/pooling_models/embed.md).
 - [Transcriptions API](./speech_to_text.md#transcriptions-api) (`/v1/audio/transcriptions`)
@@ -24,7 +25,7 @@ We currently support the following OpenAI APIs:
 
 ## Anthropic APIs
 
-- Anthropic messages API (`/v1/messages`)
+- Anthropic messages API (`/v1/messages`, `/v1/messages/count_tokens`)
 
 ## Cohere APIs
 
@@ -35,10 +36,6 @@ We currently support the following OpenAI APIs:
     - Implements [Jina AI's v1 rerank API](https://jina.ai/reranker/)
     - compatible with [Cohere's v1 & v2 rerank APIs](https://docs.cohere.com/v2/reference/rerank)
 
-## SageMaker APIs
-
-- `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints)
-
 ## Pooling APIs
 
 For further details on pooling models, please refer to [this page](../../models/pooling_models/README.md).
@@ -51,7 +48,7 @@ For further details on pooling models, please refer to [this page](../../models/
     - [OpenAI-compatible Embeddings API](../../models/pooling_models/embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`)
     - Only applicable to [embedding models](../../models/pooling_models/embed.md).
 - [Scoring Usages](../../models/pooling_models/scoring.md)
-    - [Score API](../../models/pooling_models/scoring.md#score-api) (`/score`)
+    - [Score API](../../models/pooling_models/scoring.md#score-api) (`/score`, `/v1/score`)
     - [Cohere Rerank API](../../models/pooling_models/scoring.md#rerank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
     - Applicable to [score models](../../models/pooling_models/scoring.md) (cross-encoder, bi-encoder, late-interaction).
 - [Pooling API](../../models/pooling_models/README.md#pooling-api) (`/pooling`)
@@ -68,17 +65,6 @@ For further details on speech to text, please refer to [this page](speech_to_tex
 - [Realtime API](./speech_to_text.md#realtime-api) (`/v1/realtime`)
     - Only applicable to [Automatic Speech Recognition (ASR) models](../../models/supported_models.md#realtime-transcription).
 
-## Disaggregated APIs
-
-### Renderer APIs
-
-For further details on renderer APIs, please refer to [this page](renderer.md).
-
-- [Completions Render API](renderer.md) (`/v1/completions/render`)
-    - Render completion requests
-- [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`)
-    - Render chat completions
-
 ## Custom APIs
 
 - [Classification API](../../models/pooling_models/classify.md#classification-api) (`/classify`)
@@ -91,14 +77,79 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
     - Applicable to [CausalLM models](../../models/generative_models.md) (task `"generate"`).
     - Computes next-token probabilities for specified `label_token_ids`.
 
-## Utility APIs
+## Instrumentator APIs
+
+### Basic APIs
+
+- `/version` - Version information
+- `/load` - Server load metrics
+- `/v1/models` - List available models
+- `/health` - Health check
+
+### Metrics APIs
+
+For further details on metrics, please refer to [this page](../../design/metrics.md).
+
+- `/metrics` - Prometheus-compatible metrics HTTP endpoint
+
+### Offline API Documentation
+
+The FastAPI `/docs` endpoint requires an internet connection by default. To enable offline access in air-gapped environments, use the `--enable-offline-docs` flag:
+
+```bash
+vllm serve NousResearch/Meta-Llama-3-8B-Instruct --enable-offline-docs
+```
+
+### LoRA dynamic loading
+
+LoRA dynamic loading & unloading is enabled in the API server. This should ONLY be used for local development!
+
+- `/v1/load_lora_adapter` - LoRA dynamic loading
+- `/v1/unload_lora_adapter` - LoRA dynamic unloading
+
+### Profiling APIs
+
+For further details on profiling vLLM, please refer to [this page](../../contributing/profiling.md).
+
+- `/start_profile` - Start PyTorch profiler
+- `/stop_profile` - Stop PyTorch profiler
+
+### SageMaker APIs
+
+- `/ping` - SageMaker health check
+- `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints)
+
+## Disaggregated Everything
+
+### Tokens IN <> Tokens OUT
+
+- `/inference/v1/generate` - Generate completions
+- `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set)
+
+### Renderer APIs
+
+For further details on renderer APIs, please refer to [this page](renderer.md).
+
+- [Completions Render API](renderer.md) (`/v1/completions/render`)
+    - Render completion requests
+- [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`)
+    - Render chat completions
+
+### Derenderer APIs
+
+- `/v1/completions/derender` - Derenderer completion requests
+- `/v1/chat/completions/derender` - Derenderer chat completion requests
+
+## Tokenize APIs
 
 - `/tokenize` - Tokenize text
 - `/detokenize` - Detokenize tokens
-- `/health` - Health check
-- `/ping` - SageMaker health check
-- `/version` - Version information
-- `/load` - Server load metrics
+- `/tokenizer_info` - Get comprehensive tokenizer information including chat templates and configuration
+
+## Elastic Expert Parallelism (EEP)
+
+- `/scale_elastic_ep` - Trigger scaling operations
+- `/is_scaling_elastic_ep` - Check if scaling is in progress
 
 ## Server in development mode
 
@@ -120,7 +171,9 @@ For further details on Weight Transfer, please refer to [this page](../../traini
 - `/resume` - Resume generation
 - `/is_paused` - Check if generation is paused
 - `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
+- `/start_weight_update` - Prepares the inference engine for a weight update.
 - `/update_weights` - Update model weights (can alter model behavior)
+- `/finish_weight_update` - Finalizes the weight update
 - `/get_world_size` - Get distributed world size
 
 ### Collective RPC
@@ -189,14 +242,6 @@ the detected format, which can be one of:
 If the result is not what you expect, you can set the `--chat-template-content-format` CLI argument
 to override which format to use.
 
-## Offline API Documentation
-
-The FastAPI `/docs` endpoint requires an internet connection by default. To enable offline access in air-gapped environments, use the `--enable-offline-docs` flag:
-
-```bash
-vllm serve NousResearch/Meta-Llama-3-8B-Instruct --enable-offline-docs
-```
-
 ## Ray Serve LLM
 
 Ray Serve LLM enables scalable, production-grade serving of the vLLM engine. It integrates tightly with vLLM and extends it with features such as auto-scaling, load balancing, and back-pressure.

docs/serving/online_serving/openai_compatible_server.md

@@ -9,12 +9,13 @@ We currently support the following OpenAI APIs:
 - [Completions API](#completions-api) (`/v1/completions`)
     - Only applicable to [text generation models](../../models/generative_models.md).
     - *Note: `suffix` parameter is not supported.*
-- [Responses API](#responses-api) (`/v1/responses`)
-    - Only applicable to [text generation models](../../models/generative_models.md).
 - [Chat Completions API](#chat-api) (`/v1/chat/completions`)
     - Only applicable to [text generation models](../../models/generative_models.md) with a [chat template](../online_serving/README.md#chat-template).
     - *Note: `user` parameter is ignored.*
     - *Note:* Setting the `parallel_tool_calls` parameter to `false` ensures vLLM only returns zero or one tool call per request. Setting it to `true` (the default) allows returning more than one tool call per request. There is no guarantee more than one tool call will be returned if this is set to `true`, as that behavior is model dependent and not all models are designed to support parallel tool calls.
+- [Chat Completions batch API](#chat-api) (`/v1/chat/completions/batch`)
+- [Responses API](#responses-api) (`/v1/responses`, `/v1/responses/{response_id}`, `/v1/responses/{response_id}/cancel`)
+    - Only applicable to [text generation models](../../models/generative_models.md).
 - [Embeddings API](../../models/pooling_models/embed.md#openai-compatible-embeddings-api) (`/v1/embeddings`)
     - Only applicable to [embedding models](../../models/pooling_models/embed.md).
 - [Transcriptions API](./speech_to_text.md#transcriptions-api) (`/v1/audio/transcriptions`)

docs/training/layerwise.md

@@ -28,9 +28,9 @@ For more information on implementation, see [Low Level `layerwise` API](#low-lev
 Online quantization refers to when a user provides full precision weights and those weights are quantized on-the-fly as they are loaded into the model. The layerwise reloading system handles this by treating online quantization as a **processing** step, which is then handled in an online way both during first-time load and during reload. A typical online quantization method implementation should look like this:
 
 ```python
-class Fp8OnlineLinearMethod(Fp8LinearMethod):
-    """Online version of Fp8LinearMethod which loads a full precision checkpoint
-    and quantizes weights during loading."""
+class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
+    """Online version of FP8 per-tensor quantization which loads a full
+    precision checkpoint and quantizes weights during loading."""
 
     uses_meta_device: bool = True
 

examples/generate/multimodal/vision_language_offline.py

@@ -2533,15 +2533,17 @@ MODELS_NEED_VIDEO_METADATA = [
 
 MODELS_SUPPORT_VIT_CUDA_GRAPH = [
     "llama4",
-    "internvl_chat",
+    "qwen2_vl",
     "qwen2_5_vl",
     "qwen3_vl",
     "qwen3_vl_moe",
-    "qwen2_vl",
+    "kimi_vl",
     "qwen3_5",
     "qwen3_5_moe",
+    "internvl_chat",
     "stepvl",
     "glm4_1v",
+    "deepseek_ocr",
 ]
 
 

examples/tool_chat_template_gemma4.jinja

@@ -116,7 +116,9 @@
     }
 {%- endmacro -%}
 {%- macro format_argument(argument, escape_keys=True) -%}
-    {%- if argument is string -%}
+    {%- if argument is none -%}
+        {{- 'null' -}}
+    {%- elif argument is string -%}
         {{- '<|"|>' + argument + '<|"|>' -}}
     {%- elif argument is boolean -%}
         {{- 'true' if argument else 'false' -}}
@@ -172,18 +174,21 @@
     {{- '<tool_response|>' -}}
 {%- endmacro -%}
 
-{%- set ns = namespace(prev_message_type=None) -%}
+{#- ===== SETUP ===== -#}
+{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
 {%- set loop_messages = messages -%}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set preserve_thinking = preserve_thinking | default(false) -%}
 {{- bos_token -}}
 {#- Handle System/Tool Definitions Block -#}
-{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
     {{- '<|turn>system\n' -}}
     {#- Inject Thinking token at the very top of the FIRST system turn -#}
-    {%- if enable_thinking is defined and enable_thinking -%}
+    {%- if enable_thinking -%}
         {{- '<|think|>\n' -}}
         {%- set ns.prev_message_type = 'think' -%}
     {%- endif -%}
-    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+    {%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
         {%- if messages[0]['content'] is string -%}
             {{- messages[0]['content'] | trim -}}
         {%- elif messages[0]['content'] is sequence -%}
@@ -217,31 +222,24 @@
     {%- if message['role'] != 'tool' -%}
     {%- set ns.prev_message_type = None -%}
     {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
-    {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
-    {%- set prev_nt = namespace(role=None, found=false) -%}
-    {%- if loop.index0 > 0 -%}
-        {%- for j in range(loop.index0 - 1, -1, -1) -%}
-            {%- if not prev_nt.found -%}
-                {%- if loop_messages[j]['role'] != 'tool' -%}
-                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
-                    {%- set prev_nt.found = true -%}
-                {%- endif -%}
-            {%- endif -%}
-        {%- endfor -%}
-    {%- endif -%}
-    {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+    {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#}
+    {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
     {%- if not continue_same_model_turn -%}
         {{- '<|turn>' + role + '\n' }}
+        {%- if role == 'model' and not enable_thinking and not (message.get('reasoning') or message.get('reasoning_content')) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
     {%- endif -%}
 
     {#- Render reasoning/reasoning_content as thinking channel -#}
     {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
-    {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+    {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%}
+    {%- if thinking_text and thinking_gate -%}
         {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
     {%- endif -%}
 
-            {%- if message['tool_calls'] -%}
-                {%- for tool_call in message['tool_calls'] -%}
+            {%- if message.get('tool_calls') -%}
+                {%- for tool_call in message.get('tool_calls') -%}
                     {%- set function = tool_call['function'] -%}
                     {{- '<|tool_call>call:' + function['name'] + '{' -}}
                     {%- if function['arguments'] is mapping -%}
@@ -251,8 +249,13 @@
                             {%- set ns_args.found_first = true -%}
                             {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
                         {%- endfor -%}
-                    {%- elif function['arguments'] is string -%}
-                        {{- function['arguments'] -}}
+                    {%- elif function['arguments'] is none -%}
+                    {%- else -%}
+                        {{- raise_exception(
+                            "chat_template: tool_calls[].function.arguments must be a "
+                            "JSON object (mapping), not a string. Deserialize arguments "
+                            "before passing to the template."
+                        ) -}}
                     {%- endif -%}
                     {{- '}<tool_call|>' -}}
                 {%- endfor -%}
@@ -262,7 +265,7 @@
             {%- set ns_tr_out = namespace(flag=false) -%}
             {%- if message.get('tool_responses') -%}
                 {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
-                {%- for tool_response in message['tool_responses'] -%}
+                {%- for tool_response in message.get('tool_responses') -%}
                     {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
                     {%- set ns_tr_out.flag = true -%}
                     {%- set ns.prev_message_type = 'tool_response' -%}
@@ -277,8 +280,8 @@
                     {%- else -%}
                         {%- set follow = loop_messages[k] -%}
                         {#- Resolve tool_call_id to function name -#}
-                        {%- set ns_tname = namespace(name=follow.get('name') | default('unknown', true)) -%}
-                        {%- for tc in message['tool_calls'] -%}
+                        {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
+                        {%- for tc in message.get('tool_calls') -%}
                             {%- if tc.get('id') == follow.get('tool_call_id') -%}
                                 {%- set ns_tname.name = tc['function']['name'] -%}
                             {%- endif -%}
@@ -296,9 +299,9 @@
                             {%- endfor -%}
                             {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
                             {%- for part in tool_body -%}
-                                {%- if part.get('type') == 'image' -%}
+                                {%- if part.get('type') in ['image', 'image_url'] -%}
                                     {{- '<|image|>' -}}
-                                {%- elif part.get('type') == 'audio' -%}
+                                {%- elif part.get('type') in ['audio', 'input_audio'] -%}
                                     {{- '<|audio|>' -}}
                                 {%- elif part.get('type') == 'video' -%}
                                     {{- '<|video|>' -}}
@@ -314,29 +317,26 @@
             {%- endif -%}
 
             {%- set captured_content -%}
-            {%- if message['content'] is string -%}
+            {%- if message.get('content') is string -%}
                 {%- if role == 'model' -%}
                     {{- strip_thinking(message['content']) -}}
                 {%- else -%}
                     {{- message['content'] | trim -}}
                 {%- endif -%}
-            {%- elif message['content'] is sequence -%}
+            {%- elif message.get('content') is sequence -%}
                 {%- for item in message['content'] -%}
-                    {%- if item['type'] == 'text' -%}
+                    {%- if item.get('type') == 'text' -%}
                         {%- if role == 'model' -%}
                             {{- strip_thinking(item['text']) -}}
                         {%- else -%}
                             {{- item['text'] | trim -}}
                         {%- endif -%}
-                    {%- elif item['type'] == 'image' -%}
+                    {%- elif item.get('type') in ['image', 'image_url'] -%}
                         {{- '<|image|>' -}}
-                        {%- set ns.prev_message_type = 'image' -%}
-                    {%- elif item['type'] == 'audio' -%}
+                    {%- elif item.get('type') in ['audio', 'input_audio'] -%}
                         {{- '<|audio|>' -}}
-                        {%- set ns.prev_message_type = 'audio' -%}
-                    {%- elif item['type'] == 'video' -%}
+                    {%- elif item.get('type') == 'video' -%}
                         {{- '<|video|>' -}}
-                        {%- set ns.prev_message_type = 'video' -%}
                     {%- endif -%}
                 {%- endfor -%}
             {%- endif -%}
@@ -345,19 +345,43 @@
             {{- captured_content -}}
             {%- set has_content = captured_content | trim | length > 0 -%}
 
+        {#- Forward-scan: find next non-tool message role for continuation detection -#}
+        {%- set next_nt = namespace(role=None, found=false) -%}
+        {%- for j in range(loop.index0 + 1, loop_messages | length) -%}
+            {%- if not next_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set next_nt.role = loop_messages[j]['role'] -%}
+                    {%- set next_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+
+        {%- set continues_into_next = (
+            role == 'model'
+            and next_nt.role == 'assistant'
+            and (not message.get('tool_calls') or ns_tr_out.flag)
+        ) -%}
+
         {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
             {{- '<|tool_response>' -}}
+        {%- elif continues_into_next -%}
+            {{- '\n' -}}
         {%- elif not (ns_tr_out.flag and not has_content) -%}
             {{- '<turn|>\n' -}}
         {%- endif -%}
+
+    {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
+    {%- set ns.prev_non_tool_role = message['role'] -%}
     {%- endif -%}
 {%- endfor -%}
 
 {%- if add_generation_prompt -%}
     {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
         {{- '<|turn>model\n' -}}
-        {%- if not enable_thinking | default(false) -%}
+        {%- if not enable_thinking -%}
             {{- '<|channel>thought\n<channel|>' -}}
         {%- endif -%}
+    {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%}
+        {{- '<|channel>thought\n' -}}
     {%- endif -%}
-{%- endif -%}
+{%- endif -%}

pyproject.toml

@@ -162,6 +162,8 @@ dout = "dout"
 Pn = "Pn"
 arange = "arange"
 thw = "thw"
+# temporal position ids (parallels hpos/wpos in vision RoPE)
+tpos = "tpos"
 subtile = "subtile"
 HSA = "HSA"
 setp = "setp"

requirements/common.txt

@@ -29,6 +29,7 @@ xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
 typing_extensions >= 4.10
 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
 partial-json-parser # used for parsing partial JSON outputs
+jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
 pyzmq >= 25.0.0
 msgspec
 mistral_common[image] >= 1.11.3

requirements/kv_connectors.txt

@@ -2,5 +2,5 @@ lmcache >= 0.3.9
 # CuPy 14.1.0 imports pytest from cupy.testing._random. Use <14.1.0
 # until a fixed newer release is verified for runtime images.
 cupy-cuda13x < 14.1.0
-nixl >= 1.1.0 # Required for disaggregated prefill
+nixl == 1.2.0 # Required for disaggregated prefill
 mooncake-transfer-engine >= 0.3.8

requirements/rocm.txt

@@ -11,7 +11,7 @@ numba == 0.65.0 # Required for N-gram speculative decoding
 datasets
 peft
 pytest-asyncio
-tensorizer==2.10.1
+tensorizer==2.12.1
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8

requirements/test/cuda.in

@@ -1,6 +1,6 @@
 # testing
 pytest
-tensorizer==2.10.1
+tensorizer==2.12.1
 pytest-forked
 pytest-asyncio
 pytest-rerunfailures

requirements/test/cuda.txt

@@ -360,6 +360,7 @@ jsonpointer==3.0.0
     # via jsonschema
 jsonschema==4.23.0
     # via
+    #   -c requirements/common.txt
     #   hypothesis-jsonschema
     #   mistral-common
     #   ray
@@ -965,7 +966,7 @@ tenacity==9.1.2
     #   gpt-oss
     #   lm-eval
     #   plotly
-tensorizer==2.10.1
+tensorizer==2.12.1
     # via -r requirements/test/cuda.in
 termcolor==3.1.0
     # via gpt-oss

requirements/test/nightly-torch.txt

@@ -1,6 +1,6 @@
 # testing
 pytest
-tensorizer==2.10.1
+tensorizer==2.12.1
 pytest-forked
 pytest-asyncio
 pytest-rerunfailures

requirements/test/rocm.in

@@ -2,7 +2,7 @@
 
 # testing
 pytest
-tensorizer==2.10.1
+tensorizer==2.12.1
 pytest-forked
 pytest-asyncio
 pytest-rerunfailures

requirements/test/rocm.txt

@@ -439,6 +439,8 @@ jsonpointer==3.1.0
     # via jsonschema
 jsonschema==4.26.0
     # via
+    #   -c requirements/common.txt
+    #   -r requirements/test/../common.txt
     #   hypothesis-jsonschema
     #   mcp
     #   mistral-common
@@ -1180,7 +1182,7 @@ tenacity==9.1.4
     # via
     #   gpt-oss
     #   lm-eval
-tensorizer==2.10.1
+tensorizer==2.12.1
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in

requirements/test/xpu.in

@@ -13,7 +13,7 @@ pytest-shard
 absl-py
 accelerate
 arctic-inference
-lm_eval[api]
+lm_eval[api]>=0.4.12
 modelscope
 
 # --- Audio Processing ---

requirements/test/xpu.txt

@@ -33,7 +33,6 @@ arctic-inference==0.1.1
 attrs==26.1.0
     # via
     #   aiohttp
-    #   jsonlines
     #   jsonschema
     #   referencing
 audioread==3.0.1
@@ -225,10 +224,9 @@ joblib==1.5.3
     #   librosa
     #   nltk
     #   scikit-learn
-jsonlines==4.0.0
-    # via lm-eval
 jsonschema==4.26.0
     # via
+    #   -c requirements/common.txt
     #   hypothesis-jsonschema
     #   mistral-common
     #   schemathesis
@@ -246,7 +244,7 @@ librosa==0.10.2.post1
     # via -r requirements/test/xpu.in
 llvmlite==0.47.0
     # via numba
-lm-eval==0.4.11
+lm-eval==0.4.12
     # via -r requirements/test/xpu.in
 lxml==6.0.2
     # via
@@ -733,5 +731,3 @@ xxhash==3.6.0
     #   evaluate
 yarl==1.23.0
     # via aiohttp
-zstandard==0.25.0
-    # via lm-eval

requirements/tpu.txt

@@ -12,4 +12,4 @@ ray[data]
 setuptools==78.1.0
 setuptools-rust>=1.9.0
 nixl==0.3.0
-tpu-inference==0.21.0
+tpu-inference==0.22.1

requirements/xpu.txt

@@ -16,5 +16,5 @@ torch==2.12.0
 torchaudio
 torchvision
 
-auto_round_lib>=0.13.0
+auto_round_lib>=0.13.3
 vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9.1/vllm_xpu_kernels-0.1.9.1-cp38-abi3-manylinux_2_28_x86_64.whl

rust/Cargo.lock

@@ -5804,6 +5804,7 @@ dependencies = [
  "enum-as-inner",
  "expect-test",
  "futures",
+ "parking_lot",
  "rmp-serde",
  "serde",
  "serde_json",
@@ -5879,6 +5880,7 @@ dependencies = [
  "expect-test",
  "futures",
  "http-body",
+ "indexmap 2.13.0",
  "itertools 0.14.0",
  "libc",
  "llm-multimodal",

rust/Cargo.toml

@@ -105,7 +105,7 @@ tonic-prost = "0.14.5"
 tonic-prost-build = "0.14.5"
 tool-parser = "1.2.0"
 tower = { version = "0.5.3", features = ["util"] }
-tower-http = { version = "0.6.8", features = ["trace"] }
+tower-http = { version = "0.6.8", features = ["cors", "trace"] }
 tracing = { version = "0.1.44", features = ["release_max_level_debug"] }
 tracing-futures = { version = "0.2.5", features = ["futures-03"] }
 tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] }

rust/src/chat/src/lib.rs

@@ -145,8 +145,8 @@ impl ChatLlm {
         self.text.tokenizer_vocab_size()
     }
 
-    /// Model vocabulary size, else `None`.
-    pub fn model_vocab_size(&self) -> Option<usize> {
+    /// Model vocabulary size from the model config.
+    pub fn model_vocab_size(&self) -> usize {
         self.text.model_vocab_size()
     }
 
@@ -235,6 +235,12 @@ impl ChatLlm {
         Ok(token_ids)
     }
 
+    /// Abort in-flight requests by their external (user-supplied) request ids.
+    pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
+        self.text.abort(external_ids).await?;
+        Ok(())
+    }
+
     /// Shut down the underlying LLM client and its background tasks.
     pub async fn shutdown(self) -> Result<()> {
         self.text.shutdown().await?;
@@ -271,7 +277,7 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, granite4, hermes, hy_v3, internlm, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
     }
 
     #[test]
@@ -282,6 +288,6 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, gemma4, glm45, kimi, kimi_k2, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string());
     }
 }

rust/src/chat/src/output/default/mod.rs

@@ -37,6 +37,7 @@ trait_set! {
 pub struct DefaultChatOutputProcessor {
     reasoning_parser: Option<Box<dyn ReasoningParser>>,
     tool_parser: Option<Box<dyn ToolParser>>,
+    parallel_tool_calls: bool,
 }
 
 impl DefaultChatOutputProcessor {
@@ -74,6 +75,7 @@ impl DefaultChatOutputProcessor {
         Ok(Self {
             reasoning_parser,
             tool_parser,
+            parallel_tool_calls: request.parallel_tool_calls,
         })
     }
 
@@ -86,6 +88,7 @@ impl DefaultChatOutputProcessor {
         Self {
             reasoning_parser: None,
             tool_parser: None,
+            parallel_tool_calls: true,
         }
     }
 
@@ -159,7 +162,7 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
     fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
         let reasoning = reasoning_event_stream(decoded, self.reasoning_parser);
         let tool = tool_event_stream(reasoning, self.tool_parser);
-        let structured = structured_chat_event_stream(tool);
+        let structured = structured_chat_event_stream(tool, self.parallel_tool_calls);
 
         Ok(structured.boxed())
     }

rust/src/chat/src/output/default/tool.rs

@@ -473,7 +473,7 @@ mod tests {
             })));
         let parser = DeepSeekV4ToolParser::create(&deepseek_v4_test_tools()).unwrap();
         let assistant_events = tool_event_stream(stream::iter(events), Some(parser));
-        let chat_events = structured_chat_event_stream(assistant_events);
+        let chat_events = structured_chat_event_stream(assistant_events, true);
 
         ChatEventStream::new("req_deepseek_v4".to_string(), Box::pin(chat_events))
             .collect_message()
@@ -717,9 +717,10 @@ mod tests {
 
         let message = ChatEventStream::new(
             "req_fallback".to_string(),
-            Box::pin(structured_chat_event_stream(stream::iter(
-                events.into_iter().map(Ok),
-            ))),
+            Box::pin(structured_chat_event_stream(
+                stream::iter(events.into_iter().map(Ok)),
+                true,
+            )),
         )
         .collect_message()
         .await
@@ -968,9 +969,10 @@ mod tests {
         ));
         let collected = ChatEventStream::new(
             "req_final_only".to_string(),
-            Box::pin(structured_chat_event_stream(stream::iter(
-                events.into_iter().map(Ok),
-            ))),
+            Box::pin(structured_chat_event_stream(
+                stream::iter(events.into_iter().map(Ok)),
+                true,
+            )),
         )
         .collect_message()
         .await

rust/src/chat/src/output/harmony/mod.rs

@@ -35,6 +35,7 @@ use crate::request::ChatRequest;
 pub struct HarmonyChatOutputProcessor {
     encoding: &'static HarmonyEncoding,
     tool_calls_enabled: bool,
+    parallel_tool_calls: bool,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -76,6 +77,7 @@ impl HarmonyChatOutputProcessor {
         Ok(Self {
             encoding: harmony_encoding()?,
             tool_calls_enabled: request.tool_parsing_enabled(),
+            parallel_tool_calls: request.parallel_tool_calls,
         })
     }
 }
@@ -110,7 +112,11 @@ impl ChatOutputProcessor for HarmonyChatOutputProcessor {
     fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
         let assistant =
             harmony_assistant_event_stream(decoded, self.encoding, self.tool_calls_enabled);
-        Ok(crate::output::structured::structured_chat_event_stream(assistant).boxed())
+        Ok(crate::output::structured::structured_chat_event_stream(
+            assistant,
+            self.parallel_tool_calls,
+        )
+        .boxed())
     }
 }
 

rust/src/chat/src/output/structured.rs

@@ -53,16 +53,22 @@ struct StructuredEventState {
     open_tool_call: Option<OpenToolCall>,
     /// Next OpenAI-compatible tool-call ordinal.
     next_tool_call_index: usize,
+    /// Whether more than one tool call may be surfaced northbound.
+    parallel_tool_calls: bool,
+    /// Whether the current tool-call parse is being suppressed.
+    suppressing_tool_call: bool,
 }
 
 impl StructuredEventState {
     /// Create one fresh assembly state for a new streamed response.
-    fn new() -> Self {
+    fn new(parallel_tool_calls: bool) -> Self {
         Self {
             message: AssistantMessage::default(),
             open_text_block: None,
             open_tool_call: None,
             next_tool_call_index: 0,
+            parallel_tool_calls,
+            suppressing_tool_call: false,
         }
     }
 
@@ -98,6 +104,12 @@ impl StructuredEventState {
 
         let index = self.next_tool_call_index;
         self.next_tool_call_index += 1;
+        if !self.parallel_tool_calls && index >= 1 {
+            self.suppressing_tool_call = true;
+            return Ok(events);
+        }
+
+        self.suppressing_tool_call = false;
         self.open_tool_call = Some(OpenToolCall {
             index,
             id: id.clone(),
@@ -110,6 +122,10 @@ impl StructuredEventState {
 
     /// Append one incremental tool-call arguments delta.
     fn push_tool_call_arguments(&mut self, delta: String) -> Result<Vec<ChatEvent>> {
+        if self.suppressing_tool_call {
+            return Ok(Vec::new());
+        }
+
         let mut events = Vec::new();
         let Some(open_tool_call) = self.open_tool_call.as_mut() else {
             return Err(Error::ToolCallStreamInvariant {
@@ -207,6 +223,11 @@ impl StructuredEventState {
 
     /// Finalize the currently open tool call, if present.
     fn close_open_tool_call(&mut self, events: &mut Vec<ChatEvent>) {
+        if self.suppressing_tool_call {
+            self.suppressing_tool_call = false;
+            return;
+        }
+
         let Some(open_tool_call) = self.open_tool_call.take() else {
             return;
         };
@@ -229,11 +250,12 @@ impl StructuredEventState {
 #[try_stream]
 pub(crate) async fn structured_chat_event_stream(
     stream: impl AssistantEventStream,
+    parallel_tool_calls: bool,
     mut y: TryYielder<ChatEvent, Error>,
 ) -> Result<()> {
     pin_mut!(stream);
 
-    let mut state = StructuredEventState::new();
+    let mut state = StructuredEventState::new(parallel_tool_calls);
 
     while let Some(event) = stream.next().await.transpose()? {
         match event {
@@ -315,7 +337,7 @@ mod tests {
             }),
         ]);
 
-        let events = structured_chat_event_stream(events)
+        let events = structured_chat_event_stream(events, true)
             .collect::<Vec<_>>()
             .await
             .into_iter()
@@ -369,7 +391,7 @@ mod tests {
             }),
         ]);
 
-        let events = structured_chat_event_stream(events)
+        let events = structured_chat_event_stream(events, true)
             .collect::<Vec<_>>()
             .await
             .into_iter()
@@ -420,7 +442,7 @@ mod tests {
             }),
         ]);
 
-        let events = structured_chat_event_stream(events)
+        let events = structured_chat_event_stream(events, true)
             .collect::<Vec<_>>()
             .await
             .into_iter()
@@ -471,7 +493,7 @@ mod tests {
             }),
         ]);
 
-        let events = structured_chat_event_stream(events)
+        let events = structured_chat_event_stream(events, true)
             .collect::<Vec<_>>()
             .await
             .into_iter()
@@ -499,7 +521,7 @@ mod tests {
             delta: "{}".to_string(),
         })]);
 
-        let err = structured_chat_event_stream(events)
+        let err = structured_chat_event_stream(events, true)
             .collect::<Vec<_>>()
             .await
             .into_iter()
@@ -509,4 +531,56 @@ mod tests {
 
         assert!(matches!(err, Error::ToolCallStreamInvariant { .. }));
     }
+
+    #[tokio::test]
+    async fn structured_stream_suppresses_later_tool_calls_when_parallel_disabled() {
+        let events = stream::iter(vec![
+            Ok(AssistantEvent::ToolCallStart {
+                id: "call_1".to_string(),
+                name: "first".to_string(),
+            }),
+            Ok(AssistantEvent::ToolCallArgumentsDelta {
+                delta: r#"{"a":1}"#.to_string(),
+            }),
+            Ok(AssistantEvent::ToolCallStart {
+                id: "call_2".to_string(),
+                name: "second".to_string(),
+            }),
+            Ok(AssistantEvent::ToolCallArgumentsDelta {
+                delta: r#"{"b":2}"#.to_string(),
+            }),
+            Ok(AssistantEvent::Done {
+                usage: vllm_llm::TokenUsage {
+                    prompt_token_count: 1,
+                    output_token_count: 1,
+                    cached_token_count: 0,
+                },
+                finish_reason: FinishReason::stop_eos(),
+                kv_transfer_params: None,
+            }),
+        ]);
+
+        let events = structured_chat_event_stream(events, false)
+            .collect::<Vec<_>>()
+            .await
+            .into_iter()
+            .collect::<crate::Result<Vec<_>>>()
+            .unwrap();
+
+        assert!(matches!(
+            events[0],
+            ChatEvent::ToolCallStart { index: 0, .. }
+        ));
+        assert!(matches!(
+            events[1],
+            ChatEvent::ToolCallArgumentsDelta { index: 0, .. }
+        ));
+        assert!(matches!(events[2], ChatEvent::ToolCallEnd { index: 0, .. }));
+        let ChatEvent::Done { message, .. } = &events[3] else {
+            panic!("expected done");
+        };
+        let tool_calls = message.tool_calls().collect::<Vec<_>>();
+        assert_eq!(tool_calls.len(), 1);
+        assert_eq!(tool_calls[0].name, "first");
+    }
 }

rust/src/chat/src/parser/reasoning/mod.rs

@@ -5,9 +5,9 @@ use std::sync::LazyLock;
 pub use vllm_reasoning_parser::{
     CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
     DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
-    KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
-    ReasoningDelta, ReasoningError, ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser,
-    Step3p5ReasoningParser,
+    KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
+    NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
+    ReasoningParser, SeedOssReasoningParser, Step3ReasoningParser, Step3p5ReasoningParser,
 };
 use vllm_tokenizer::DynTokenizer;
 
@@ -24,6 +24,7 @@ pub mod names {
     pub const KIMI: &str = "kimi";
     pub const KIMI_K2: &str = "kimi_k2";
     pub const MINIMAX_M2: &str = "minimax_m2";
+    pub const MINIMAX_M3: &str = "minimax_m3";
     pub const NEMOTRON_V3: &str = "nemotron_v3";
     pub const QWEN3: &str = "qwen3";
     pub const SEED_OSS: &str = "seed_oss";
@@ -62,6 +63,7 @@ impl ReasoningParserFactory {
             .register_parser::<KimiReasoningParser>(names::KIMI)
             .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
             .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
+            .register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
             .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
             .register_parser::<Qwen3ReasoningParser>(names::QWEN3)
             .register_parser::<SeedOssReasoningParser>(names::SEED_OSS)
@@ -90,6 +92,8 @@ impl ReasoningParserFactory {
             .register_pattern("step3", names::STEP3)
             .register_pattern("seed-oss", names::SEED_OSS)
             .register_pattern("seedoss", names::SEED_OSS)
+            .register_pattern("minimax-m3", names::MINIMAX_M3)
+            .register_pattern("mm-m3", names::MINIMAX_M3)
             .register_pattern("minimax", names::MINIMAX_M2)
             .register_pattern("mm-m2", names::MINIMAX_M2)
             .register_pattern("cohere", names::COHERE_CMD)

rust/src/chat/src/parser/reasoning/tests.rs

@@ -34,10 +34,12 @@ fn factory_contains_and_lists_registered_parsers() {
     assert!(factory.contains(names::DEEPSEEK_V4));
     assert!(factory.contains(names::SEED_OSS));
     assert!(factory.contains(names::STEP3P5));
+    assert!(factory.contains(names::MINIMAX_M3));
     assert!(factory.list().contains(&names::QWEN3.to_string()));
     assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
     assert!(factory.list().contains(&names::SEED_OSS.to_string()));
     assert!(factory.list().contains(&names::STEP3P5.to_string()));
+    assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
 }
 
 #[test]
@@ -88,6 +90,19 @@ fn factory_routes_seed_oss_models() {
     );
 }
 
+#[test]
+fn factory_resolves_minimax_m3_before_generic_minimax() {
+    let factory = ReasoningParserFactory::new();
+    assert_eq!(
+        factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
+        Some(names::MINIMAX_M3)
+    );
+    assert_eq!(
+        factory.resolve_name_for_model("mm-m3"),
+        Some(names::MINIMAX_M3)
+    );
+}
+
 #[test]
 fn factory_rejects_unknown_parser_names() {
     let tokenizer = Arc::new(FakeTokenizer);

rust/src/chat/src/parser/tool/mod.rs

@@ -6,8 +6,9 @@ pub use vllm_tool_parser::{
     DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
     Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, Granite4ToolParser, HermesToolParser,
     HyV3ToolParser, Internlm2ToolParser, KimiK2ToolParser, Llama3JsonToolParser,
-    MinimaxM2ToolParser, MistralToolParser, Phi4MiniJsonToolParser, Qwen3CoderToolParser,
-    Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError, ToolParserOutput,
+    MinimaxM2ToolParser, MinimaxM3ToolParser, MistralToolParser, Phi4MiniJsonToolParser,
+    Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
+    ToolParserOutput,
 };
 
 use crate::parser::ParserFactory;
@@ -32,6 +33,7 @@ pub mod names {
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
     pub const MINIMAX_M2: &str = "minimax_m2";
+    pub const MINIMAX_M3: &str = "minimax_m3";
     pub const MISTRAL: &str = "mistral";
     pub const PHI4_MINI_JSON: &str = "phi4_mini_json";
     pub const QWEN3_CODER: &str = "qwen3_coder";
@@ -73,6 +75,7 @@ impl ToolParserFactory {
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
             .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
+            .register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
             .register_parser::<MistralToolParser>(names::MISTRAL)
             .register_parser::<Phi4MiniJsonToolParser>(names::PHI4_MINI_JSON)
             .register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
@@ -111,6 +114,8 @@ impl ToolParserFactory {
             .register_pattern("gemma-4", names::GEMMA4)
             .register_pattern("granite-4", names::GRANITE4)
             .register_pattern("kimi-k2", names::KIMI_K2)
+            .register_pattern("minimax-m3", names::MINIMAX_M3)
+            .register_pattern("mm-m3", names::MINIMAX_M3)
             .register_pattern("minimax", names::MINIMAX_M2)
             .register_pattern("mm-m2", names::MINIMAX_M2);
 

rust/src/chat/src/parser/tool/tests.rs

@@ -157,6 +157,14 @@ fn factory_new_resolves_default_patterns() {
         factory.resolve_name_for_model("tencent/Hy3-preview"),
         Some(names::HY_V3)
     );
+    assert_eq!(
+        factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
+        Some(names::MINIMAX_M3)
+    );
+    assert_eq!(
+        factory.resolve_name_for_model("org/mm-m3-base"),
+        Some(names::MINIMAX_M3)
+    );
     assert_eq!(
         factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
         Some(names::MINIMAX_M2)

rust/src/chat/src/request.rs

@@ -406,6 +406,10 @@ pub struct ChatRequest {
     pub tools: Vec<ChatTool>,
     /// Tool-choice behavior for this request.
     pub tool_choice: ChatToolChoice,
+    /// Whether the model may return more than one tool call per response.
+    ///
+    /// When `false`, only the first parsed tool call is surfaced northbound.
+    pub parallel_tool_calls: bool,
     /// Text decode options for incremental detokenization.
     pub decode_options: TextDecodeOptions,
     /// Whether to emit intermediate northbound content deltas before the
@@ -442,6 +446,7 @@ impl ChatRequest {
             chat_options: ChatOptions::default(),
             tools: Vec::new(),
             tool_choice: ChatToolChoice::None,
+            parallel_tool_calls: true,
             decode_options: TextDecodeOptions::default(),
             intermediate: true,
             priority: 0,

rust/src/cmd/src/cli.rs

@@ -23,8 +23,8 @@ use vllm_engine_core_client::TransportMode;
 use vllm_managed_engine::ManagedEngineConfig;
 use vllm_managed_engine::cli::{ManagedEngineArgs, repartition_managed_engine_args};
 use vllm_server::{
-    ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, HttpListenerMode,
-    ParserSelection, RendererSelection,
+    ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode, CorsConfig,
+    HttpListenerMode, ParserSelection, RendererSelection,
 };
 
 use crate::cli::unsupported::UnsupportedArgs;
@@ -84,6 +84,13 @@ pub enum Command {
     Serve(ServeArgs),
 }
 
+/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
+/// the CORS list arguments (e.g. `--allowed-origins '["*"]'`). Parsing the whole
+/// value as one item keeps clap from treating the field as a repeated flag.
+#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
+#[serde(transparent)]
+pub struct JsonStringList(pub Vec<String>);
+
 /// Runtime arguments shared by the external-engine and managed-engine paths.
 #[serde_as]
 #[derive(Educe, Clone, Args, PartialEq, Eq, Deserialize)]
@@ -127,6 +134,11 @@ pub struct SharedRuntimeArgs {
     /// `config.json`.
     #[arg(long)]
     pub max_model_len: Option<u32>,
+    /// Maximum number of log probabilities to return when `logprobs` is
+    /// specified in sampling parameters. `-1` means no cap.
+    #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)]
+    #[serde(default)]
+    pub max_logprobs: Option<i32>,
     /// TCP port for the gRPC Generate service. When not set, no gRPC server is
     /// started.
     #[arg(long)]
@@ -215,6 +227,30 @@ pub struct SharedRuntimeArgs {
     #[serde(default)]
     pub served_model_name: Vec<String>,
 
+    /// CORS allowed origins as a JSON list. `["*"]` allows any origin.
+    #[arg(long, value_parser = parse_json::<JsonStringList>, value_name = "JSON", default_value = r#"["*"]"#)]
+    #[serde(default = "default_cors_wildcard")]
+    pub allowed_origins: JsonStringList,
+
+    /// CORS allowed methods as a JSON list. `["*"]` allows the standard set.
+    #[arg(long, value_parser = parse_json::<JsonStringList>, value_name = "JSON", default_value = r#"["*"]"#)]
+    #[serde(default = "default_cors_wildcard")]
+    pub allowed_methods: JsonStringList,
+
+    /// CORS allowed request headers as a JSON list. `["*"]` mirrors the request.
+    #[arg(long, value_parser = parse_json::<JsonStringList>, value_name = "JSON", default_value = r#"["*"]"#)]
+    #[serde(default = "default_cors_wildcard")]
+    pub allowed_headers: JsonStringList,
+
+    /// Allow CORS credentials (cookies, authorization headers).
+    #[arg(
+        long,
+        default_missing_value = "true",
+        num_args = 0..=1
+    )]
+    #[serde(default)]
+    pub allow_credentials: bool,
+
     /// Unsupported Python vLLM frontend arguments recognized but not yet
     /// implemented in Rust.
     #[educe(Debug(ignore))]
@@ -254,16 +290,19 @@ impl SharedRuntimeArgs {
         input_address: String,
         output_address: String,
         coordinator_address: Option<String>,
+        engine_start_index: u32,
         engine_count: usize,
     ) -> Config {
         let ready_timeout = self.ready_timeout();
         let shutdown_timeout = self.shutdown_timeout();
         let api_server_options = self.api_server_options();
+        let cors = self.cors_config();
 
         Config {
             transport_mode: TransportMode::Bootstrapped {
                 input_address,
                 output_address,
+                engine_start_index,
                 engine_count,
                 ready_timeout,
             },
@@ -281,7 +320,9 @@ impl SharedRuntimeArgs {
             chat_template: self.chat_template,
             default_chat_template_kwargs: self.default_chat_template_kwargs,
             chat_template_content_format: self.chat_template_content_format,
+            max_logprobs: self.max_logprobs,
             api_server_options,
+            cors,
             api_keys: self.api_key,
             disable_log_stats: self.disable_log_stats,
             grpc_port: self.grpc_port,
@@ -303,6 +344,7 @@ impl SharedRuntimeArgs {
         let ready_timeout = self.ready_timeout();
         let shutdown_timeout = self.shutdown_timeout();
         let api_server_options = self.api_server_options();
+        let cors = self.cors_config();
 
         Config {
             transport_mode: TransportMode::HandshakeOwner {
@@ -324,7 +366,9 @@ impl SharedRuntimeArgs {
             chat_template: self.chat_template,
             default_chat_template_kwargs: self.default_chat_template_kwargs,
             chat_template_content_format: self.chat_template_content_format,
+            max_logprobs: self.max_logprobs,
             api_server_options,
+            cors,
             api_keys: self.api_key,
             disable_log_stats: self.disable_log_stats,
             grpc_port: self.grpc_port,
@@ -339,12 +383,25 @@ impl SharedRuntimeArgs {
             enable_request_id_headers: self.enable_request_id_headers,
         }
     }
+
+    fn cors_config(&self) -> CorsConfig {
+        CorsConfig {
+            allow_origins: self.allowed_origins.0.clone(),
+            allow_methods: self.allowed_methods.0.clone(),
+            allow_headers: self.allowed_headers.0.clone(),
+            allow_credentials: self.allow_credentials,
+        }
+    }
 }
 
 fn default_engine_ready_timeout_secs() -> u64 {
     600
 }
 
+fn default_cors_wildcard() -> JsonStringList {
+    JsonStringList(vec!["*".to_string()])
+}
+
 fn parse_json<T: DeserializeOwned>(value: &str) -> Result<T, String> {
     serde_json::from_str(value).map_err(|e| format!("invalid JSON object: {}", e.as_report()))
 }
@@ -380,6 +437,10 @@ pub struct FrontendArgs {
     /// `stats_update_address`.
     #[arg(long)]
     pub coordinator_address: Option<String>,
+    /// First data-parallel engine rank expected to register with this
+    /// bootstrapped frontend.
+    #[arg(long, default_value_t = 0)]
+    pub engine_start_index: u32,
     /// Total number of data-parallel engines expected for this frontend.
     #[arg(long, default_value_t = 1)]
     pub engine_count: usize,
@@ -397,6 +458,7 @@ impl FrontendArgs {
             self.input_address,
             self.output_address,
             self.coordinator_address,
+            self.engine_start_index,
             self.engine_count,
         )
     }
@@ -467,6 +529,7 @@ impl ServeArgs {
         self.managed_engine.clone().into_config(
             self.runtime.model.clone(),
             self.runtime.max_model_len,
+            self.runtime.max_logprobs,
             self.runtime.language_model_only,
             self.runtime.disable_log_stats,
             self.runtime.shutdown_timeout,

rust/src/cmd/src/cli/tests.rs

@@ -38,6 +38,7 @@ fn serve_args_forward_python_flags_with_separator() {
                         max_model_len: Some(
                             512,
                         ),
+                        max_logprobs: None,
                         grpc_port: None,
                         shutdown_timeout: 0,
                         chat_template: None,
@@ -48,6 +49,22 @@ fn serve_args_forward_python_flags_with_separator() {
                         enable_request_id_headers: false,
                         disable_log_stats: false,
                         served_model_name: [],
+                        allowed_origins: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allowed_methods: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allowed_headers: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allow_credentials: false,
                     },
                     managed_engine: ManagedEngineArgs {
                         python: "../vllm/.venv/bin/python",
@@ -134,6 +151,29 @@ fn serve_args_forward_disable_log_stats_to_managed_engine() {
     assert_eq!(config.python_args, vec!["--disable-log-stats"]);
 }
 
+#[test]
+fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--max-logprobs",
+        "-1",
+    ])
+    .unwrap();
+
+    let Command::Serve(args) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert_eq!(args.runtime.max_logprobs, Some(-1));
+
+    let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string());
+    assert_eq!(frontend_config.max_logprobs, Some(-1));
+
+    let engine_config = args.to_managed_engine_config(5555);
+    assert_eq!(engine_config.python_args, vec!["--max-logprobs", "-1"]);
+}
+
 #[test]
 fn serve_args_auto_forward_python_multi_char_alias_without_separator() {
     let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "-tp", "2"]).unwrap();
@@ -312,11 +352,17 @@ fn serve_args_reject_unknown_renderer_value() {
 
 #[test]
 fn serve_args_reject_unsupported_flag_arg() {
-    let error = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B", "--allow-credentials"])
-        .unwrap_err();
+    let error = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--ssl-keyfile",
+        "/tmp/key.pem",
+    ])
+    .unwrap_err();
 
     expect![[r#"
-        error: invalid value 'true' for '--allow-credentials [<ALLOW_CREDENTIALS>]': argument is not implemented in Rust frontend yet
+        error: invalid value '/tmp/key.pem' for '--ssl-keyfile <SSL_KEYFILE>': argument is not implemented in Rust frontend yet
 
         Remove this unsupported argument to continue.
 
@@ -324,8 +370,7 @@ fn serve_args_reject_unsupported_flag_arg() {
         This may lead to unexpected behavior as the Rust frontend will completely ignore that argument.
 
         For more information, try '--help'.
-    "#]]
-    .assert_eq(&error.to_string());
+    "#]].assert_eq(&error.to_string());
 }
 
 #[test]
@@ -379,6 +424,7 @@ fn frontend_args_accept_json() {
                     coordinator_address: Some(
                         "tcp://127.0.0.1:7000",
                     ),
+                    engine_start_index: 0,
                     engine_count: 1,
                     runtime: SharedRuntimeArgs {
                         model: "Qwen/Qwen3-0.6B",
@@ -388,6 +434,7 @@ fn frontend_args_accept_json() {
                         renderer: Auto,
                         language_model_only: false,
                         max_model_len: None,
+                        max_logprobs: None,
                         grpc_port: None,
                         shutdown_timeout: 0,
                         chat_template: None,
@@ -398,6 +445,22 @@ fn frontend_args_accept_json() {
                         enable_request_id_headers: false,
                         disable_log_stats: false,
                         served_model_name: [],
+                        allowed_origins: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allowed_methods: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allowed_headers: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allow_credentials: false,
                     },
                 },
             ),
@@ -431,6 +494,7 @@ fn frontend_args_json_applies_defaults() {
     assert_eq!(args.runtime.reasoning_parser, ParserSelection::Auto);
     assert_eq!(args.runtime.renderer, RendererSelection::Auto);
     assert_eq!(args.runtime.max_model_len, None);
+    assert_eq!(args.runtime.max_logprobs, None);
     assert_eq!(args.runtime.shutdown_timeout, 0);
 }
 
@@ -446,7 +510,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
         "--output-address",
         "ipc:///tmp/output.sock",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"shutdown_timeout":3}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","engine_ready_timeout_secs":42,"tool_call_parser":"hermes","reasoning_parser":"qwen3_thinking","tokenizer_mode":"deepseek_v32","language_model_only":true,"max_model_len":8192,"max_logprobs":-1,"shutdown_timeout":3}"#,
     ])
     .unwrap();
 
@@ -465,6 +529,7 @@ fn frontend_args_json_accepts_supported_non_default_fields() {
     assert_eq!(args.runtime.renderer, RendererSelection::DeepSeekV32);
     assert!(args.runtime.language_model_only);
     assert_eq!(args.runtime.max_model_len, Some(8192));
+    assert_eq!(args.runtime.max_logprobs, Some(-1));
     assert_eq!(args.runtime.shutdown_timeout, 3);
 }
 
@@ -531,6 +596,71 @@ fn frontend_args_json_sets_prompt_tokens_details_flag() {
     assert!(args.runtime.enable_prompt_tokens_details);
 }
 
+#[test]
+fn serve_args_parse_cors_flags() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--allowed-origins",
+        r#"["http://a.com","http://b.com"]"#,
+        "--allowed-methods",
+        r#"["GET","POST"]"#,
+        "--allow-credentials",
+    ])
+    .unwrap();
+
+    let Command::Serve(serve) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert_eq!(
+        serve.runtime.allowed_origins.0,
+        ["http://a.com", "http://b.com"]
+    );
+    assert_eq!(serve.runtime.allowed_methods.0, ["GET", "POST"]);
+    assert!(serve.runtime.allow_credentials);
+    // Unspecified lists keep the permissive default.
+    assert_eq!(serve.runtime.allowed_headers.0, ["*"]);
+}
+
+#[test]
+fn serve_args_cors_defaults_are_permissive() {
+    let cli = Cli::try_parse_from(["vllm-rs", "serve", "Qwen/Qwen3-0.6B"]).unwrap();
+
+    let Command::Serve(serve) = cli.command else {
+        panic!("expected serve args");
+    };
+    assert_eq!(serve.runtime.allowed_origins.0, ["*"]);
+    assert_eq!(serve.runtime.allowed_methods.0, ["*"]);
+    assert_eq!(serve.runtime.allowed_headers.0, ["*"]);
+    assert!(!serve.runtime.allow_credentials);
+}
+
+#[test]
+fn frontend_args_json_parses_cors_fields() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "frontend",
+        "--listen-fd",
+        "3",
+        "--input-address",
+        "ipc:///tmp/input.sock",
+        "--output-address",
+        "ipc:///tmp/output.sock",
+        "--args-json",
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","allowed_origins":["http://a.com"],"allow_credentials":true}"#,
+    ])
+    .unwrap();
+
+    let Command::Frontend(args) = cli.command else {
+        panic!("expected frontend args");
+    };
+    assert_eq!(args.runtime.allowed_origins.0, ["http://a.com"]);
+    assert!(args.runtime.allow_credentials);
+    // Unspecified lists fall back to the permissive default via serde.
+    assert_eq!(args.runtime.allowed_methods.0, ["*"]);
+}
+
 #[test]
 fn frontend_args_json_rejects_unsupported_fields() {
     let error = Cli::try_parse_from([
@@ -543,14 +673,14 @@ fn frontend_args_json_rejects_unsupported_fields() {
         "--output-address",
         "ipc:///tmp/output.sock",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","ssl_keyfile":"/tmp/key.pem"}"#,
     ])
     .unwrap_err();
 
     expect![[r#"
-        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true}' for '--args-json <JSON>': 
+        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","ssl_keyfile":"/tmp/key.pem"}' for '--args-json <JSON>': 
         The following arguments are not implemented in Rust frontend yet:
-        - allow_credentials
+        - ssl_keyfile
 
         Remove these arguments to continue.
 
@@ -570,15 +700,15 @@ fn frontend_args_json_aggregates_multiple_unsupported_fields() {
         "--output-address",
         "ipc:///tmp/output.sock",
         "--args-json",
-        r#"{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}"#,
+        r#"{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","ssl_keyfile":"/tmp/key.pem"}"#,
     ])
     .unwrap_err();
 
     let actual = error.to_string().replace(": \n", ":\n");
     expect![[r#"
-        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","allow_credentials":true,"ssl_keyfile":"/tmp/key.pem"}' for '--args-json <JSON>':
+        error: invalid value '{"model_tag":"Qwen/Qwen3-0.6B","response_role":"assistant","ssl_keyfile":"/tmp/key.pem"}' for '--args-json <JSON>':
         The following arguments are not implemented in Rust frontend yet:
-        - allow_credentials
+        - response_role
         - ssl_keyfile
 
         Remove these arguments to continue.
@@ -792,6 +922,7 @@ fn serve_args_accept_handshake_aliases() {
                         renderer: Auto,
                         language_model_only: false,
                         max_model_len: None,
+                        max_logprobs: None,
                         grpc_port: None,
                         shutdown_timeout: 0,
                         chat_template: None,
@@ -802,6 +933,22 @@ fn serve_args_accept_handshake_aliases() {
                         enable_request_id_headers: false,
                         disable_log_stats: false,
                         served_model_name: [],
+                        allowed_origins: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allowed_methods: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allowed_headers: JsonStringList(
+                            [
+                                "*",
+                            ],
+                        ),
+                        allow_credentials: false,
                     },
                     managed_engine: ManagedEngineArgs {
                         python: "python3",
@@ -917,11 +1064,24 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
             chat_template: None,
             default_chat_template_kwargs: None,
             chat_template_content_format: Auto,
+            max_logprobs: None,
             api_server_options: ApiServerOptions {
                 enable_log_requests: false,
                 enable_prompt_tokens_details: false,
                 enable_request_id_headers: false,
             },
+            cors: CorsConfig {
+                allow_origins: [
+                    "*",
+                ],
+                allow_methods: [
+                    "*",
+                ],
+                allow_headers: [
+                    "*",
+                ],
+                allow_credentials: false,
+            },
             api_keys: [],
             disable_log_stats: false,
             grpc_port: None,
@@ -985,11 +1145,24 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
             chat_template: None,
             default_chat_template_kwargs: None,
             chat_template_content_format: Auto,
+            max_logprobs: None,
             api_server_options: ApiServerOptions {
                 enable_log_requests: false,
                 enable_prompt_tokens_details: false,
                 enable_request_id_headers: false,
             },
+            cors: CorsConfig {
+                allow_origins: [
+                    "*",
+                ],
+                allow_methods: [
+                    "*",
+                ],
+                allow_headers: [
+                    "*",
+                ],
+                allow_credentials: false,
+            },
             api_keys: [],
             disable_log_stats: false,
             grpc_port: None,
@@ -1033,8 +1206,10 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
         "ipc:///tmp/output.sock",
         "--coordinator-address",
         "tcp://127.0.0.1:7000",
+        "--engine-start-index",
+        "3",
         "--engine-count",
-        "2",
+        "1",
         "--args-json",
         r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#,
     ])
@@ -1050,7 +1225,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
             transport_mode: Bootstrapped {
                 input_address: "ipc:///tmp/input.sock",
                 output_address: "ipc:///tmp/output.sock",
-                engine_count: 2,
+                engine_start_index: 3,
+                engine_count: 1,
                 ready_timeout: 600s,
             },
             coordinator_mode: External {
@@ -1068,11 +1244,24 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
             chat_template: None,
             default_chat_template_kwargs: None,
             chat_template_content_format: Auto,
+            max_logprobs: None,
             api_server_options: ApiServerOptions {
                 enable_log_requests: false,
                 enable_prompt_tokens_details: false,
                 enable_request_id_headers: false,
             },
+            cors: CorsConfig {
+                allow_origins: [
+                    "*",
+                ],
+                allow_methods: [
+                    "*",
+                ],
+                allow_headers: [
+                    "*",
+                ],
+                allow_credentials: false,
+            },
             api_keys: [],
             disable_log_stats: false,
             grpc_port: None,

rust/src/cmd/src/cli/unsupported.rs

@@ -202,14 +202,6 @@ pub struct EngineUnsupportedArgs {
     #[arg(long)]
     pub tokenizer_revision: Option<Unsupported>,
 
-    /// Maximum number of log probabilities to return when `logprobs` is
-    /// specified in `SamplingParams`. The default value comes the default for
-    /// the OpenAI Chat Completions API. -1 means no cap, i.e. all
-    /// (output_length * vocab_size) logprobs are allowed to be returned and
-    /// it may cause OOM.
-    #[arg(long)]
-    pub max_logprobs: Option<Unsupported>,
-
     /// Skip initialization of tokenizer and detokenizer. Expects valid
     /// `prompt_token_ids` and `None` for prompt from the input. The generated
     /// output will contain token ids.
@@ -534,27 +526,6 @@ pub struct ServerUnsupportedArgs {
     #[arg(long)]
     pub disable_access_log_for_endpoints: Option<Noop>,
 
-    /// Allow credentials.
-    #[arg(
-        long,
-        visible_alias = "no-allow-credentials",
-        default_missing_value = "true",
-        num_args = 0..=1
-    )]
-    pub allow_credentials: Option<Unsupported>,
-
-    /// Allowed origins.
-    #[arg(long)]
-    pub allowed_origins: Option<Unsupported>,
-
-    /// Allowed methods.
-    #[arg(long)]
-    pub allowed_methods: Option<Unsupported>,
-
-    /// Allowed headers.
-    #[arg(long)]
-    pub allowed_headers: Option<Unsupported>,
-
     /// The file path to the SSL key file.
     #[arg(long)]
     pub ssl_keyfile: Option<Unsupported>,

rust/src/engine-core-client/src/client.rs

@@ -56,6 +56,9 @@ pub enum TransportMode {
         /// Output PULL socket address that engines will connect to for
         /// responses.
         output_address: String,
+        /// First data-parallel engine rank expected to register on this
+        /// transport.
+        engine_start_index: u32,
         /// Total number of engines expected to register on this transport.
         engine_count: usize,
         /// Maximum time to wait for all expected engines to register.
@@ -246,6 +249,7 @@ impl EngineCoreClient {
             TransportMode::Bootstrapped {
                 input_address,
                 output_address,
+                engine_start_index,
                 engine_count,
                 ready_timeout,
             } => {
@@ -256,6 +260,7 @@ impl EngineCoreClient {
                 transport::connect_bootstrapped(
                     input_address,
                     output_address,
+                    *engine_start_index,
                     *engine_count,
                     *ready_timeout,
                 )
@@ -409,6 +414,24 @@ impl EngineCoreClient {
             .expect("engine core client requires at least one engine")
     }
 
+    /// Return the world size (TP * PP) from the parallel config, if available.
+    pub fn world_size(&self) -> u64 {
+        self.engines
+            .first()
+            .expect("engine core client requires at least one engine")
+            .ready_response
+            .world_size
+    }
+
+    /// Return the data parallel size from the parallel config, if available.
+    pub fn data_parallel_size(&self) -> u64 {
+        self.engines
+            .first()
+            .expect("engine core client requires at least one engine")
+            .ready_response
+            .data_parallel_size
+    }
+
     /// Get the model name associated with this client used for metrics
     /// labeling.
     pub fn model_name(&self) -> &str {
@@ -490,6 +513,10 @@ impl EngineCoreClient {
             return Ok(());
         }
 
+        // Finalize the consumer streams first, before the engine round-trip.
+        let all_request_ids: Vec<String> = abortable.values().flatten().cloned().collect();
+        self.inner.abort_requests_locally(&all_request_ids);
+
         for (engine_id, request_ids) in abortable {
             self.inner.do_abort_requests(&engine_id, &request_ids).await?;
         }

rust/src/engine-core-client/src/client/imp.rs

@@ -1,5 +1,6 @@
 use std::collections::{BTreeMap, BTreeSet};
 use std::sync::Arc;
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use arc_swap::ArcSwapOption;
 use parking_lot::Mutex;
@@ -126,6 +127,20 @@ impl ClientInner {
         self.request_reg.lock().finish_many(request_ids)
     }
 
+    /// Finalize client-initiated aborts by pushing a terminal `Abort` output
+    /// down each request's stream and removing it from the registry. Returns
+    /// the request ids that were still active. See [`RequestRegistry::abort_many`].
+    pub fn abort_requests_locally<'a>(
+        &self,
+        request_ids: impl IntoIterator<Item = &'a String>,
+    ) -> Vec<String> {
+        let timestamp = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .map(|d| d.as_secs_f64())
+            .unwrap_or(0.0);
+        self.request_reg.lock().abort_many(request_ids, timestamp)
+    }
+
     /// Apply one scheduler stats update for the given engine to the local
     /// routing state. Returns `false` if the engine is unknown to the
     /// client.

rust/src/engine-core-client/src/client/state.rs

@@ -9,7 +9,7 @@ use crate::client::stream::EngineCoreStreamOutput;
 use crate::error::{Error, Result};
 use crate::protocol::stats::SchedulerStats;
 use crate::protocol::utility::UtilityOutput;
-use crate::protocol::{EngineCoreEventType, EngineCoreOutput};
+use crate::protocol::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
 use crate::transport::ConnectedEngine;
 
 pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
@@ -289,6 +289,34 @@ impl RequestRegistry {
             .collect()
     }
 
+    /// Finalize client-initiated aborts: remove each request and push a
+    /// terminal output with `finish_reason = Abort` down its stream before the
+    /// sender drops. Returns the request ids that were still active.
+    pub fn abort_many<'a>(
+        &mut self,
+        request_ids: impl IntoIterator<Item = &'a String>,
+        timestamp: f64,
+    ) -> Vec<String> {
+        let mut aborted = Vec::new();
+        for request_id in request_ids {
+            let Some((sender, engine_id)) = self.remove(request_id) else {
+                continue;
+            };
+            let output = EngineCoreStreamOutput {
+                engine_index: engine_id.engine_index().unwrap_or(0),
+                timestamp,
+                output: EngineCoreOutput {
+                    request_id: request_id.clone(),
+                    finish_reason: Some(EngineCoreFinishReason::Abort),
+                    ..EngineCoreOutput::default()
+                },
+            };
+            let _ = sender.send(Ok(output));
+            aborted.push(request_id.clone());
+        }
+        aborted
+    }
+
     /// Remove one request from the local registry. Returns the tracked entry if
     /// it exists.
     #[must_use]

rust/src/engine-core-client/src/mock_engine.rs

@@ -15,6 +15,8 @@ use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
 pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
 /// Default KV block count advertised by reusable mock engine helpers.
 pub const DEFAULT_MOCK_NUM_GPU_BLOCKS: u64 = 0;
+/// Default KV block size (tokens per block)
+pub const DEFAULT_MOCK_BLOCK_SIZE: u64 = 16;
 
 /// Startup behavior for one mock engine joining a frontend.
 #[derive(Debug, Clone)]
@@ -46,9 +48,14 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
     EngineCoreReadyResponse {
         max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
         num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
+        block_size: DEFAULT_MOCK_BLOCK_SIZE,
         dp_stats_address: None,
         dtype: ModelDtype::Float32,
         vllm_version: "test-vllm-version".to_string(),
+        world_size: 1,
+        data_parallel_size: 1,
+        kv_cache_size_tokens: None,
+        kv_cache_max_concurrency: None,
     }
 }
 

rust/src/engine-core-client/src/protocol/handshake.rs

@@ -28,7 +28,7 @@ pub struct ReadyMessage {
 /// profiling).
 ///
 /// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/c8d98f81f6/vllm/v1/engine/__init__.py#L67-L77>
+/// <https://github.com/vllm-project/vllm/blob/c9340e6f35/vllm/v1/engine/__init__.py#L68-L80>
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct EngineCoreReadyResponse {
     /// Engine-reported maximum model context length (auto-fitted after
@@ -36,12 +36,22 @@ pub struct EngineCoreReadyResponse {
     pub max_model_len: u64,
     /// Number of GPU blocks available for KV cache on this engine.
     pub num_gpu_blocks: u64,
+    /// KV cache block size (tokens per block).
+    pub block_size: u64,
     /// DP coordinator stats publish address, if applicable.
     pub dp_stats_address: Option<String>,
     /// Effective model dtype after Python vLLM resolves `--dtype`.
     pub dtype: ModelDtype,
     /// Python vLLM version reported by the engine process.
     pub vllm_version: String,
+    /// World size (TP * PP) from the parallel config.
+    pub world_size: u64,
+    /// Data parallelism size from the parallel config.
+    pub data_parallel_size: u64,
+    /// Total KV cache capacity in tokens, if reported.
+    pub kv_cache_size_tokens: Option<u64>,
+    /// Maximum achievable request concurrency given the KV cache, if reported.
+    pub kv_cache_max_concurrency: Option<f64>,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup

rust/src/engine-core-client/src/protocol/mod.rs

@@ -44,6 +44,14 @@ fn default_repetition_penalty() -> f32 {
     1.0
 }
 
+fn default_temperature() -> f32 {
+    1.0
+}
+
+fn default_max_tokens() -> u32 {
+    16
+}
+
 mod classified_outputs;
 pub mod dtype;
 pub mod handshake;
@@ -246,24 +254,28 @@ pub struct StructuredOutputsParams {
 ///
 /// Original Python definition:
 /// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L155-L291>
+// Python's SamplingParams is `omit_defaults=True`, so msgpack drops
+// default-valued keys; default the whole struct. Per-field fns cover the
+// non-zero defaults.
 #[serde_with::skip_serializing_none]
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, DefaultFromSerde)]
+#[serde(default)]
 pub struct EngineCoreSamplingParams {
     /// Controls randomness. Lower values are more deterministic; zero means
     /// greedy sampling.
+    #[serde(default = "default_temperature")]
     pub temperature: f32,
     /// Cumulative probability threshold for nucleus sampling.
     #[serde(default = "default_top_p")]
     pub top_p: f32,
     /// Maximum number of top tokens to consider. `0` means all tokens.
-    #[serde(default)]
     pub top_k: u32,
     /// Random seed used by the sampler when present.
     pub seed: Option<i64>,
     /// Maximum number of tokens to generate per output sequence.
+    #[serde(default = "default_max_tokens")]
     pub max_tokens: u32,
     /// Minimum number of tokens to generate before EOS or stop-token handling.
-    #[serde(default)]
     pub min_tokens: u32,
     /// Number of log probabilities to return per generated token.
     ///
@@ -274,7 +286,6 @@ pub struct EngineCoreSamplingParams {
     /// `None` disables prompt logprobs. `-1` requests the full vocabulary.
     pub prompt_logprobs: Option<i32>,
     /// Minimum probability threshold for token sampling.
-    #[serde(default)]
     pub min_p: f32,
     /// Frequency penalty applied by the sampler.
     pub frequency_penalty: f32,
@@ -301,16 +312,13 @@ pub struct EngineCoreSamplingParams {
     pub all_stop_token_ids: BTreeSet<u32>,
     /// Logit biases to apply during sampling.
     /// Keys are token IDs
-    #[serde(default)]
     pub logit_bias: Option<HashMap<u32, f32>>,
     /// Restrict output to these token IDs only.
-    #[serde(default)]
     pub allowed_token_ids: Option<Vec<u32>>,
     /// Tokenized bad words to avoid during generation.
-    #[serde(default, rename = "_bad_words_token_ids")]
+    #[serde(rename = "_bad_words_token_ids")]
     pub bad_words_token_ids: Option<Vec<Vec<u32>>>,
     /// Parameters for configuring structured outputs (guided decoding).
-    #[serde(default)]
     pub structured_outputs: Option<StructuredOutputsParams>,
     /// Specific token IDs for which log probabilities should be returned at
     /// each position.
@@ -318,15 +326,12 @@ pub struct EngineCoreSamplingParams {
     /// When set, the engine returns logprobs for exactly these tokens in
     /// addition to the sampled/scored token. Mutually exclusive with the
     /// `logprobs` count field in practice.
-    #[serde(default)]
     pub logprob_token_ids: Option<Vec<u32>>,
     /// If `Some(true)`, the request will not attempt to read from the prefix
     /// cache; newly computed blocks may still populate the cache. `None`
     /// defers to engine-core defaults.
-    #[serde(default)]
     pub skip_reading_prefix_cache: Option<bool>,
     /// Additional request parameters for custom extensions (from `vllm_xargs`).
-    #[serde(default)]
     pub extra_args: Option<HashMap<String, serde_json::Value>>,
 }
 
@@ -640,4 +645,58 @@ mod tests {
         let value = serde_json::to_value(params).unwrap();
         assert_eq!(value["_backend"], "guidance");
     }
+
+    /// A real `sampling_params` is a sparse `omit_defaults` map; absent fields
+    /// must fall back to defaults. `python_compat` can't catch this since Rust
+    /// encodes full maps (see `engine_core_request_serializes_as_full_array`).
+    #[test]
+    fn decodes_sampling_params_with_omitted_defaults() {
+        let sampling_params = Value::Map(vec![
+            (
+                Value::from("stop_token_ids"),
+                Value::Array(vec![Value::from(151643u32)]),
+            ),
+            (Value::from("skip_reading_prefix_cache"), Value::from(false)),
+        ]);
+        let request = Value::Array(vec![
+            Value::from("req-omit-defaults"),
+            Value::Array(vec![
+                Value::from(1u32),
+                Value::from(2u32),
+                Value::from(3u32),
+            ]),
+            Value::Nil,
+            sampling_params,
+            Value::Nil,
+            Value::from(1.0f64),
+        ]);
+
+        let mut bytes = Vec::new();
+        rmpv::encode::write_value(&mut bytes, &request).unwrap();
+
+        let decoded: EngineCoreRequest = decode_msgpack(&bytes)
+            .expect("a real omit_defaults request must decode (regression: missing field)");
+
+        assert_eq!(decoded.request_id, "req-omit-defaults");
+        let sampling = decoded.sampling_params.expect("sampling params present");
+
+        assert_eq!(sampling.stop_token_ids, vec![151643]);
+        assert_eq!(sampling.skip_reading_prefix_cache, Some(false));
+
+        // Omitted fields -> Python defaults.
+        assert_eq!(sampling.temperature, 1.0);
+        assert_eq!(sampling.top_p, 1.0);
+        assert_eq!(sampling.top_k, 0);
+        assert_eq!(sampling.seed, None);
+        assert_eq!(sampling.max_tokens, 16);
+        assert_eq!(sampling.min_tokens, 0);
+        assert_eq!(sampling.min_p, 0.0);
+        assert_eq!(sampling.frequency_penalty, 0.0);
+        assert_eq!(sampling.presence_penalty, 0.0);
+        assert_eq!(sampling.repetition_penalty, 1.0);
+        assert_eq!(sampling.logprobs, None);
+        assert_eq!(sampling.prompt_logprobs, None);
+        assert_eq!(sampling.eos_token_id, None);
+        assert!(sampling.all_stop_token_ids.is_empty());
+    }
 }