use newtype for logprobs count (-1 for all)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

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

@@ -0,0 +1,80 @@
+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
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
+    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
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
+    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
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
+    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/basic_correctness.yaml

@@ -5,6 +5,10 @@ steps:
 - label: XPU Sleep Mode
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/engine_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: Engine (1 GPU)
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/expert_parallelism_intel.yaml

@@ -6,6 +6,10 @@ steps:
   key: eplb-algorithm
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/kernels_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: vLLM IR Tests
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/lora_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: LoRA Runtime + Utils
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -34,6 +38,10 @@ steps:
 - label: LoRA Fused/MoE Kernels
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -54,6 +62,10 @@ steps:
 - label: LoRA Punica Kernels
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -74,6 +86,10 @@ steps:
 - label: LoRA Punica FP8/XPU Ops
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -94,6 +110,10 @@ steps:
 - label: LoRA Models
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -117,6 +137,10 @@ steps:
 - label: LoRA Multimodal
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/misc_intel.yaml

@@ -5,6 +5,10 @@ steps:
 - label: V1 Core + KV + Metrics
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -31,6 +35,10 @@ steps:
 - label: V1 Sample + Logits
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -57,17 +65,24 @@ 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
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -92,6 +107,10 @@ steps:
   key: regression
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -123,6 +142,10 @@ steps:
   timeout_in_minutes: 30
   num_devices: 2
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -154,6 +177,10 @@ steps:
   key: async-engine-inputs-utils-worker
   timeout_in_minutes: 30
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/model_runner_v2_intel.yaml

@@ -0,0 +1,62 @@
+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
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
+  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
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
+  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/models_multimodal_intel.yaml

@@ -6,6 +6,10 @@ steps:
   key: multi-modal-models-standard-1-qwen2
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -27,6 +31,10 @@ steps:
   key: multi-modal-models-standard-2-qwen3-gemma
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -47,6 +55,10 @@ steps:
   key: multi-modal-models-standard-3-llava-qwen2-vl
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
   no_plugin: true
   working_dir: "."
   env:
@@ -68,6 +80,10 @@ steps:
   key: multi-modal-models-standard-4-other-whisper
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:
@@ -88,6 +104,10 @@ steps:
   key: multi-modal-processor
   timeout_in_minutes: 45
   device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
   no_plugin: true
   working_dir: "."
   env:

.buildkite/intel_jobs/test-intel.yaml

@@ -19,6 +19,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -49,6 +53,10 @@ steps:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -60,19 +68,24 @@ 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" &&
         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/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 --ignore=v1/spec_decode/test_speculators_correctness.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'
   - label: "XPU server test"
     depends_on:
       - image-build-xpu
     timeout_in_minutes: 30
     device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
     no_plugin: true
     env:
       REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -87,3 +100,24 @@ 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
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
+    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/scripts/install-kv-connectors.sh

@@ -4,6 +4,11 @@
 
 set -euo pipefail
 
+if python3 -c "import torch; raise SystemExit(0 if torch.version.hip is not None else 1)"; then
+    uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+    exit 0
+fi
+
 REQUIREMENTS_FILE="${KV_CONNECTORS_REQUIREMENTS:-/vllm-workspace/requirements/kv_connectors.txt}"
 
 uv pip install --system -r "${REQUIREMENTS_FILE}"

.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
@@ -658,7 +647,7 @@ steps:
   - pytest -v -s v1/cudagraph/test_cudagraph_mode.py
 
 - label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 35
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
   agent_pool: mi250_1
   optional: true
@@ -1605,9 +1594,10 @@ steps:
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   parallelism: 2
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
@@ -1638,10 +1628,11 @@ steps:
   - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
 
 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
-  parallelism: 4
+  optional: true
+  parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/
@@ -1694,6 +1685,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 +1783,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 +1797,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
@@ -2074,19 +2077,6 @@ steps:
   - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
   - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
 
-- label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/
-  - tests/v1/e2e/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
-
 - label: e2e Scheduling (1 GPU) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2132,9 +2122,10 @@ steps:
   - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
 
 - label: Spec Decode Eagle # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_1
+  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - vllm/v1/spec_decode/
@@ -2754,7 +2745,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  --------------------------------------------------------#
 
@@ -3052,7 +3043,7 @@ steps:
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#
 
 - label: Kernels (B200-MI355) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 15
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   working_dir: "/vllm-workspace/"
@@ -3076,11 +3067,10 @@ steps:
   - pytest -v -s tests/kernels/attention/test_attention_selector.py
 
 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 60
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
   parallelism: 2
-  optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/attention/
@@ -3094,10 +3084,10 @@ steps:
   - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
 
 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_1
-  parallelism: 4
+  parallelism: 5
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
   - csrc/quantization/cutlass_w8a8/moe/

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

@@ -74,6 +74,16 @@ steps:
     - tests/v1/e2e/general/
   commands:
     - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+  mirror:
+    amd:
+      device: mi250_1
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/
+      - tests/v1/e2e/general/
+      - vllm/platforms/rocm.py
 
 - label: V1 e2e (2 GPUs)
   key: v1-e2e-2-gpus

.buildkite/test_areas/kernels.yaml

@@ -74,6 +74,20 @@ steps:
   commands:
     - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py
 
 - label: Kernels Attention DiffKV Test (H100)
   key: kernels-attention-diffkv-test-h100
@@ -104,6 +118,7 @@ steps:
       source_file_dependencies:
       - csrc/quantization/
       - vllm/model_executor/layers/quantization
+      - vllm/config/
       - tests/kernels/quantization
       - tests/kernels/quantization/test_rocm_skinny_gemms.py
       - vllm/_aiter_ops.py
@@ -127,6 +142,22 @@ steps:
     - pytest -v -s kernels/moe --ignore=kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
     - pytest -v -s kernels/moe/test_modular_oai_triton_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
   parallelism: 5
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 50
+      source_file_dependencies:
+      - csrc/quantization/cutlass_w8a8/moe/
+      - csrc/moe/
+      - tests/kernels/moe
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/distributed/device_communicators/
+      - vllm/envs.py
+      - vllm/config
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd
 
 - label: Kernels Mamba Test
   key: kernels-mamba-test

.buildkite/test_areas/lm_eval.yaml

@@ -101,6 +101,16 @@ steps:
   num_devices: 8
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
+  mirror:
+    amd:
+      device: mi300_8
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
 
 - label: MoE Refactor Integration Test (H100 - TEMPORARY)
   key: moe-refactor-integration-test-h100-temporary

.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
@@ -99,6 +105,12 @@ steps:
     # Integration test for streaming correctness (requires special branch).
     - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
     - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd
 
 - label: V1 Others (CPU)
   key: v1-others-cpu

.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,15 @@ 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:
+      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/pytorch.yaml

@@ -107,6 +107,12 @@ steps:
     - tests/compile/passes
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
 
 - label: PyTorch Fullgraph Smoke Test
   key: pytorch-fullgraph-smoke-test

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

.buildkite/test_areas/spec_decode.yaml

@@ -12,6 +12,20 @@ steps:
     - tests/v1/e2e/spec_decode/
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/worker/gpu/spec_decode/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/sample/
+      - vllm/model_executor/layers/
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py
 
 - label: Spec Decode Eagle Nightly B200
   key: spec-decode-eagle-nightly-b200

.claude/skills/ci-fails-buildkite/SKILL.md

@@ -0,0 +1,35 @@
+---
+name: ci-fails-buildkite
+description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
+---
+
+# Diagnosing vLLM Buildkite CI Failures
+
+Buildkite logs are public; no login needed.
+
+`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
+
+## Fetching logs
+
+```bash
+# All failed jobs in a PR's latest build (current branch's PR if omitted):
+.buildkite/scripts/ci-fetch-log.sh --pr <PR>
+
+# All failed jobs in a build (--soft also includes soft-failed jobs;
+# --all fetches every finished job):
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
+
+# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
+# work; pass "-" as a second argument to stream to stdout:
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
+```
+
+To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
+
+```bash
+./ci-clean-log.sh ci.log
+```
+
+## Reference
+
+See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.

.github/CODEOWNERS

@@ -2,15 +2,14 @@
 # for more info about CODEOWNERS file
 
 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
-/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
-/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
+/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/model_loader @22quinn
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/ir @ProExpertProg
@@ -120,16 +119,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

@@ -199,7 +199,9 @@ cython_debug/
 .vscode/
 
 # Claude
-.claude/
+.claude/*
+!.claude/skills/
+!.claude/skills/**
 
 # Codex
 .codex/

AGENTS.md

@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
 - Keep comments and docstrings minimal and concise.
 - Assume the reader is familiar with vLLM.
 
-### Diagnosing CI failures
-
-Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
-
-```bash
-# All failed-job logs for a PR's latest build (current branch's PR if omitted):
-.buildkite/scripts/ci-fetch-log.sh --pr <PR>
-# Any Buildkite build or job URL also works:
-.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
-```
-
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:

CMakeLists.txt

@@ -319,111 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
 endif()
 
 #
-# _C extension
+# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
 #
 
-set(VLLM_EXT_SRC
-  "csrc/quantization/activation_kernels.cu"
-  "csrc/torch_bindings.cpp")
-
-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
-
-  # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
-  set(CUTLASS_REVISION "v4.4.2")
-
-  # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
-  if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
-    set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
-  endif()
-
-  if(VLLM_CUTLASS_SRC_DIR)
-    if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
-      get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
-    endif()
-    message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
-    FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
-  else()
-    FetchContent_Declare(
-        cutlass
-        GIT_REPOSITORY https://github.com/nvidia/cutlass.git
-        # Please keep this in sync with CUTLASS_REVISION line above.
-        GIT_TAG ${CUTLASS_REVISION}
-        GIT_PROGRESS TRUE
-
-        # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
-        # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
-        # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
-        GIT_SHALLOW TRUE
-    )
-  endif()
-  FetchContent_MakeAvailable(cutlass)
-
-  set_gencode_flags_for_srcs(
-    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()
-
-if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
-  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
-  list(APPEND VLLM_EXT_SRC
+if(VLLM_GPU_LANG STREQUAL "HIP")
+  set(VLLM_EXT_SRC
+    "csrc/torch_bindings.cpp"
     "csrc/custom_quickreduce.cu"
     "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu"
-  )
-# if ROCM endif
-endif()
+    "csrc/libtorch_stable/cuda_utils_kernels.cu")
 
-message(STATUS "Enabling C extension.")
-define_extension_target(
-  _C
-  DESTINATION vllm
-  LANGUAGE ${VLLM_GPU_LANG}
-  SOURCES ${VLLM_EXT_SRC}
-  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-  ARCHITECTURES ${VLLM_GPU_ARCHES}
-  INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
-  INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
-  USE_SABI 3
-  WITH_SOABI)
+  message(STATUS "Enabling C extension.")
+  define_extension_target(
+    _C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${VLLM_EXT_SRC}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
+    INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
+    USE_SABI 3
+    WITH_SOABI)
 
-# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
-# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
-# driver API. This causes problems when linking with earlier versions of CUDA.
-# Setting this variable sidesteps the issue by calling the driver directly.
-target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  # If CUTLASS is compiled on NVCC >= 12.5, it by default uses
+  # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
+  # driver API. This causes problems when linking with earlier versions of CUDA.
+  # Setting this variable sidesteps the issue by calling the driver directly.
+  target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+endif() # _C HIP endif
 
 if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   #
@@ -432,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
   set(VLLM_STABLE_EXT_SRC
     "csrc/libtorch_stable/torch_bindings.cpp"
     "csrc/libtorch_stable/activation_kernels.cu"
+    "csrc/libtorch_stable/quantization/activation_kernels.cu"
     "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
     "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -458,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
+    SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
+
+    # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
+    set(CUTLASS_REVISION "v4.4.2")
+
+    # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
+    if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
+      set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
+    endif()
+
+    if(VLLM_CUTLASS_SRC_DIR)
+      if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
+        get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
+      endif()
+      message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
+      FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
+    else()
+      FetchContent_Declare(
+          cutlass
+          GIT_REPOSITORY https://github.com/nvidia/cutlass.git
+          # Please keep this in sync with CUTLASS_REVISION line above.
+          GIT_TAG ${CUTLASS_REVISION}
+          GIT_PROGRESS TRUE
+
+          # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
+          # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
+          # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
+          GIT_SHALLOW TRUE
+      )
+    endif()
+    FetchContent_MakeAvailable(cutlass)
+
     list(APPEND VLLM_STABLE_EXT_SRC
       "csrc/libtorch_stable/cuda_view.cu"
       "csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -958,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${FP4_SM120_SRCS}"
       CUDA_ARCHS "${FP4_SM120_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
     message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -991,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${FP4_SM100_SRCS}"
       CUDA_ARCHS "${FP4_SM100_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
     message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")

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

@@ -32,19 +32,17 @@ 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_SOURCE_DIR}/python/fmha_sm100/__init__.py"
-  "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100/sparse.py"
+  "${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_SOURCE_DIR}/python/fmha_sm100/cute/"
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
   DESTINATION vllm/third_party/fmha_sm100/cute
   COMPONENT fmha_sm100
-  FILES_MATCHING
-    REGEX "/__pycache__(/.*)?$" EXCLUDE
-    REGEX ".*\\.pyc$" EXCLUDE
-    PATTERN "example.py" EXCLUDE
-    PATTERN "test_*.py" EXCLUDE
-    PATTERN "*.py"
-    PATTERN "build_k2q_csr.cu")
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)

cmake/external_projects/qutlass.cmake

@@ -60,6 +60,7 @@ endif()
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
   set(QUTLASS_SOURCES
+    csrc/qutlass_registration.cpp
     ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
     ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
     ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
 
   if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
     list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
+    if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
+       EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
+      list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
+    else()
+      get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
+      if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
+        list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
+      endif()
+    endif()
   elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
-    list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
+    list(APPEND QUTLASS_INCLUDES
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
     message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
   else()
     message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
     CUDA_ARCHS "${QUTLASS_ARCHS}"
   )
 
-  target_sources(_C PRIVATE ${QUTLASS_SOURCES})
-  target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
-  target_compile_definitions(_C PRIVATE
+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
+  define_extension_target(
+    _qutlass_C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${QUTLASS_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_definitions(_qutlass_C PRIVATE
     QUTLASS_DISABLE_PYBIND=1
     TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-  )
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
 
   set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
     $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
       "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
       "CUDA_ARCHS='${CUDA_ARCHS}'.")
   endif()
+  add_custom_target(_qutlass_C)
 endif()

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/attention/mla/sm100_cutlass_mla_kernel.cu

@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
   using TileShapeD = typename MlaSm100Type::TileShapeD;
   arguments.problem_shape =
       cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
-  // Assumes device 0 when getting sm_count.
-  arguments.hw_info.sm_count =
-      sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
+  if (sm_count <= 0) {
+    int current_device = 0;
+    cudaGetDevice(&current_device);
+    arguments.hw_info.sm_count =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
+  } else {
+    arguments.hw_info.sm_count = sm_count;
+  }
   arguments.split_kv = static_cast<int>(num_kv_splits);
   MlaSm100Type::Fmha::set_split_kv(arguments);
 

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

@@ -58,8 +58,15 @@
 
 #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
@@ -186,6 +193,21 @@ __device__ __forceinline__ void storeElems(
   *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
 // ────────────────────────────────────────────────────────────────────────────
@@ -202,7 +224,8 @@ __device__ __forceinline__ void storeElems(
 //     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, bool kIsSparse, bool kInsertKV>
+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
@@ -215,7 +238,7 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
     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
-    scalar_t* __restrict__ kv_cache,     // [nb,2,bs,nkv,128] or 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,
@@ -355,7 +378,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
           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;
-          storeElems<scalar_t>(kv_cache + off + dim_base, elems);
+          storeCacheElems<scalar_t, cache_t, kv_dt>(kv_cache + off + dim_base,
+                                                    elems);
         }
       }
     }
@@ -373,13 +397,13 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 // ────────────────────────────────────────────────────────────────────────────
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t>
+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, scalar_t* kv_cache,
+                          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,
@@ -419,7 +443,8 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
   #define LAUNCH(IS_SPARSE, INSERT)                                           \
     cudaLaunchKernelEx(                                                       \
         &config,                                                              \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>,   \
+        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,   \
@@ -428,7 +453,8 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
   // ROCm: standard kernel launch syntax (no PDL/stream serialization).
   // clang-format off
   #define LAUNCH(IS_SPARSE, INSERT)                                          \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, 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,               \
@@ -455,6 +481,33 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
 }  // 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
 // ────────────────────────────────────────────────────────────────────────────
@@ -475,9 +528,14 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     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
+        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,
@@ -510,6 +568,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   // (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;
@@ -552,8 +612,14 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                  torch::headeronly::ScalarType::Long &&
              index_slot_mapping->numel() == slot_mapping->numel()),
         "index_slot_mapping must be int64 CUDA with slot_mapping length");
-    STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
-                    "kv_cache dtype must match qkv (bf16 cache only)");
+    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");
@@ -601,35 +667,9 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
   VLLM_STABLE_DISPATCH_HALF_TYPES(
       qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
         using st = scalar_t;
-        vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
-            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<st*>(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);
+        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];
+      if constexpr (HasWeight) {
+        float w = static_cast<float>(src2.val[j]);
+        dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
+      } else {
+        dst.val[j] = static_cast<scalar_t>(x * s_variance);
+      }
     }
     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,22 @@ 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]);
+        float wf = Converter::convert(w.data[j]);
+        out.data[j] = Converter::convert(x * s_variance * wf);
+      }
+    } 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 +168,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 +198,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) {
+      float w = (float)weight[idx];
+      input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
+    } 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 +238,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 +315,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/layernorm_quant_kernels.cu

@@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      // Multiply in weight's native dtype to match rms_norm_kernel.
-      scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
+      float w = static_cast<float>(src2.val[j]);
+      // Round normalized result through scalar_t to match the precision of the
+      // unfused composite (rms_norm writes scalar_t, then
+      // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
+      // Without this round, the fused path is strictly more accurate and
+      // disagrees with the composite at exact E4M3 quantization tie boundaries.
+      scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
       out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
           scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
                                                 scale_inv);
@@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 #pragma unroll
     for (int i = 0; i < width; ++i) {
       float x = Converter::convert(res.data[i]);
-      // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
-      HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
+      float wf = Converter::convert(w.data[i]);
+      // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
+      // to match the unfused composite path at FP8 boundaries. We use the
+      // backend's hip_type for the intermediate since c10::Half/BFloat16 has
+      // ambiguous conversions on CUDA and no implicit conversion on ROCm.
+      HipT out_norm_h = Converter::convert(x * s_variance * wf);
       out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
           Converter::convert(out_norm_h), scale_inv);
     }
@@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
-    scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
+    float w = (float)weight[idx];
+    // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
+    // to match the unfused composite path at FP8 boundaries.
+    scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
     out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
         static_cast<float>(out_norm), scale_inv);
   }

csrc/libtorch_stable/moe/moe_align_sum_kernels.cu

@@ -9,7 +9,7 @@
 #include <torch/headeronly/core/ScalarType.h>
 
 #include "../../cuda_compat.h"
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/dispatch_utils.h"
 #include "libtorch_stable/torch_utils.h"
 

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

@@ -2,6 +2,25 @@
 
 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/util/Exception.h>
+
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <torch/csrc/stable/ops.h>
+
+inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
+  // Ensure tensor is on CUDA
+  STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
+
+  // Get the raw data pointer
+  void* data_ptr = tensor.mutable_data_ptr();
+
+  /// Create a new tensor from the raw data pointer
+  return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
+                                  tensor.device(), tensor.scalar_type());
+}
 
 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                                torch::stable::Tensor& output_q,
@@ -185,11 +204,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,
@@ -297,7 +317,8 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
     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);
+    std::optional<torch::stable::Tensor> index_q_out,
+    const std::string& kv_cache_dtype);
 
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
@@ -366,6 +387,18 @@ 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,
                         double alpha = 1.0, double beta = 0.0);
+
+void silu_and_mul_quant(torch::stable::Tensor& out,
+                        torch::stable::Tensor& input,
+                        torch::stable::Tensor& scale);
+
+void persistent_masked_m_silu_mul_quant(
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
+    bool use_ue8m0);
+
 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/activation_kernels.cu

@@ -1,16 +1,12 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "libtorch_stable/torch_utils.h"
 
 #include <cmath>
-#include "core/math.hpp"
-#include "../cuda_compat.h"
-#include "dispatch_utils.h"
 
+#include "libtorch_stable/core/math.hpp"
+#include "cuda_compat.h"
+#include "libtorch_stable/dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"
 
-#include <c10/util/Float8_e4m3fn.h>
-
 #ifndef USE_ROCM
   #include <cuda_bf16.h>
   #include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
   #endif
 #endif
 
-#include "core/registration.h"
 namespace vllm {
 
 template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
 }  // namespace vllm
 
 // Launch activation, gating, and quantize kernel.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                               \
-  int d = input.size(-1) / 2;                                               \
-  int64_t num_tokens = input.numel() / input.size(-1);                      \
-  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);     \
-  dim3 block(std::min(d, 512));                                             \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
-  VLLM_DISPATCH_FLOATING_TYPES(                                             \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                      \
-        VLLM_DISPATCH_FP8_TYPES(                                            \
-            out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] {  \
-              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,    \
-                                             fp8_t>                         \
-                  <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),       \
-                                               input.data_ptr<scalar_t>(),  \
-                                               scale.data_ptr<float>(), d); \
-            });                                                             \
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                             \
+  int d = input.size(-1) / 2;                                             \
+  int64_t num_tokens = input.numel() / input.size(-1);                    \
+  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);   \
+  dim3 block(std::min(d, 512));                                           \
+  const torch::stable::accelerator::DeviceGuard device_guard(             \
+      input.get_device_index());                                          \
+  const cudaStream_t stream =                                             \
+      get_current_cuda_stream(input.get_device_index());                  \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                    \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                    \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                   \
+            out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
+              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,  \
+                                             fp8_t>                       \
+                  <<<grid, block, 0, stream>>>(                           \
+                      out.mutable_data_ptr<fp8_t>(),                      \
+                      input.const_data_ptr<scalar_t>(),                   \
+                      scale.const_data_ptr<float>(), d);                  \
+            });                                                           \
       });
 
-void silu_and_mul_quant(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        torch::Tensor& scale) {
-  TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
-              out.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.dtype() == torch::kFloat16 ||
-              input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(input.size(-1) % 2 == 0);
+void silu_and_mul_quant(torch::stable::Tensor& out,    // [..., d]
+                        torch::stable::Tensor& input,  // [..., 2 * d]
+                        torch::stable::Tensor& scale) {
+  STD_TORCH_CHECK(
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::Half ||
+          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Input must be FP16 or BF16");
+  STD_TORCH_CHECK(input.size(-1) % 2 == 0);
   LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
 }
 
 void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,              // (E, T, 2*H)
-    const at::Tensor& tokens_per_expert,  // (E)
-    at::Tensor& y_q,                      // (E, T, H) [OUT]
-    at::Tensor& y_s,                      // (E, T, H//group_size) [OUT]
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
     bool cast_scale_ue8m0) {
 #ifndef USE_ROCM
 
@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
   // fixed GROUP_SIZE of 128.
   static constexpr int GROUP_SIZE = 128;
 
-  TORCH_CHECK(input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
-              y_q.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
+  STD_TORCH_CHECK(input.scalar_type() ==
+                  torch::headeronly::ScalarType::BFloat16);
+  STD_TORCH_CHECK(
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
 
   bool const is_packed_ue8m0 =
-      (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
+      (y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
+       cast_scale_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
+                  is_packed_ue8m0);
 
   using Idx_t = int64_t;
 
@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(
 
   int const NUM_GROUPS = H / GROUP_SIZE;
 
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
 
   // TODO: Get this from cuda_arch ?
   static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
     static constexpr int max_shared_mem_bytes =                                \
         GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2;                               \
     dim3 grid(sms), block(THREAD_COUNT);                                       \
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));          \
-    VLLM_DISPATCH_FP8_TYPES(                                                   \
+    const torch::stable::accelerator::DeviceGuard device_guard(                \
+        input.get_device_index());                                             \
+    VLLM_STABLE_DISPATCH_FP8_TYPES(                                            \
         y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] {        \
           vllm::silu_mul_fp8_quant_deep_gemm_kernel<                           \
               BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
               Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES>                           \
               <<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>(  \
-                  reinterpret_cast<__nv_bfloat16*>(input.data_ptr()),          \
-                  (fp8_t*)y_q.data_ptr(),                                      \
-                  reinterpret_cast<scale_t*>(y_s.data_ptr()),                  \
-                  reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
-                  T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,       \
+                  reinterpret_cast<const __nv_bfloat16*>(                      \
+                      input.const_data_ptr()),                                 \
+                  y_q.mutable_data_ptr<fp8_t>(),                               \
+                  reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()),          \
+                  reinterpret_cast<const int32_t*>(                            \
+                      tokens_per_expert.const_data_ptr()),                     \
+                  E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,    \
                   stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T,          \
                   STRIDE_YS_G, STRIDE_YS_P, stride_counts_e);                  \
         });
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
   Idx_t stride_ys_g = y_s.stride(2);
   Idx_t stride_ys_p = 0;
   if (!cast_scale_ue8m0) {
-    TORCH_CHECK(!is_packed_ue8m0);
+    STD_TORCH_CHECK(!is_packed_ue8m0);
     LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
                 false);
     return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
     return;
   }
 
-  TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kInt32);
+  STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);
 
   // Int32 packed ue8m0 scales tensor.
   // Let E, T, G be the number to experts, number of tokens and number of groups

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu

@@ -31,7 +31,7 @@
 
 #include "cutlass/util/packed_stride.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"
 
 using namespace cute;

csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu

@@ -31,7 +31,7 @@
 
 #include "cutlass/util/packed_stride.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"
 
 using namespace cute;

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/cutlass/c3x/cutlass_gemm_caller.cuh

@@ -19,7 +19,7 @@
 #include "cutlass/gemm/collective/collective_builder.hpp"
 #include "cutlass/util/packed_stride.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on
 

csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh

@@ -14,7 +14,7 @@
 #include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/collective/collective_builder.hpp"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on
 

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cuh

@@ -22,7 +22,7 @@
 #include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
 
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on
 

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

@@ -301,12 +301,21 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
 
   const int sf_k_local = local_group_id % kGroupsPerBlockX;
   const int row_local = local_group_id / kGroupsPerBlockX;
-  const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
-  const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
+  // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
+  const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.x * 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 +426,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.
@@ -484,34 +497,66 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
                   " is not a multiple of 4.");
   const int kx = GetGroupsPerBlockX(padded_groups_per_row);
   const int ry = 16 / kx;
-  const int64_t blocks_x = padded_groups_per_row / kx;
-  const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t sf_k_blocks = padded_groups_per_row / kx;
   const int num_threads = (kx * ry) * THREADS_PER_GROUP;
-  // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
-  STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
-                      blocks_y <= static_cast<int64_t>(INT32_MAX),
+  // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
+  constexpr int64_t kMaxGridDimYZ = 65535;
+  STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
+                      sf_k_blocks <= kMaxGridDimYZ,
                   "per_token_group_quant_8bit_packed grid too large: (",
-                  blocks_x, ", ", blocks_y, ").");
+                  row_blocks, ", ", sf_k_blocks, ").");
 
   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>(row_blocks),             \
+                            static_cast<unsigned int>(sf_k_blocks));           \
+      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>(row_blocks),                         \
+                static_cast<unsigned int>(sf_k_blocks));                       \
+      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

@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   // TODO: Remove this once ROCm upgrade to torch 2.11.
   ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
 
+  // Note about marlin kernel 'workspace' arguments:
+  // Technically these should be mutable since they are modified by the kernel.
+  // But since they are set back to zero once the kernel is finished we can
+  // hand wave and say that they have no net effect.
+  //
+  // The reason to mark 'workspace' as immutable is so that they don't interfere
+  // with using ScalarType arguments in the ops. If they are marked as mutable,
+  // pytorch throws an assert in
+  // 'torch._higher_order_ops._register_effectful_op' that prevents these
+  // kernels from being torch.compile'd.
+  // See the following document for more info on custom types and ops that use
+  // custom types:
+  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
+
   // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
   ops.def(
       "machete_supported_schedules("
@@ -308,22 +322,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 +367,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
@@ -471,7 +470,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "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) -> ()");
+      "int block_size, Tensor!? q_out, Tensor!? index_q_out, "
+      "str kv_cache_dtype) -> ()");
 
   // Apply repetition penalties to logits in-place.
   ops.def(
@@ -494,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor workspace, int k, int max_seq_len) -> ()");
 
   // Activation ops
+  ops.def(
+      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
+      "y_q, Tensor! y_s, bool use_ue8m0) -> ()");
+  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
+
   // Activation function used in SwiGLU.
   ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
 
@@ -506,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
       "float alpha=1.0, float beta=0.0) -> ()");
 
+  // SwiGLU activation with FP8 quantization.
+  ops.def(
+      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
+
   // Activation function used in GeGLU with `none` approximation.
   ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
 
@@ -704,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
 
   // Activation kernels (shared CUDA/ROCm)
+  ops.impl("persistent_masked_m_silu_mul_quant",
+           TORCH_BOX(&persistent_masked_m_silu_mul_quant));
+  ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
+  ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
   ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
   ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
   ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));

csrc/ops.h

@@ -9,36 +9,14 @@
 
 #include <vector>
 
-torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
-  // Ensure tensor is on CUDA
-  if (!tensor.is_cuda()) {
-    throw std::runtime_error("Tensor must be on CUDA device");
-  }
-
-  // Get the raw data pointer
-  void* data_ptr = tensor.data_ptr();
-
-  // Get tensor sizes and strides
-  std::vector<int64_t> sizes = tensor.sizes().vec();
-  std::vector<int64_t> strides = tensor.strides().vec();
-
-  // Get tensor options (dtype, device)
-  auto options = tensor.options();
-
-  // Create a new tensor from the raw data pointer
-  auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
-
-  return new_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
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 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);
-
-void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,   // (E, T, 2*H)
-    const at::Tensor& counts,  // (E)
-    at::Tensor& y_q,           // (E, T, H) [OUT]
-    at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
-    bool use_ue8m0);
-
 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);
 
 void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);

csrc/qutlass_registration.cpp

@@ -0,0 +1,5 @@
+#include "core/registration.h"
+
+// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
+// This stub lets Python import vllm._qutlass_C to trigger op registration.
+REGISTER_EXTENSION(_qutlass_C)

csrc/torch_bindings.cpp

@@ -20,17 +20,6 @@
 
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   // vLLM custom ops
-  //
-
-  ops.def(
-      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
-      "y_q, Tensor! y_s,"
-      "bool use_ue8m0) -> ()");
-  ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
-           &persistent_masked_m_silu_mul_quant);
-
-  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
-  ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
 
 #ifdef USE_ROCM
   // TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
   ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
            &get_cuda_view_from_cpu_tensor);
-#endif
-
-  // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
-  ops.def(
-      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
-  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
-
-  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
-  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
-  // full-cache bf16/fp8 variants).
-
-  // Quantization ops
-#ifndef USE_ROCM
-
-  // Note about marlin kernel 'workspace' arguments:
-  // Technically these should be mutable since they are modified by the kernel.
-  // But since they are set back to zero once the kernel is finished we can
-  // hand wave and say that they have no net effect.
-  //
-  // The reason to mark 'workspace' as immutable is so that they don't interfere
-  // with using ScalarType arguments in the ops. If they are marked as mutable,
-  // pytorch throws an assert in
-  // 'torch._higher_order_ops._register_effectful_op' that prevents these
-  // kernels from being torch.compile'd.
-  // See the following document for more info on custom types and ops that use
-  // custom types:
-  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
-
 #endif
 }
 

docker/Dockerfile.xpu

@@ -75,13 +75,13 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
 # Install UMD
 RUN mkdir neo && \
     cd neo && \
-    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
-    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
-    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
-    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
-    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
-    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
-    wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-core-2_2.34.4+21428_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.34.4/intel-igc-opencl-2_2.34.4+21428_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-ocloc_26.18.38308.1-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/intel-opencl-icd_26.18.38308.1-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libigdgmm12_22.10.0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/26.18.38308.1/libze-intel-gpu1_26.18.38308.1-0_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero_1.28.2+u24.04_amd64.deb && \
     dpkg -i *.deb && \
     cd .. && \
     rm -rf neo
@@ -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/contributing/model/basic.md

@@ -133,10 +133,10 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
 For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
 The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.
 
-For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
+For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
 These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).
 
-For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
+For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`.
 Please follow the same guidelines as case (2) for implementing these models.
 We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
 For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
@@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles
 Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
 It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
 Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
-Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
+Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
 The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.

docs/deployment/integrations/production-stack.md

@@ -4,7 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
 
 * **Upstream vLLM compatibility** – It wraps around upstream vLLM without modifying its code.
 * **Ease of use** – Simplified deployment via Helm charts and observability through Grafana dashboards.
-* **High performance** – Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache), among others.
+* **High performance** – Optimized for LLM workloads with features like multimodel support, model-aware and prefix-aware routing, fast vLLM bootstrapping, and KV cache offloading with [LMCache](https://github.com/LMCache/LMCache) (wired up in vLLM via `--kv-offloading-backend lmcache`; see the [LMCache examples](https://github.com/vllm-project/vllm/tree/main/examples/disaggregated/lmcache) and [docs.lmcache.ai](https://docs.lmcache.ai)), among others.
 
 If you are new to Kubernetes, don't worry: in the vLLM production stack [repo](https://github.com/vllm-project/production-stack), we provide a step-by-step [guide](https://github.com/vllm-project/production-stack/blob/main/tutorials/00-install-kubernetes-env.md) and a [short video](https://www.youtube.com/watch?v=EsTJbQtzj0g) to set up everything and get started in **4 minutes**!
 

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, 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 |
+| `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 |
@@ -214,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
@@ -252,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/disagg_prefill.md

@@ -20,7 +20,7 @@ Two main reasons:
 Now supports 9 types of connectors:
 
 - **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
-- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
+- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission. LMCache also offers a multi-process (MP) mode via `LMCacheMPConnector`, where a standalone `lmcache server` holds the KV cache shared by one or more vLLM instances; see the [LMCache examples](../../examples/disaggregated/lmcache/README.md) and the [LMCache docs](https://docs.lmcache.ai) for setup.
 - **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md). You may specify one or multiple NIXL transfer backends, such as:
 
   ```bash

docs/features/kv_offloading_usage.md

@@ -74,6 +74,7 @@ vllm serve <model> \
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
+| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
 
 ## Secondary Tiers
@@ -127,6 +128,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

@@ -203,8 +203,10 @@ the vLLM JSON config.
 ### kv_connector_extra_config
 
 - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
+- `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
 - `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

@@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup
 
 Flags: `--tool-call-parser hermes`
 
-### MiniMax Models (`minimax_m1`)
-
-Supported models:
-
-* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
-* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
-
-Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja`
-
 ### DeepSeek-V3 Models (`deepseek_v3`)
 
 Supported models:

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.xpu.inc.md

@@ -27,6 +27,7 @@ Currently, there are no pre-built XPU wheels.
 
 - First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
 - Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
+- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
 
 ```bash
 git clone https://github.com/vllm-project/vllm.git

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

@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve
 | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
 | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
 | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
-| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
+| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
 | `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
 | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |

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` | ✅︎ | ✅︎ |
@@ -442,7 +441,6 @@ th {
 | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
 | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
 | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
-| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
 | `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
 | `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
 | `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
@@ -488,8 +486,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. | ✅︎ | ✅︎ |
-| `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. | | |
 
 !!! note
@@ -596,6 +592,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
 | `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
 | `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
+| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |

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
 

docs/usage/v1_guide.md

@@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor
 Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
 `Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).
 
-Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
+Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).
 
 Please note that prefix caching is not yet supported for any of the above models.
 

examples/disaggregated/lmcache/README.md

@@ -1,10 +1,38 @@
 # LMCache Examples
 
-This folder demonstrates how to use LMCache for disaggregated prefilling, CPU offloading and KV cache sharing.
+This folder demonstrates how to use LMCache with vLLM v1 for KV cache
+offloading, disaggregated prefilling, and KV cache sharing.
 
-## 1. Disaggregated Prefill in vLLM v1
+## Integration modes
 
-This example demonstrates how to run LMCache with disaggregated prefill using NIXL on a single node.
+LMCache integrates with vLLM v1 in two ways:
+
+- **In-process mode** (`LMCacheConnectorV1`): LMCache runs inside the vLLM
+  process and is configured through environment variables or a YAML config
+  file (`LMCACHE_CONFIG_FILE`). This is the simplest way to add single-node
+  CPU/disk offloading.
+- **Multi-process (MP) mode** (`LMCacheMPConnector`): LMCache runs as a
+  standalone server (`lmcache server`) that owns the KV cache storage; one or
+  more vLLM instances connect to it. This is the recommended mode for
+  distributed KV storage and for sharing KV cache across instances. See the
+  [LMCache docs](https://docs.lmcache.ai) for the full MP setup.
+
+## 1. CPU offload (in-process)
+
+- `python cpu_offload_lmcache.py` - CPU offloading with `LMCacheConnectorV1`
+  for vLLM v1.
+
+## 2. CPU offload (multi-process)
+
+- `bash cpu_offload_lmcache_mp.sh` - CPU offloading with `LMCacheMPConnector`,
+  using a standalone `lmcache server`. vLLM provides a built-in shortcut for
+  this setup via `--kv-offloading-backend lmcache` and
+  `--kv-offloading-size <GiB>`.
+
+## 3. Disaggregated Prefill in vLLM v1
+
+This example demonstrates how to run LMCache with disaggregated prefill using
+NIXL on a single node.
 
 ### Prerequisites
 
@@ -46,15 +74,7 @@ The main script generates several log files:
 - `decoder.log` - Logs from the decode server
 - `proxy.log` - Logs from the proxy server
 
-## 2. CPU Offload Examples
+## 4. KV Cache Sharing
 
-- `python cpu_offload_lmcache.py -v v0` - CPU offloading implementation for vLLM v0
-- `python cpu_offload_lmcache.py -v v1` - CPU offloading implementation for vLLM v1
-
-## 3. KV Cache Sharing
-
-The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV caches between vLLM v1 instances.
-
-## 4. Disaggregated Prefill in vLLM v0
-
-The `disaggregated_prefill_lmcache_v0.py` provides an example of how to run disaggregated prefill in vLLM v0.
+The `kv_cache_sharing_lmcache_v1.py` example demonstrates how to share KV
+caches between vLLM v1 instances through a centralized LMCache server.

examples/disaggregated/lmcache/cpu_offload_lmcache.py

@@ -1,20 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-This file demonstrates the example usage of cpu offloading
-with LMCache in vLLM v1 or v0.
-
-Usage:
-
-    Specify vLLM version
-
-    -v v0 : Use LMCacheConnector
-            model = mistralai/Mistral-7B-Instruct-v0.2
-            (Includes enable_chunked_prefill = True)
-
-    -v v1 : Use LMCacheConnectorV1 (default)
-            model = meta-llama/Meta-Llama-3.1-8B-Instruct
-            (Without enable_chunked_prefill)
+This file demonstrates the example usage of CPU offloading
+with LMCache in vLLM v1.
 
 Note that `lmcache` is needed to run this example.
 Requirements:
@@ -23,7 +11,6 @@ Learn more about LMCache environment setup, please refer to:
 https://docs.lmcache.ai/getting_started/installation.html
 """
 
-import argparse
 import contextlib
 import os
 import time
@@ -39,8 +26,6 @@ from vllm.engine.arg_utils import EngineArgs
 
 def setup_environment_variables():
     # LMCache-related environment variables
-    # Use experimental features in LMCache
-    os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
     # LMCache is set to use 256 tokens per chunk
     os.environ["LMCACHE_CHUNK_SIZE"] = "256"
     # Enable local CPU backend in LMCache
@@ -50,9 +35,9 @@ def setup_environment_variables():
 
 
 @contextlib.contextmanager
-def build_llm_with_lmcache(lmcache_connector: str, model: str):
+def build_llm_with_lmcache(model: str):
     ktc = KVTransferConfig(
-        kv_connector=lmcache_connector,
+        kv_connector="LMCacheConnectorV1",
         kv_role="kv_both",
     )
     # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
@@ -92,23 +77,10 @@ def print_output(
     print("-" * 50)
 
 
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-v",
-        "--version",
-        choices=["v0", "v1"],
-        default="v1",
-        help="Specify vLLM version (default: v1)",
-    )
-    return parser.parse_args()
-
-
 def main():
-    lmcache_connector = "LMCacheConnectorV1"
     model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
     setup_environment_variables()
-    with build_llm_with_lmcache(lmcache_connector, model) as llm:
+    with build_llm_with_lmcache(model) as llm:
         # This example script runs two requests with a shared prefix.
         # Define the shared prompt and specific prompts
         shared_prompt = "Hello, how are you?" * 1000

examples/disaggregated/lmcache/cpu_offload_lmcache_mp.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# CPU offloading with LMCache in multi-process (MP) mode.
+#
+# In MP mode, LMCache runs as a standalone server process (`lmcache server`)
+# that owns the KV cache storage. One or more vLLM instances connect to it via
+# the `LMCacheMPConnector`. This is the recommended way to run LMCache for
+# distributed KV storage and for sharing KV cache across vLLM instances.
+#
+# vLLM ships a built-in shortcut for this setup: pass `--kv-offloading-backend
+# lmcache` together with `--kv-offloading-size <GiB>` and vLLM wires up the
+# `LMCacheMPConnector` for you (it defaults to the LMCache server at
+# tcp://localhost:5555, matching the `lmcache server` default).
+#
+# Requires `lmcache` to be installed (`pip install lmcache`).
+# Learn more: https://docs.lmcache.ai
+set -euo pipefail
+
+MODEL=${MODEL:-meta-llama/Meta-Llama-3.1-8B-Instruct}
+
+# 1. Launch the standalone LMCache server (binds tcp://localhost:5555 by
+#    default). `--l1-size-gb` sets the CPU memory budget for the L1 cache.
+echo "Starting LMCache server..."
+lmcache server --host localhost --port 5555 --l1-size-gb 5 &
+LMCACHE_SERVER_PID=$!
+trap 'kill $LMCACHE_SERVER_PID 2>/dev/null || true' EXIT
+
+# 2. Launch vLLM and offload KV cache to the LMCache server.
+#    The MP connector currently requires the non-hybrid KV cache manager.
+echo "Starting vLLM server with LMCache MP offloading..."
+vllm serve "$MODEL" \
+    --port 8000 \
+    --kv-offloading-size 5 \
+    --kv-offloading-backend lmcache \
+    --disable-hybrid-kv-cache-manager
+
+# Equivalent explicit configuration (instead of the two flags above):
+#   --kv-transfer-config \
+#   '{"kv_connector":"LMCacheMPConnector","kv_role":"kv_both",
+#     "kv_connector_extra_config":{"lmcache.mp.host":"tcp://localhost",
+#                                  "lmcache.mp.port":5555}}'

examples/disaggregated/lmcache/disagg_prefill_lmcache_v0.py

@@ -1,144 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-This file demonstrates the example usage of disaggregated prefilling
-with LMCache.
-We will launch 2 vllm instances (GPU 0 for prefill and GPU 1 for decode),
-and launch an additional LMCache server.
-KV cache is transferred in the following manner:
-vLLM prefill node -> LMCache server -> vLLM decode node.
-
-Note that `pip install lmcache` is needed to run this example.
-Learn more about LMCache in https://github.com/LMCache/LMCache.
-"""
-
-import os
-import subprocess
-import time
-from multiprocessing import Event, Process
-
-from lmcache.experimental.cache_engine import LMCacheEngineBuilder
-from lmcache.integration.vllm.utils import ENGINE_NAME
-
-from vllm import LLM, SamplingParams
-from vllm.config import KVTransferConfig
-
-# LMCache-related environment variables
-# The port to start LMCache server
-port = 8100
-# Use experimental features in LMCache
-os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
-# LMCache is set to use 256 tokens per chunk
-os.environ["LMCACHE_CHUNK_SIZE"] = "256"
-# Disable local CPU backend in LMCache
-os.environ["LMCACHE_LOCAL_CPU"] = "False"
-# Set local CPU memory buffer limit to 5.0 GB
-os.environ["LMCACHE_MAX_LOCAL_CPU_SIZE"] = "5.0"
-# Set the remote URL for LMCache server
-os.environ["LMCACHE_REMOTE_URL"] = f"lm://localhost:{port}"
-# Set the serializer/deserializer between vllm and LMCache server
-# `naive` indicates using raw bytes of the tensor without any compression
-os.environ["LMCACHE_REMOTE_SERDE"] = "naive"
-
-prompts = [
-    "Hello, how are you?" * 1000,
-]
-
-
-def run_prefill(prefill_done, prompts):
-    # We use GPU 0 for prefill node.
-    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
-
-    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
-
-    ktc = KVTransferConfig(
-        kv_connector="LMCacheConnector",
-        kv_role="kv_producer",
-        kv_rank=0,
-        kv_parallel_size=2,
-    )
-    # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
-    # memory. Reduce the value if your GPU has less memory.
-    llm = LLM(
-        model="mistralai/Mistral-7B-Instruct-v0.2",
-        kv_transfer_config=ktc,
-        max_model_len=8000,
-        gpu_memory_utilization=0.8,
-        enforce_eager=True,
-    )
-
-    # llm.generate(prompts, sampling_params)
-    outputs = llm.generate(prompts, sampling_params)
-    for output in outputs:
-        generated_text = output.outputs[0].text
-        print(f"Generated text: {generated_text!r}")
-    print("Prefill node is finished.")
-    prefill_done.set()
-
-    # Clean up lmcache backend
-    LMCacheEngineBuilder.destroy(ENGINE_NAME)
-
-
-def run_decode(prefill_done, prompts, timeout=1):
-    # We use GPU 1 for decode node.
-    os.environ["CUDA_VISIBLE_DEVICES"] = "1"
-
-    sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
-
-    ktc = KVTransferConfig(
-        kv_connector="LMCacheConnector",
-        kv_role="kv_consumer",
-        kv_rank=1,
-        kv_parallel_size=2,
-    )
-    # Set GPU memory utilization to 0.8 for an A40 GPU with 40GB
-    # of memory. Reduce the value if your GPU has less memory.
-    llm = LLM(
-        model="mistralai/Mistral-7B-Instruct-v0.2",
-        kv_transfer_config=ktc,
-        max_model_len=8000,
-        gpu_memory_utilization=0.8,
-        enforce_eager=True,
-    )
-
-    print("Waiting for prefill node to finish...")
-    prefill_done.wait()
-    time.sleep(timeout)
-
-    outputs = llm.generate(prompts, sampling_params)
-    for output in outputs:
-        generated_text = output.outputs[0].text
-        print(f"Generated text: {generated_text!r}")
-
-    # Clean up lmcache backend
-    LMCacheEngineBuilder.destroy(ENGINE_NAME)
-
-
-def run_lmcache_server(port):
-    server_proc = subprocess.Popen(
-        ["python", "-m", "lmcache.experimental.server", "localhost", str(port)]
-    )
-    return server_proc
-
-
-def main():
-    prefill_done = Event()
-    prefill_process = Process(target=run_prefill, args=(prefill_done, prompts))
-    decode_process = Process(target=run_decode, args=(prefill_done, prompts))
-    lmcache_server_process = run_lmcache_server(port)
-
-    # Start prefill node
-    prefill_process.start()
-
-    # Start decode node
-    decode_process.start()
-
-    # Clean up the processes
-    decode_process.join()
-    prefill_process.terminate()
-    lmcache_server_process.terminate()
-    lmcache_server_process.wait()
-
-
-if __name__ == "__main__":
-    main()

examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_vllm_launcher.sh

@@ -30,7 +30,6 @@ if [[ $1 == "prefiller" ]]; then
 
     UCX_TLS=cuda_ipc,cuda_copy,tcp \
         LMCACHE_CONFIG_FILE=$prefill_config_file \
-        LMCACHE_USE_EXPERIMENTAL=True \
         VLLM_ENABLE_V1_MULTIPROCESSING=1 \
         VLLM_WORKER_MULTIPROC_METHOD=spawn \
         CUDA_VISIBLE_DEVICES=0 \
@@ -47,7 +46,6 @@ elif [[ $1 == "decoder" ]]; then
 
     UCX_TLS=cuda_ipc,cuda_copy,tcp \
         LMCACHE_CONFIG_FILE=$decode_config_file \
-        LMCACHE_USE_EXPERIMENTAL=True \
         VLLM_ENABLE_V1_MULTIPROCESSING=1 \
         VLLM_WORKER_MULTIPROC_METHOD=spawn \
         CUDA_VISIBLE_DEVICES=1 \

examples/disaggregated/lmcache/kv_cache_sharing_lmcache_v1.py

@@ -26,8 +26,6 @@ from vllm.config import KVTransferConfig
 # LMCache-related environment variables
 # The port to start LMCache server
 port = 8100
-# Use experimental features in LMCache
-os.environ["LMCACHE_USE_EXPERIMENTAL"] = "True"
 # LMCache is set to use 256 tokens per chunk
 os.environ["LMCACHE_CHUNK_SIZE"] = "256"
 # Disable local CPU backend in LMCache

examples/generate/multimodal/vision_language_offline.py

@@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
     return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")
 
 
-def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData:
-    assert modality == "image"
-
-    model_name = "MiniMaxAI/MiniMax-VL-01"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        max_num_seqs=2,
-        limit_mm_per_prompt={modality: 1},
-        trust_remote_code=True,
-        tensor_parallel_size=8,
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    messages = [
-        [
-            {
-                "role": "user",
-                "content": [{"type": "image"}, {"type": "text", "text": question}],
-            }
-        ]
-        for question in questions
-    ]
-    prompts = tokenizer.apply_chat_template(
-        messages, add_generation_prompt=True, tokenize=False
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # Mistral-3 HF-format
 def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
     assert modality == "image"
@@ -2485,7 +2452,6 @@ model_example_map = {
     "mantis": run_mantis,
     "minicpmo": run_minicpmo,
     "minicpmv": run_minicpmv,
-    "minimax_vl_01": run_minimax_vl_01,
     "mistral3": run_mistral3,
     "molmo": run_molmo,
     "molmo2": run_molmo2,
@@ -2533,15 +2499,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/pooling/score/colqwen3_5_rerank_online.py

@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
 and reranking. Supports both text and image inputs.
 
+Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
+or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
+
 Start the server with:
-    vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
+    vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
+        --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'
 
 Then run this script:
     python colqwen3_5_rerank_online.py
+
+Parity note (matching the native colpali ColQwen3_5Processor pipeline):
+  - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
+    i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
+    (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
+    as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
+  - When you build prompts yourself (token_embed), reproduce the processor exactly:
+      image (document): wrap in the instruction template
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        "Describe the image.<|im_end|><|endoftext|>"
+      query: append the augmentation suffix  <text> + "<|endoftext|>" * 10
+    Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
 """
 
 import requests

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 -%}

examples/tool_chat_template_minimax_m1.jinja

@@ -1,91 +0,0 @@
-{{ '<begin_of_document>' -}}
-{%- if custom_tools is defined %}
-    {%- set tools = custom_tools %}
-{%- endif %}
-{%- if not tools is defined %}
-    {%- set tools = none %}
-{%- endif %}
-
-{#- Extract system message #}
-{% set ns = namespace(system_prompt='') -%}
-{%- if messages[0]['role'] == 'system' %}
-    {%- if messages[0]['content'] is string %}
-        {%- set ns.system_prompt = messages[0]['content']|trim %}
-    {%- else %}
-        {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
-    {%- endif %}
-    {%- set messages = messages[1:] %}
-{%- else %}
-    {%- if tools is not none %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- else %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- endif %}
-{%- endif %}
-
-{#- System message #}
-{%- if ns.system_prompt != '' %}
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Tools configuration #}
-{%- if tools is not none %}
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
-{%- for tool in tools %}
-{{ tool | tojson ~ '\n' -}}
-{%- endfor %}
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Process messages #}
-{%- for message in messages %}
-    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
-        {%- if message['role'] == 'user' %}
-{{ '<beginning_of_sentence>user name=user\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ content['text']|trim -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- elif message['role'] == 'assistant' %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
-{{ content['text']|trim -}}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- endif %}
-    {%- elif 'tool_calls' in message %}
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
-{%- for tool_call in message.tool_calls %}
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
-{%- endfor %}
-{{ '</tool_calls><end_of_sentence>\n' -}}
-    {%- elif message.role == "tool" or message.role == "ipython" %}
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
-{%- if message.content is string %}
-{{ 'tool result: ' + message.content + '\n\n' -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
-{%- elif content.get('name') %}
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-    {%- endif %}
-{%- endfor %}
-
-{%- if add_generation_prompt %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- endif %}

requirements/common.txt

@@ -11,13 +11,14 @@ transformers >= 5.5.3
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 safetensors >= 0.6.2  # MXFP4/MXFP6 dtype support (F8_E8M0, F4) added in 0.6.0: https://github.com/huggingface/safetensors/pull/611
 protobuf >= 5.29.6, !=6.30.*, !=6.31.*, !=6.32.*, !=6.33.0.*, !=6.33.1.*, !=6.33.2.*, !=6.33.3.*, !=6.33.4.* # Required by LlamaTokenizer, gRPC. CVE-2026-0994
-fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
+fastapi[standard] >= 0.133.0, < 0.137.0 # First version supporting Starlette 1.0; < 0.137.0 avoids route-tree change that breaks model-hosting-container-standards handler overrides.
+starlette >= 1.0.1 # CVE-2026-48710: Host header injection in < 1.0.1
 aiohttp >= 3.13.3
 openai >= 2.0.0  # For Responses API with reasoning content
 pydantic >= 2.12.0
 prometheus_client >= 0.18.0
 pillow  # Required for image processing
-prometheus-fastapi-instrumentator >= 7.0.0
+prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
 tiktoken >= 0.6.0  # Required for DBRX tokenizer
 lm-format-enforcer == 0.11.3
 llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"

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