[KV-Offloading] Fix tensors_per_block stride (#46888)

Signed-off-by: <>
Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>
(cherry picked from commit 798185d438)

.buildkite/hardware_tests/cpu.yaml

@@ -91,7 +91,7 @@ steps:
   - tests/quantization/test_cpu_wna16.py
   commands:
     - |
-      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
+      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
       pytest -x -v -s tests/quantization/test_cpu_wna16.py"
       
@@ -136,7 +136,7 @@ steps:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 3
+  parallelism: 4
 
 - label: "Arm CPU Test"
   depends_on: []

.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/image_build/image_build_arm64.sh

@@ -21,12 +21,13 @@ else
   exit 0
 fi
 
-# build (Grace/GH200 is the arm64 GPU target; sm_90)
+# build for arm64 GPU targets: Grace/GH200 (sm_90) and DGX Spark/GB10
+# (sm_121, family-covered by 12.0 under CUDA 13)
 docker build --file docker/Dockerfile \
   --platform linux/arm64 \
   --build-arg max_jobs=16 \
   --build-arg nvcc_threads=4 \
-  --build-arg torch_cuda_arch_list="9.0" \
+  --build-arg torch_cuda_arch_list="9.0 12.0" \
   --build-arg USE_SCCACHE=1 \
   --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
   --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \

.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

@@ -1,11 +1,15 @@
 group: Expert Parallelism
-depends_on: 
+depends_on:
   - image-build-xpu
 steps:
 - label: EPLB Algorithm
   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

@@ -1,11 +1,15 @@
 group: Models - Multimodal
-depends_on: 
+depends_on:
   - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
   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/lm-eval-harness/configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml

@@ -6,9 +6,7 @@ tasks:
     value: 0.7142
   - name: "exact_match,flexible-extract"
     value: 0.4579
-env_vars:
-  VLLM_USE_FLASHINFER_MOE_FP8: "1"
-  VLLM_FLASHINFER_MOE_BACKEND: "throughput"
+moe_backend: "flashinfer_cutlass"
 limit: 1319
 num_fewshot: 5
 max_model_len: 262144

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

@@ -68,6 +68,10 @@ def launch_lm_eval(eval_config, tp_size):
     if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]:
         model_args += "attention_backend=TRITON_ATTN"
 
+    moe_backend = eval_config.get("moe_backend", None)
+    if moe_backend is not None:
+        model_args += f"moe_backend={moe_backend},"
+
     env_vars = eval_config.get("env_vars", None)
     with scoped_env_vars(env_vars):
         results = lm_eval.simple_evaluate(

.buildkite/release-pipeline.yaml

@@ -1,12 +1,25 @@
 # CUDA architecture lists — following PyTorch RELEASE.md
 # (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
 # SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
+# These requested arches are filtered by CMake's CUDA_SUPPORTED_ARCHS before
+# per-kernel arch selection. Do not add +PTX here: top-level +PTX is stripped
+# during that filtering, so kernels that need PTX must request it locally.
 env:
-  CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
-  # aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
-  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
+  # for CUDA >=13, sm_100+ targets have family specifiers (see CMakeLists.txt)
+  # so targets like 10.3 and 12.1 are automatically supported with this list
+  CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
+  # aarch64-only targets: Orin (8.7), Thor (11.0, CUDA 13+)
+  CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0"
+
+  # for CUDA <13, we need to specify all needed targets
+  # some targets (10.3, 12.1) are skipped to limit the wheel size (< 500MB)
+  # please use CUDA 13 wheels or compile yourself on these new devices
   CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
   CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
+  
+  # pre-built mooncake wheels
+  # the manylinux_2_35 wheel has compatibility issue on Ubuntu 24.04
+  # so we use different wheels for the time being
   MOONCAKE_WHEEL_AARCH64_2_35: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_aarch64.whl"
   MOONCAKE_WHEEL_AARCH64_2_39: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_39_aarch64.whl"
   MOONCAKE_WHEEL_X86_64: "https://vllm-wheels.s3.amazonaws.com/mooncake/mooncake_transfer_engine-0.3.10.post2-0da9dfea3-cp312-cp312-manylinux_2_35_x86_64.whl"
@@ -846,7 +859,6 @@ steps:
         allow_failure: true
       - step: build-cpu-release-image-arm64
         allow_failure: true
-    if: build.env("NIGHTLY") != "1"
 
   - label: "Publish release images to DockerHub"
     depends_on:

.buildkite/scripts/ci-clean-log.sh

@@ -13,5 +13,8 @@ INPUT_FILE="$1"
 # Strip timestamps
 sed -i 's/^\[[0-9]\{4\}-[0-9]\{2\}-[0-9]\{2\}T[0-9]\{2\}:[0-9]\{2\}:[0-9]\{2\}Z\] //' "$INPUT_FILE"
 
+# Strip Buildkite inline timestamp markers (ESC _bk;t=<ms> BEL)
+sed -i 's/\x1B_bk;t=[0-9]*\x07//g' "$INPUT_FILE"
+
 # Strip colorization
 sed -i -r 's/\x1B\[[0-9;]*[mK]//g' "$INPUT_FILE"

.buildkite/scripts/ci-fetch-log.sh

@@ -1,74 +1,178 @@
 #!/bin/bash
-# Usage: ./ci-fetch-log.sh <buildkite_job_url> [output_file]
-#        ./ci-fetch-log.sh <build_number> <job_uuid> [output_file]
+# Fetch vLLM Buildkite CI logs (public; no login required).
 #
-# Downloads the raw log for a Buildkite job from the public, unauthenticated
-# /organizations/<org>/pipelines/<pipeline>/builds/<n>/jobs/<uuid>/download
-# endpoint, then strips ANSI/timestamps via ci-clean-log.sh.
+# Usage:
+#   ci-fetch-log.sh [--soft|--all] --pr [<PR>]  failed jobs in the PR's latest
+#                                               build (current branch if omitted)
+#   ci-fetch-log.sh [--soft|--all] <build_url>  failed jobs in that build
+#   ci-fetch-log.sh <job_url> [output]          one job; both #<job_uuid> and
+#                                               ?sid=<id> URL forms work
+#   ci-fetch-log.sh <build> <job_uuid> [output]
 #
-# Find <build_number> and <job_uuid> via:
-#   gh pr checks <PR> --repo vllm-project/vllm
-# Each failing row's URL is .../builds/<build_number>#<job_uuid>.
-#
-# Default output path: ci-<build>-<uuid_first_13_chars>.log (e.g.
-# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's
-# first 8 chars, so the second segment is needed for uniqueness when
-# fetching multiple jobs in parallel. The script refuses to overwrite an
-# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1
-# to override.
+# --soft also fetches soft-failed jobs; --all fetches every finished job.
+# Saves each log as ci-<build>-<job-name>.log (ANSI/timestamps stripped) and
+# prints "<file>\t<job name>" per job. [output] is single-job only; "-"
+# streams to stdout. Existing files are kept; CI_FETCH_LOG_FORCE=1 refetches.
 
 set -euo pipefail
 
 ORG="vllm"
 PIPELINE="ci"
+UA="vllm-ci-fetch-log"
+UUID_RE='[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}'
 
 usage() {
-    echo "Usage: $0 <buildkite_job_url> [output_file]"
-    echo "       $0 <build_number> <job_uuid> [output_file]"
+    sed -n '2,15p' "$0" | sed 's/^# \{0,1\}//'
     exit 1
 }
 
-if [ $# -lt 1 ]; then usage; fi
+die() {
+    echo "$1" >&2
+    exit 1
+}
 
-if [[ "$1" == https://* ]]; then
+BUILD="" JOB="" SID="" OUT=""
+SCOPE="failed"
+
+while :; do
+    case "${1:-}" in
+    --soft) SCOPE="soft" ;;
+    --all) SCOPE="all" ;;
+    *) break ;;
+    esac
+    shift
+done
+
+case "${1:-}" in
+--pr)
+    PR="${2:-}"
+    # gh pr checks exits non-zero when checks are failing; that is the
+    # expected case here.
+    URL=$(gh pr checks ${PR:+"$PR"} --repo vllm-project/vllm 2>/dev/null |
+        grep -oE "https://buildkite.com/${ORG}/${PIPELINE}/builds/[0-9]+" |
+        sort -t/ -k7 -n | tail -1 || true)
+    [ -n "$URL" ] || die "No Buildkite build found via: gh pr checks ${PR:-<current branch>}"
+    BUILD="${URL##*/}"
+    ;;
+https://*)
     BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
-    JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
+    JOB=$(echo "$1" | grep -oE "#${UUID_RE}" | head -n 1 | cut -c2- || true)
+    SID=$(echo "$1" | grep -oE "[?&]sid=${UUID_RE}" | head -n 1 | sed 's/.*sid=//' || true)
     OUT="${2:-}"
-else
-    if [ $# -lt 2 ]; then usage; fi
+    [ -n "$BUILD" ] || die "Could not parse build number from: $1"
+    ;;
+[0-9]*)
+    [ $# -ge 2 ] || usage
     BUILD="$1"
     JOB="$2"
     OUT="${3:-}"
-fi
-
-if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
-    echo "Could not parse build number or job UUID from: $1" >&2
+    ;;
+*)
     usage
-fi
-
-# Jobs in the same build share the UUID's first segment, so include the
-# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames
-# unique when fetching multiple jobs from one build in parallel.
-if [ -z "$OUT" ]; then
-    OUT="ci-${BUILD}-${JOB:0:13}.log"
-fi
-
-if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
-    echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2
-    exit 1
-fi
+    ;;
+esac
 
 COOKIES=$(mktemp)
-trap 'rm -f "$COOKIES"' EXIT
+JOBS_TSV=$(mktemp)
+trap 'rm -f "$COOKIES" "$JOBS_TSV"' EXIT
 
-# Buildkite issues a session cookie on first hit; subsequent /download needs it.
-curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \
+# Buildkite issues a session cookie on first hit; later requests need it.
+curl -fsSL -c "$COOKIES" -A "$UA" \
     "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null
 
-curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \
-    "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \
-    -o "$OUT"
+# The build's job list (id, step uuid, state, name) is served as JSON from
+# the user-facing /data/jobs endpoint. Flatten it to TSV for easy filtering:
+#   job_id  step_uuid  failed  soft_failed  finished  slug  name
+curl -fsSL -b "$COOKIES" -A "$UA" \
+    "https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}/data/jobs" |
+    python3 -c '
+import json, re, sys
 
-bash "$(dirname "$0")/ci-clean-log.sh" "$OUT"
+data = json.load(sys.stdin)
+if data.get("has_next_page"):
+    print("warning: job list is paginated; some jobs not shown", file=sys.stderr)
+for r in data["records"]:
+    if r.get("type") != "script":
+        continue
+    name = (r.get("name") or "").replace("\t", " ").replace("\n", " ")
+    slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:60]
+    print("\t".join([
+        r["id"],
+        r.get("step_uuid") or "",
+        str(r.get("passed") is False),
+        str(bool(r.get("soft_failed"))),
+        str(bool(r.get("finished_at"))),
+        slug,
+        name,
+    ]))
+' >"$JOBS_TSV" || die "Could not list jobs for build ${BUILD}"
 
-echo "$OUT"
+if [ -n "$SID" ] && [ -z "$JOB" ]; then
+    # The ?sid= in builds/<N>/list URLs is the *step* uuid, not the job uuid.
+    JOB=$(awk -F'\t' -v s="$SID" '$1 == s || $2 == s {print $1; exit}' "$JOBS_TSV")
+    [ -n "$JOB" ] || die "No job matching sid=${SID} in build ${BUILD}"
+fi
+
+fetch_job() { # <job_uuid> <output_file>
+    curl -fsSL -b "$COOKIES" -A "$UA" \
+        "https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/$1/download" \
+        -o "$2"
+    bash "$(dirname "$0")/ci-clean-log.sh" "$2"
+}
+
+if [ -n "$JOB" ]; then
+    # Single-job mode.
+    NAME=$(awk -F'\t' -v j="$JOB" '$1 == j {print $7; exit}' "$JOBS_TSV")
+    SLUG=$(awk -F'\t' -v j="$JOB" '$1 == j {print $6; exit}' "$JOBS_TSV")
+    [ -n "$OUT" ] || OUT="ci-${BUILD}-${SLUG:-${JOB:0:13}}.log"
+    if [ "$OUT" = "-" ]; then
+        TMP=$(mktemp)
+        fetch_job "$JOB" "$TMP"
+        cat "$TMP"
+        rm -f "$TMP"
+        exit 0
+    fi
+    if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
+        die "Refusing to overwrite existing ${OUT} (set CI_FETCH_LOG_FORCE=1 or pass an output path)."
+    fi
+    fetch_job "$JOB" "$OUT"
+    printf '%s\t%s\n' "$OUT" "${NAME:-$JOB}"
+    exit 0
+fi
+
+# Build-wide mode: fetch finished jobs matching $SCOPE.
+[ -z "$OUT" ] || die "[output_file] is only valid when fetching a single job."
+
+case "$SCOPE" in
+failed) FILTER='$3 == "True" && $4 == "False" && $5 == "True"' ;;
+soft) FILTER='$3 == "True" && $5 == "True"' ;;
+all) FILTER='$5 == "True"' ;;
+esac
+
+if [ "$SCOPE" = "failed" ]; then
+    SOFT=$(awk -F'\t' '$3 == "True" && $4 == "True"' "$JOBS_TSV" | wc -l)
+    [ "$SOFT" -eq 0 ] || echo "Skipping ${SOFT} soft-failed job(s); use --soft to include them." >&2
+fi
+
+FOUND=0
+EMITTED=" "
+while IFS=$'\t' read -r job_id _ _ _ _ slug name; do
+    FOUND=$((FOUND + 1))
+    out="ci-${BUILD}-${slug:-${job_id:0:13}}.log"
+    # Retries share a name with the original job; disambiguate by uuid.
+    case "$EMITTED" in
+    *" $out "*) out="ci-${BUILD}-${slug:-job}-${job_id:0:13}.log" ;;
+    esac
+    EMITTED="${EMITTED}${out} "
+    if [ -e "$out" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
+        echo "Keeping existing ${out} (set CI_FETCH_LOG_FORCE=1 to refetch)." >&2
+    elif ! fetch_job "$job_id" "$out"; then
+        echo "Failed to download log for job ${job_id} (${name})." >&2
+        continue
+    fi
+    printf '%s\t%s\n' "$out" "$name"
+done < <(awk -F'\t' "$FILTER" "$JOBS_TSV")
+
+if [ "$FOUND" -eq 0 ]; then
+    echo "No matching jobs in build ${BUILD} (scope: ${SCOPE})." >&2
+fi

.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 --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
+    ;;
+  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

@@ -161,24 +161,6 @@ steps:
   commands:
   - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
 
-- label: PyTorch Fullgraph # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/model_executor/
-  - vllm/v1/attention/
-  - vllm/config/compilation.py
-  - csrc/
-  - tests/compile
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
-
 - label: PyTorch Fullgraph Smoke Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -263,37 +245,6 @@ steps:
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 
-- label: Elastic EP Scaling Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/compilation/
-  - tests/distributed/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s distributed/test_elastic_ep.py
-
-- label: EPLB Execution # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_execute.py
-  - tests/distributed/test_eplb_spec_decode.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s distributed/test_eplb_execute.py
-  - pytest -v -s distributed/test_eplb_spec_decode.py
-
 - label: Pipeline + Context Parallelism (4 GPUs) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -315,78 +266,8 @@ steps:
   - pytest -v -s distributed/test_pp_cudagraph.py
   - pytest -v -s distributed/test_pipeline_parallel.py
 
-#-----------------------------------------------------------  mi250 · evals  -----------------------------------------------------------#
-
-- label: Multi-Modal Accuracy Eval (Small Models) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - vllm/multimodal/
-  - vllm/inputs/
-  - vllm/v1/core/
-  - vllm/platforms/rocm.py
-  - vllm/model_executor/model_loader/
-  commands:
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
-
-#---------------------------------------------------------  mi250 · examples  ----------------------------------------------------------#
-
-- label: Examples # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/examples"
-  source_file_dependencies:
-  - vllm/entrypoints
-  - vllm/multimodal
-  - examples/
-  - vllm/platforms/rocm.py
-  commands:
-    - pip install tensorizer
-    # Basic
-    - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 basic/offline_inference/classify.py
-    - python3 basic/offline_inference/embed.py
-    - python3 basic/offline_inference/score.py
-    # Multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
-    # Pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
-    # Features demo
-    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
-
 #----------------------------------------------------------  mi250 · kernels  ----------------------------------------------------------#
 
-- label: Kernels Core Operation Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - csrc/
-  - tests/kernels/core
-  - tests/kernels/test_top_k_per_row.py
-  - tests/kernels/test_concat_mla_q.py
-  - vllm/model_executor/layers/rotary_embedding/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - 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 Helion Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -398,7 +279,7 @@ steps:
   - tests/kernels/helion/
   - vllm/platforms/rocm.py
   commands:
-  - pip install helion==1.0.0
+  - pip install helion==1.1.0
   - pytest -v -s kernels/helion/
 
 - label: Kernels Mamba Test # TBD
@@ -415,22 +296,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
@@ -448,50 +313,6 @@ steps:
   commands:
   - pytest -v -s models/test_utils.py models/test_vision.py
 
-- label: Basic Models Tests (Extra Initialization) %N # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  torch_nightly: true
-  parallelism: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/layers/
-  - tests/models/test_initialization.py
-  - tests/models/registry.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
-- label: Basic Models Tests (Initialization) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/test_initialization.py
-  - tests/models/registry.py
-  commands:
-  - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
-
-- label: Basic Models Tests (Other) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/test_terratorch.py
-  - tests/models/test_transformers.py
-  - tests/models/test_registry.py
-  commands:
-  - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
-
 #-----------------------------------------------------  mi250 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (MTEB) # TBD
@@ -516,53 +337,8 @@ steps:
   commands:
   - pytest -v -s models/language/generation_ppl_test
 
-- label: Language Models Tests (Extra Standard) %N # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  torch_nightly: true
-  parallelism: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/model_executor/models/
-  - vllm/model_executor/model_loader/
-  - vllm/model_executor/layers/
-  - vllm/v1/attention/backends/
-  - vllm/v1/attention/selector.py
-  - tests/models/language/pooling/test_embedding.py
-  - tests/models/language/generation/test_common.py
-  - tests/models/language/pooling/test_classification.py
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pip freeze | grep -E 'torch'
-  - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
-
 #----------------------------------------------------  mi250 · models / multimodal  ----------------------------------------------------#
 
-- label: Multi-Modal Models (Extended Generation 2) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal/generation
-  commands:
-  - 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 Pooling) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/
-  - tests/models/multimodal/pooling
-  commands:
-  - pytest -v -s models/multimodal/pooling -m 'not core_model'
-
 - label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl" # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -608,6 +384,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 +439,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
@@ -791,21 +572,6 @@ steps:
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
   - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
 
-- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_2
-  num_gpus: 2
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-  - vllm/v1/worker/kv_connector_model_runner_mixin.py
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
-
 - label: V1 e2e (2 GPUs) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -818,20 +584,6 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
 
-- label: Distributed NixlConnector PD accuracy (4 GPUs)  # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 #-------------------------------------------------------------  mi250 · misc  ------------------------------------------------------------#
 
 - label: Async Engine, Inputs, Utils, Worker, Config (CPU) # TBD
@@ -991,6 +743,24 @@ steps:
   commands:
   - pytest -s -v compile/passes --ignore compile/passes/distributed
 
+- label: PyTorch Fullgraph # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  torch_nightly: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/compilation/
+  - vllm/model_executor/
+  - vllm/v1/attention/
+  - vllm/config/compilation.py
+  - csrc/
+  - tests/compile
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
+
 - label: Pytorch Nightly Dependency Override Check # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1068,6 +838,21 @@ steps:
   - pytest -v -s distributed/test_eplb_algo.py
   - pytest -v -s distributed/test_eplb_utils.py
 
+- label: EPLB Execution # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_4
+  num_gpus: 4
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/eplb
+  - tests/distributed/test_eplb_execute.py
+  - tests/distributed/test_eplb_spec_decode.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s distributed/test_eplb_execute.py
+  - pytest -v -s distributed/test_eplb_spec_decode.py
+
 - label: Distributed Tests (2xH100-2xMI250) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1421,6 +1206,21 @@ steps:
   commands:
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-small-rocm.txt
 
+- label: Multi-Modal Accuracy Eval (Small Models) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  optional: true
+  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
+  source_file_dependencies:
+  - vllm/multimodal/
+  - vllm/inputs/
+  - vllm/v1/core/
+  - vllm/platforms/rocm.py
+  - vllm/model_executor/model_loader/
+  commands:
+  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
+
 - label: GPQA Eval (GPT-OSS) (2xH100-2xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -1605,9 +1405,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 +1439,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 +1496,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]
@@ -1737,6 +1553,52 @@ steps:
   - pytest -v -s model_executor -m '(not slow_test)'
   - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py
 
+#------------------------------------------------------  mi300 · models / basic  -------------------------------------------------------#
+
+- label: Basic Models Tests (Extra Initialization) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  torch_nightly: true
+  parallelism: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/layers/
+  - tests/models/test_initialization.py
+  - tests/models/registry.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s models/test_initialization.py -k 'not test_can_initialize_small_subset' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
+- label: Basic Models Tests (Initialization) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  torch_nightly: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/test_initialization.py
+  - tests/models/registry.py
+  commands:
+  - pytest -v -s models/test_initialization.py::test_can_initialize_small_subset
+
+- label: Basic Models Tests (Other) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  torch_nightly: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/
+  - tests/models/test_terratorch.py
+  - tests/models/test_transformers.py
+  - tests/models/test_registry.py
+  commands:
+  - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
+
 #-----------------------------------------------------  mi300 · models / language  -----------------------------------------------------#
 
 - label: Language Models Test (Extended Pooling)  # TBD
@@ -1765,6 +1627,28 @@ steps:
   - pip freeze | grep -E 'torch'
   - pytest -v -s models/language -m 'core_model and (not slow_test)'
 
+- label: Language Models Tests (Extra Standard) %N # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_1
+  torch_nightly: true
+  parallelism: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/model_executor/models/
+  - vllm/model_executor/model_loader/
+  - vllm/model_executor/layers/
+  - vllm/v1/attention/backends/
+  - vllm/v1/attention/selector.py
+  - tests/models/language/pooling/test_embedding.py
+  - tests/models/language/generation/test_common.py
+  - tests/models/language/pooling/test_classification.py
+  - vllm/_aiter_ops.py
+  - vllm/platforms/rocm.py
+  commands:
+  - pip freeze | grep -E 'torch'
+  - pytest -v -s models/language -m 'core_model and slow_test' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+
 #----------------------------------------------------  mi300 · models / multimodal  ----------------------------------------------------#
 
 - label: Multi-Modal Models (Extended Generation 1) # TBD
@@ -1778,10 +1662,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 +1676,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 +1956,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 +2001,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/
@@ -2345,6 +2215,21 @@ steps:
   - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
   - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
 
+- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
+  timeout_in_minutes: 180
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+  agent_pool: mi300_2
+  num_gpus: 2
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+  - vllm/v1/worker/kv_connector_model_runner_mixin.py
+  - tests/v1/kv_connector/nixl_integration/
+  - vllm/platforms/rocm.py
+  commands:
+  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
+  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+
 - label: Distributed Tests (2xH100-2xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2754,7 +2639,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  --------------------------------------------------------#
 
@@ -2946,7 +2831,6 @@ steps:
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/transformers_utils/configs/qwen3_5.py
   - vllm/transformers_utils/configs/qwen3_5_moe.py
-  - vllm/model_executor/models/qwen.py
   - vllm/model_executor/models/qwen2.py
   - vllm/model_executor/models/qwen3.py
   - vllm/model_executor/models/qwen3_next.py
@@ -3053,7 +2937,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/"
@@ -3077,11 +2961,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/
@@ -3095,10 +2978,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/
@@ -3184,7 +3067,6 @@ steps:
   - vllm/model_executor/models/qwen3_5_mtp.py
   - vllm/transformers_utils/configs/qwen3_5.py
   - vllm/transformers_utils/configs/qwen3_5_moe.py
-  - vllm/model_executor/models/qwen.py
   - vllm/model_executor/models/qwen2.py
   - vllm/model_executor/models/qwen3.py
   - vllm/model_executor/models/qwen3_next.py
@@ -3423,52 +3305,6 @@ steps:
   commands:
   - pytest -v -s -m 'not slow_test' v1/spec_decode
 
-- label: NixlConnector PD + Spec Decode acceptance (2 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-  agent_pool: mi355_2
-  num_gpus: 2
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-  - vllm/v1/worker/kv_connector_model_runner_mixin.py
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
-
-- label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-  agent_pool: mi355_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
-- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-  agent_pool: mi355_4
-  num_gpus: 4
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
-  - tests/v1/kv_connector/nixl_integration/
-  - vllm/platforms/rocm.py
-  commands:
-  - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
-  - DP_EP=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
-
 #------------------------------------------------------  mi355 · weight_loading  -------------------------------------------------------#
 
 - label: Weight Loading Multiple GPU # TBD

.buildkite/test_areas/attention.yaml

@@ -15,7 +15,7 @@ steps:
     - pytest -v -s v1/attention
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 70
       depends_on:
       - image-build-amd

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

@@ -61,6 +61,20 @@ steps:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
 
+- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs)
+  key: hybrid-ssm-nixlconnector-pd-prefix-cache-2-gpus
+  timeout_in_minutes: 25
+  working_dir: "/vllm-workspace/tests"
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
+    - vllm/v1/core/sched/
+    - vllm/v1/core/kv_cache_coordinator.py
+    - tests/v1/kv_connector/nixl_integration/
+  commands:
+    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
+
 - label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
   key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
   timeout_in_minutes: 30

.buildkite/test_areas/engine.yaml

@@ -28,7 +28,7 @@ steps:
   - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -44,7 +44,7 @@ steps:
     - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 40
       depends_on:
       - image-build-amd
@@ -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/entrypoints.yaml

@@ -46,7 +46,7 @@ steps:
   - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -63,7 +63,7 @@ steps:
   - pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -82,7 +82,7 @@ steps:
   - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 80
       depends_on:
       - image-build-amd
@@ -104,7 +104,7 @@ steps:
   - pytest -v -s entrypoints/anthropic
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd
@@ -165,7 +165,7 @@ steps:
   - pytest -s entrypoints/openai/correctness/
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
       source_file_dependencies:

.buildkite/test_areas/kernels.yaml

@@ -47,8 +47,10 @@ steps:
   - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
   - vllm/models/deepseek_v4/common/ops/
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
   commands:
     - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_top_k_per_row.py
 
 - label: Deepseek V4 Kernel Test (B200)
   key: deepseek-v4-kernel-test-b200
@@ -74,6 +76,33 @@ 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
+  timeout_in_minutes: 20
+  device: h100
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/v1/attention/ops/triton_unified_attention_diffkv.py
+  - vllm/v1/attention/backends/triton_attn_diffkv.py
+  - vllm/v1/attention/backends/flash_attn_diffkv.py
+  - tests/kernels/attention/test_triton_unified_attention_diffkv.py
+  commands:
+    - pytest -v -s kernels/attention/test_triton_unified_attention_diffkv.py
 
 - label: Kernels Quantization Test %N
   key: kernels-quantization-test
@@ -87,10 +116,11 @@ steps:
   parallelism: 2
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       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
@@ -114,6 +144,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
@@ -224,7 +270,7 @@ steps:
   - vllm/utils/import_utils.py
   - tests/kernels/helion/
   commands:
-    - pip install helion==1.0.0
+    - pip install helion==1.1.0
     - pytest -v -s kernels/helion/
 
  

.buildkite/test_areas/lm_eval.yaml

@@ -14,7 +14,7 @@ steps:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 55
       depends_on:
       - image-build-amd
@@ -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
@@ -126,6 +136,49 @@ steps:
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
+- label: LM Eval Humming (A100 - TEMPORARY)
+  key: lm-eval-humming-a100
+  timeout_in_minutes: 30
+  device: a100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+
+- label: LM Eval Humming (H100 - TEMPORARY)
+  key: lm-eval-humming-h100
+  timeout_in_minutes: 30
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
+
+- label: LM Eval Humming (B200 - TEMPORARY)
+  key: lm-eval-humming-b200
+  timeout_in_minutes: 30
+  device: b200-k8s
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
 
 - label: LM Eval TurboQuant KV Cache
   key: lm-eval-turboquant-kv-cache
@@ -167,6 +220,20 @@ steps:
     - uv pip install --system 'gpt-oss[eval]==0.0.5'
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
 
+- label: GPQA Eval (GPT-OSS) (DGX Spark)
+  key: gpqa-eval-gpt-oss-spark
+  timeout_in_minutes: 120
+  device: dgx-spark
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - tests/evals/gpt_oss/
+  commands:
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
+
 - label: MRCR Eval Small Models
   device: h200_35gb
   timeout_in_minutes: 30

.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
@@ -138,11 +150,26 @@ steps:
     - vllm/v1/spec_decode/extract_hidden_states.py
     - vllm/model_executor/models/extract_hidden_states.py
     - vllm/transformers_utils/configs/extract_hidden_states.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
     - tests/v1/kv_connector/extract_hidden_states_integration
   commands:
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -v -s v1/kv_connector/extract_hidden_states_integration
 
+- label: Extract Hidden States Integration (2 GPUs)
+  key: extract-hidden-states-integration-2-gpus
+  timeout_in_minutes: 20
+  num_devices: 2
+  source_file_dependencies:
+    - vllm/v1/spec_decode/extract_hidden_states.py
+    - vllm/model_executor/models/extract_hidden_states.py
+    - vllm/transformers_utils/configs/extract_hidden_states.py
+    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+    - tests/v1/kv_connector/extract_hidden_states_integration
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
+
 - label: Regression
   key: regression
   timeout_in_minutes: 20
@@ -293,6 +320,7 @@ steps:
   - vllm/transformers_utils/
   - vllm/utils/
   - vllm/v1/
+  - tests/test_envs.py
   - tests/test_inputs.py
   - tests/test_outputs.py
   - tests/test_pooling_params.py
@@ -309,6 +337,7 @@ steps:
   device: cpu-small
   commands:
   - python3 standalone_tests/lazy_imports.py
+  - pytest -v -s test_envs.py
   - pytest -v -s test_inputs.py
   - pytest -v -s test_outputs.py
   - pytest -v -s test_pooling_params.py

.buildkite/test_areas/models_multimodal.yaml

@@ -15,7 +15,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -30,10 +30,9 @@ 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: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -50,7 +49,7 @@ steps:
     - pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       depends_on:
       - image-build-amd
 
@@ -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
@@ -155,7 +160,7 @@ steps:
     - pytest -v -s models/multimodal/pooling -m 'not core_model'
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 60
       depends_on:
       - image-build-amd

.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
@@ -40,3 +44,17 @@ steps:
   - pytest -v -s plugins_tests/test_oot_registration_online.py # it needs a clean process
   - pytest -v -s plugins_tests/test_oot_registration_offline.py # it needs a clean process
   - pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
+
+
+- label: GGUF Plugin
+  key: gguf-plugin
+  device: h200_18gb
+  timeout_in_minutes: 30
+  soft_fail: true
+  optional: true
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization
+  - tests/plugins_tests/test_gguf_plugin.py
+  commands:
+    - pip install "vllm-gguf-plugin >= 0.0.2"
+    - pytest -v -s plugins_tests/gguf

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

@@ -21,6 +21,18 @@ steps:
   - uv pip install --system conch-triton-kernels
   - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
 
+- label: Quantized Fusions
+  key: quantized-fusions
+  timeout_in_minutes: 30
+  source_file_dependencies:
+  - tests/fusion
+  - vllm/model_executor/layers/fusion
+  - vllm/model_executor/kernels/linear
+  - vllm/model_executor/layers/quantization/compressed_tensors
+  - vllm/model_executor/layers/quantization/modelopt.py
+  commands:
+    - pytest -v -s fusion/
+
 - label: Quantized MoE Test (B200)
   key: quantized-moe-test-b200
   timeout_in_minutes: 60

.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
@@ -39,7 +53,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
@@ -78,7 +92,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 65
       depends_on:
       - image-build-amd
@@ -103,7 +117,7 @@ steps:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
   mirror:
     amd:
-      device: mi300_1
+      device: mi325_1
       timeout_in_minutes: 50
       depends_on:
       - image-build-amd

.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
@@ -80,7 +79,7 @@
 /vllm/v1/worker/kv_connector_model_runner_mixin.py @orozery @NickLucche
 
 # Model runner V2
-/vllm/v1/worker/gpu @WoosukKwon @njhill
+/vllm/v1/worker/gpu @WoosukKwon @njhill @yewentao256
 /vllm/v1/worker/gpu/kv_connector.py @orozery
 
 # CI & building
@@ -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/dependabot.yml

@@ -21,7 +21,6 @@ updates:
       - dependency-name: "torchvision"
       - dependency-name: "xformers"
       - dependency-name: "lm-format-enforcer"
-      - dependency-name: "gguf"
       - dependency-name: "compressed-tensors"
       - dependency-name: "ray[cgraph]" # Ray Compiled Graph
       - dependency-name: "lm-eval"

.github/mergify.yml

@@ -144,12 +144,12 @@ pull_request_rules:
     - label != stale
     - or:
       - files~=^examples/.*mistral.*\.py
-      - files~=^tests/.*mistral.*\.py
-      - files~=^vllm/model_executor/models/.*mistral.*\.py
+      - files~=^tests/.*(?:mistral|voxtral|mixtral|pixtral).*\.py
+      - files~=^vllm/model_executor/models/.*(?:mistral|voxtral|mixtral|pixtral).*\.py
       - files~=^vllm/reasoning/.*mistral.*\.py
       - files~=^vllm/tool_parsers/.*mistral.*\.py
-      - files~=^vllm/transformers_utils/.*mistral.*\.py
-      - title~=(?i)Mistral
+      - files~=^vllm/transformers_utils/.*(?:mistral|voxtral|pixtral).*\.py
+      - title~=(?i)(?:mistral|ministral|voxtral|mixtral|pixtral)
   actions:
     label:
       add:
@@ -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"

.github/workflows/scripts/build.sh

@@ -9,7 +9,7 @@ PATH=${cuda_home}/bin:$PATH
 LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
 
 # Install requirements
-if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
+if [ "$(echo "$2" | cut -d. -f1)" = "12" ]; then
     sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
 fi
 $python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
@@ -17,7 +17,10 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
 # Limit the number of parallel jobs to avoid OOM
 export MAX_JOBS=1
 # Make sure release wheels are built for the following architectures
-export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
+# Do not add +PTX here: vLLM filters torch's top-level PTX flag when it
+# converts global gencode flags into per-kernel arch lists. If a specific
+# kernel needs PTX, add +PTX to that kernel's CMake arch list instead.
+export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0"
 
 bash tools/check_repo.sh
 

.gitignore

@@ -15,6 +15,9 @@ vllm/third_party/flashmla/flash_mla_interface.py
 # DeepGEMM vendored package built from source
 vllm/third_party/deep_gemm/
 
+# fmha_sm100 vendored package built from source
+vllm/third_party/fmha_sm100/
+
 # triton jit
 .triton
 
@@ -196,7 +199,9 @@ cython_debug/
 .vscode/
 
 # Claude
-.claude/
+.claude/*
+!.claude/skills/
+!.claude/skills/**
 
 # Codex
 .codex/
@@ -233,7 +238,7 @@ actionlint
 shellcheck*/
 
 # Ignore moe/marlin_moe gen code
-csrc/moe/marlin_moe_wna16/kernel_*
+csrc/libtorch_stable/moe/marlin_moe_wna16/kernel_*
 
 # Ignore ep_kernels_workspace folder
 ep_kernels_workspace/

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
   rev: v21.1.2
   hooks:
   - id: clang-format
-    exclude: 'csrc/(moe/topk_softmax_kernels.cu|libtorch_stable/quantization/gguf/(ggml-common.h|dequantize.cuh|vecdotq.cuh|mmq.cuh|mmvq.cuh))|vllm/third_party/.*'
+    exclude: 'csrc/libtorch_stable/moe/topk_softmax_kernels.cu|vllm/third_party/.*'
     types_or: [c++, cuda]
     args: [--style=file, --verbose]
 - repo: https://github.com/DavidAnson/markdownlint-cli2

AGENTS.md

@@ -105,6 +105,15 @@ The line length limit for Python code is 88 characters. If you are not sure, use
 
 Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings) (`Args:`/`Returns:`/`Raises:` sections), not reStructuredText/Sphinx fields (`:param:`, `:return:`, `:rtype:`).
 
+### Coding style guidelines
+
+Follow these rules for all code changes in this repository:
+
+- Try to match existing code style.
+- Code should be self-documenting and self-explanatory.
+- Keep comments and docstrings minimal and concise.
+- Assume the reader is familiar with vLLM.
+
 ### Commit messages
 
 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
@@ -126,6 +135,12 @@ Do not modify code in these areas without first reading and following the
 linked guide. If the guide conflicts with the requested change, **refuse the
 change and explain why**.
 
+Security reviewers should start with [`SECURITY.md`](SECURITY.md),
+[`docs/usage/security.md`](docs/usage/security.md), and
+[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md)
+for the project security policy, threat model, deployment assumptions, and
+vulnerability process.
+
 - **Editing these instructions**:
   [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md)
   — Rules for modifying AGENTS.md or any domain-specific guide it references.

SECURITY.md

@@ -34,6 +34,15 @@ Vulnerabilities that cause denial of service or partial disruption, but do not a
 
 Minor issues such as informational disclosures, logging errors, non-exploitable flaws, or weaknesses that require local or high-privilege access and offer negligible impact. Examples include side channel attacks or hash collisions. These issues often have CVSS scores less than 4.0
 
+## Fix disclosure policy
+
+When a security report is accepted, the fix process depends on the severity:
+
+* **CRITICAL and HIGH severity**: Fixes are developed in a private security fork and coordinated with the prenotification group before public disclosure.
+* **MODERATE and LOW severity**: Fixes are developed and submitted as public pull requests. These issues do not require embargo since they do not enable arbitrary code execution or significant data breach, and public visibility accelerates community review and adoption of the fix.
+
+The vulnerability management team reserves the right to adjust the disclosure approach on a case-by-case basis, taking into account factors such as active exploitation, unusual attack surface, or coordination requirements with downstream vendors.
+
 ## Prenotification policy
 
 For certain security issues of CRITICAL, HIGH, or MODERATE severity level, we may prenotify certain organizations or vendors that ship vLLM. The purpose of this prenotification is to allow for a coordinated release of fixes for severe issues.

benchmarks/attention_benchmarks/README.md

@@ -108,7 +108,6 @@ python benchmark.py \
     --backends flash triton flashinfer \
     --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \
     --num-layers 10 \
-    --repeats 5 \
     --output-csv results.csv
 ```
 
@@ -164,14 +163,17 @@ python benchmark.py \
 # Model configuration
 --num-layers N                      # Number of layers
 --head-dim N                        # Head dimension
+--v-head-dim N                      # Value head dimension (defaults to --head-dim)
 --num-q-heads N                     # Query heads
 --num-kv-heads N                    # KV heads
 --block-size N                      # Block size
+--kv-lora-rank N                    # MLA KV LoRA rank
+--qk-nope-head-dim N                # MLA non-RoPE QK head dim
+--qk-rope-head-dim N                # MLA RoPE QK head dim
 
 # Benchmark settings
 --device DEVICE                     # Device (default: cuda:0)
---repeats N                         # Repetitions
---warmup-iters N                    # Warmup iterations
+--warmup-ms N                       # Warmup window in ms for triton do_bench
 --profile-memory                    # Profile memory usage
 
 # Parameter sweeps
@@ -211,8 +213,6 @@ config = BenchmarkConfig(
     num_kv_heads=1,
     block_size=128,
     device="cuda:0",
-    repeats=5,
-    warmup_iters=3,
 )
 
 # CUTLASS MLA with specific num_kv_splits
@@ -253,14 +253,10 @@ formatter.save_json(results, "output.json")
 
 ## Tips
 
-**1. Warmup matters** - Use `--warmup-iters 10` for stable results
+**1. Save results** - Always use `--output-csv` or `--output-json`
 
-**2. Multiple repeats** - Use `--repeats 20` for low variance
+**2. Test incrementally** - Start with `--num-layers 1`
 
-**3. Save results** - Always use `--output-csv` or `--output-json`
+**3. Extended grammar** - Leverage spec decode, chunked prefill patterns
 
-**4. Test incrementally** - Start with `--num-layers 1 --repeats 1`
-
-**5. Extended grammar** - Leverage spec decode, chunked prefill patterns
-
-**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values
+**4. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values

benchmarks/attention_benchmarks/benchmark.py

@@ -26,6 +26,9 @@ Examples:
 """
 
 import argparse
+import os
+import shutil
+import subprocess
 import sys
 from dataclasses import replace
 from pathlib import Path
@@ -50,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
@@ -83,13 +96,15 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
         else:
             return run_standard_attention_benchmark(config)
     except Exception as e:
+        error_msg = str(e) or repr(e)
         return BenchmarkResult(
             config=config,
             mean_time=float("inf"),
+            median_time=float("inf"),
             std_time=0,
             min_time=float("inf"),
             max_time=float("inf"),
-            error=str(e),
+            error=error_msg,
         )
 
 
@@ -115,9 +130,12 @@ def run_model_parameter_sweep(
     """
     all_results = []
 
-    console.print(
-        f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]"
+    sweep_desc = (
+        f"{sweep.param_name} = {sweep.values}"
+        if sweep.param_name
+        else f"{len(sweep.values)} configurations"
     )
+    console.print(f"[yellow]Model sweep mode: testing {sweep_desc}[/]")
 
     total = len(backends) * len(batch_specs) * len(sweep.values)
 
@@ -125,9 +143,9 @@ def run_model_parameter_sweep(
         for backend in backends:
             for spec in batch_specs:
                 for value in sweep.values:
-                    # Create config with modified model parameter
+                    # Create config with modified model parameter(s)
                     config_args = base_config_args.copy()
-                    config_args[sweep.param_name] = value
+                    sweep.apply(config_args, value)
 
                     # Create config with original backend for running
                     clean_config = BenchmarkConfig(
@@ -144,13 +162,21 @@ def run_model_parameter_sweep(
                     all_results.append(result)
 
                     if not result.success:
+                        err_label = (
+                            f"{sweep.param_name}={value}"
+                            if sweep.param_name
+                            else f"{value}"
+                        )
                         console.print(
-                            f"[red]Error {backend} {spec} {sweep.param_name}="
-                            f"{value}: {result.error}[/]"
+                            f"[red]Error {backend} {spec} {err_label}"
+                            f": {result.error}[/]"
                         )
 
                     pbar.update(1)
 
+    if base_config_args.get("ncu_profile"):
+        return all_results
+
     # Display sweep results - create separate table for each parameter value
     console.print("\n[bold green]Model Parameter Sweep Results:[/]")
     formatter = ResultsFormatter(console)
@@ -184,7 +210,10 @@ def run_model_parameter_sweep(
     )
 
     for param_value in sorted_param_values:
-        console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]")
+        label = (
+            f"{sweep.param_name} = {param_value}" if sweep.param_name else param_value
+        )
+        console.print(f"\n[bold cyan]{label}[/]")
         param_results = by_param_value[param_value]
 
         # Create modified results with original backend names
@@ -200,8 +229,9 @@ def run_model_parameter_sweep(
         formatter.print_table(modified_results, backends, compare_to_fastest=True)
 
     # Show optimal backend for each (param_value, batch_spec) combination
+    sweep_name = sweep.param_name or "config"
     console.print(
-        f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]"
+        f"\n[bold cyan]Optimal backend for each ({sweep_name}, batch_spec):[/]"
     )
 
     # Group by (param_value, batch_spec)
@@ -236,7 +266,10 @@ def run_model_parameter_sweep(
     for param_value, spec in sorted_keys:
         # Print header when param value changes
         if param_value != current_param_value:
-            console.print(f"\n  [bold]{sweep.param_name}={param_value}:[/]")
+            header = (
+                f"{sweep.param_name}={param_value}" if sweep.param_name else param_value
+            )
+            console.print(f"\n  [bold]{header}:[/]")
             current_param_value = param_value
 
         results = by_param_and_spec[(param_value, spec)]
@@ -322,6 +355,9 @@ def run_parameter_sweep(
 
                     pbar.update(1)
 
+    if base_config_args.get("ncu_profile"):
+        return all_results
+
     # Display sweep results
     console.print("\n[bold green]Sweep Results:[/]")
     backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]
@@ -459,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(
@@ -474,11 +524,35 @@ def main():
     parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
     parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
     parser.add_argument("--block-size", type=int, default=16, help="Block size")
+    parser.add_argument(
+        "--v-head-dim",
+        type=int,
+        default=None,
+        help="Value head dimension (defaults to --head-dim if unset)",
+    )
+
+    # MLA-specific model dimensions
+    parser.add_argument(
+        "--kv-lora-rank", type=int, default=None, help="MLA KV LoRA rank"
+    )
+    parser.add_argument(
+        "--qk-nope-head-dim", type=int, default=None, help="MLA non-RoPE QK head dim"
+    )
+    parser.add_argument(
+        "--qk-rope-head-dim", type=int, default=None, help="MLA RoPE QK head dim"
+    )
 
     # Benchmark settings
     parser.add_argument("--device", default="cuda:0", help="Device")
-    parser.add_argument("--repeats", type=int, default=1, help="Repetitions")
-    parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")
+    parser.add_argument(
+        "--warmup-ms",
+        type=int,
+        default=None,
+        help=(
+            "Warmup window in ms for triton's do_bench (default: triton's own). "
+            "Has no effect with CUDA graphs; pass --no-cuda-graphs to use it."
+        ),
+    )
     parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
     parser.add_argument(
         "--kv-cache-dtype",
@@ -491,10 +565,33 @@ def main():
         action=argparse.BooleanOptionalAction,
         default=True,
         help=(
-            "Launch kernels with CUDA graphs to eliminate CPU overhead"
-            "in measurements (default: True)"
+            "Use triton do_bench_cudagraph (True) or do_bench (False) "
+            "for timing. CUDA graphs eliminate CPU launch overhead "
+            "(default: True)"
         ),
     )
+    parser.add_argument(
+        "--num-splits",
+        type=int,
+        default=None,
+        help="FlashAttention split-K factor (0=auto heuristic, 1=disabled, >1=force N)",
+    )
+    parser.add_argument(
+        "--ncu-profile",
+        action="store_true",
+        default=False,
+        help=(
+            "Enable Nsight Compute profiling mode. Automatically wraps the "
+            "script with ncu, capturing a profile with source correlation. "
+            "Use --ncu-output to set the output file name."
+        ),
+    )
+    parser.add_argument(
+        "--ncu-output",
+        type=str,
+        default="profile",
+        help="Output file name for ncu profile (default: 'profile').",
+    )
 
     # Parameter sweep (use YAML config for advanced sweeps)
     parser.add_argument(
@@ -545,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)
 
@@ -576,23 +679,28 @@ def main():
             model = yaml_config["model"]
             args.num_layers = model.get("num_layers", args.num_layers)
             args.head_dim = model.get("head_dim", args.head_dim)
+            args.v_head_dim = model.get("v_head_dim", args.v_head_dim)
             args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
             args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
             args.block_size = model.get("block_size", args.block_size)
+            # MLA-specific dimensions
+            args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
+            args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
+            args.qk_rope_head_dim = model.get("qk_rope_head_dim", args.qk_rope_head_dim)
 
         # Benchmark settings (top-level keys)
         if "device" in yaml_config:
             args.device = yaml_config["device"]
-        if "repeats" in yaml_config:
-            args.repeats = yaml_config["repeats"]
-        if "warmup_iters" in yaml_config:
-            args.warmup_iters = yaml_config["warmup_iters"]
+        if "warmup_ms" in yaml_config:
+            args.warmup_ms = yaml_config["warmup_ms"]
         if "profile_memory" in yaml_config:
             args.profile_memory = yaml_config["profile_memory"]
         if "kv_cache_dtype" in yaml_config:
             args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
         if "cuda_graphs" in yaml_config:
             args.cuda_graphs = yaml_config["cuda_graphs"]
+        if "ncu_profile" in yaml_config:
+            args.ncu_profile = yaml_config["ncu_profile"]
 
         # Parameter sweep configuration
         if "parameter_sweep" in yaml_config:
@@ -612,7 +720,7 @@ def main():
         if "model_parameter_sweep" in yaml_config:
             sweep_config = yaml_config["model_parameter_sweep"]
             args.model_parameter_sweep = ModelParameterSweep(
-                param_name=sweep_config["param_name"],
+                param_name=sweep_config.get("param_name"),
                 values=sweep_config["values"],
                 label_format=sweep_config.get(
                     "label_format", "{backend}_{param_name}_{value}"
@@ -631,6 +739,32 @@ def main():
 
         console.print()
 
+    # Re-exec under ncu if --ncu-profile and not already inside ncu. This runs
+    # after YAML processing so ncu_profile set via config file is honored.
+    if args.ncu_profile and "_NCU_INNER" not in os.environ:
+        ncu = shutil.which("ncu")
+        if ncu is None:
+            print("Error: 'ncu' not found in PATH", file=sys.stderr)
+            sys.exit(1)
+        cmd = [
+            ncu,
+            "--profile-from-start",
+            "off",
+            "--set",
+            "full",
+            "--import-source",
+            "yes",
+            "-o",
+            args.ncu_output,
+            sys.executable,
+            *sys.argv,
+        ]
+        env = os.environ.copy()
+        env["CUTE_DSL_LINEINFO"] = "1"
+        env["_NCU_INNER"] = "1"
+        print(f"Launching: {' '.join(cmd)}")
+        sys.exit(subprocess.call(cmd, env=env))
+
     # Handle CLI-based parameter sweep (if not from YAML)
     if (
         (not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)
@@ -655,6 +789,18 @@ def main():
     console.print(f"Batch specs: {', '.join(args.batch_specs)}")
     console.print(f"KV cache dtype: {args.kv_cache_dtype}")
     console.print(f"CUDA graphs: {args.cuda_graphs}")
+    if args.warmup_ms is not None and args.cuda_graphs:
+        console.print(
+            "[yellow]Warning: --warmup-ms is ignored with CUDA graphs "
+            "(do_bench_cudagraph warms up internally). Pass --no-cuda-graphs "
+            "to use it.[/]"
+        )
+    if args.num_splits == 0 and args.cuda_graphs:
+        console.print(
+            "[yellow]Warning: --num-splits 0 (FA3 heuristic) is not CUDA-graph "
+            "compatible and may fail or fall back. Pass --no-cuda-graphs or use "
+            "--num-splits >=1.[/]"
+        )
     console.print()
 
     init_workspace_manager(args.device)
@@ -662,8 +808,68 @@ def main():
     # Run benchmarks
     all_results = []
 
+    # Under ncu profiling the kernels run only to be captured by the profiler;
+    # timings are placeholder zeros, so the result tables and saved metrics are
+    # skipped. The Nsight Compute report (--ncu-output) holds the real data.
+    if args.ncu_profile:
+        console.print(
+            "[dim]ncu profiling enabled: result tables and saved metrics are "
+            "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[/]"
@@ -708,11 +914,11 @@ def main():
                         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,
+                        ncu_profile=args.ncu_profile,
+                        warmup_ms=args.warmup_ms,
                     )
 
                     # Add decode pipeline config
@@ -749,6 +955,7 @@ def main():
                         result = BenchmarkResult(
                             config=config,
                             mean_time=timing["mean"],
+                            median_time=timing.get("median", timing["mean"]),
                             std_time=timing["std"],
                             min_time=timing["min"],
                             max_time=timing["max"],
@@ -770,6 +977,7 @@ def main():
                         result = BenchmarkResult(
                             config=config,
                             mean_time=float("inf"),
+                            median_time=float("inf"),
                             std_time=0,
                             min_time=float("inf"),
                             max_time=float("inf"),
@@ -779,6 +987,9 @@ def main():
 
                 pbar.update(1)
 
+        if args.ncu_profile:
+            return
+
         # Display decode vs prefill results
         console.print("\n[bold green]Decode vs Prefill Results:[/]")
 
@@ -858,15 +1069,20 @@ def main():
         base_config_args = {
             "num_layers": args.num_layers,
             "head_dim": args.head_dim,
+            "v_head_dim": args.v_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,
+            "ncu_profile": args.ncu_profile,
+            "warmup_ms": args.warmup_ms,
+            "num_splits": args.num_splits,
+            "kv_lora_rank": args.kv_lora_rank,
+            "qk_nope_head_dim": args.qk_nope_head_dim,
+            "qk_rope_head_dim": args.qk_rope_head_dim,
         }
         all_results = run_model_parameter_sweep(
             backends,
@@ -882,15 +1098,17 @@ def main():
         base_config_args = {
             "num_layers": args.num_layers,
             "head_dim": args.head_dim,
+            "v_head_dim": args.v_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,
+            "ncu_profile": args.ncu_profile,
+            "warmup_ms": args.warmup_ms,
+            "num_splits": args.num_splits,
         }
         all_results = run_parameter_sweep(
             backends, args.batch_specs, base_config_args, args.parameter_sweep, console
@@ -914,15 +1132,17 @@ def main():
                             batch_spec=spec,
                             num_layers=args.num_layers,
                             head_dim=args.head_dim,
+                            v_head_dim=getattr(args, "v_head_dim", None),
                             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,
+                            ncu_profile=args.ncu_profile,
+                            warmup_ms=args.warmup_ms,
+                            num_splits=args.num_splits,
                         )
 
                         result = run_benchmark(config)
@@ -935,9 +1155,10 @@ def main():
 
                         pbar.update(1)
 
-            console.print("\n[bold green]Results:[/]")
-            formatter = ResultsFormatter(console)
-            formatter.print_table(decode_results, backends)
+            if not args.ncu_profile:
+                console.print("\n[bold green]Results:[/]")
+                formatter = ResultsFormatter(console)
+                formatter.print_table(decode_results, backends)
 
         # Run prefill backend comparison
         if prefill_backends:
@@ -962,9 +1183,8 @@ def main():
                             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,
+                            warmup_ms=args.warmup_ms,
                             prefill_backend=pb,
                         )
 
@@ -980,16 +1200,17 @@ def main():
 
                         pbar.update(1)
 
-            console.print("\n[bold green]Prefill Backend Results:[/]")
-            formatter = ResultsFormatter(console)
-            formatter.print_table(
-                prefill_results, prefill_backends, compare_to_fastest=True
-            )
+            if not args.ncu_profile:
+                console.print("\n[bold green]Prefill Backend Results:[/]")
+                formatter = ResultsFormatter(console)
+                formatter.print_table(
+                    prefill_results, prefill_backends, compare_to_fastest=True
+                )
 
         all_results = decode_results + prefill_results
 
-    # Save results
-    if all_results:
+    # Save results (skip ncu profiling runs: timings are placeholder zeros)
+    if all_results and not args.ncu_profile:
         formatter = ResultsFormatter(console)
         if args.output_csv:
             formatter.save_csv(all_results, args.output_csv)

benchmarks/attention_benchmarks/common.py

@@ -15,6 +15,8 @@ from batch_spec import get_batch_type, parse_batch_spec
 from rich.console import Console
 from rich.table import Table
 
+from vllm.triton_utils import triton
+
 
 def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
     """
@@ -34,6 +36,30 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:
         return (0, 0, 0)
 
 
+def run_do_bench(
+    benchmark_fn,
+    use_cuda_graphs: bool,
+    warmup_ms: int | None = None,
+) -> list[float]:
+    kwargs: dict[str, Any] = {"return_mode": "all"}
+    if use_cuda_graphs:
+        result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)
+    else:
+        if warmup_ms is not None:
+            kwargs["warmup"] = warmup_ms
+        result = triton.testing.do_bench(benchmark_fn, **kwargs)
+    return result
+
+
+def run_ncu_profile(benchmark_fn) -> None:
+    benchmark_fn()
+    torch.accelerator.synchronize()
+    torch.cuda.cudart().cudaProfilerStart()
+    benchmark_fn()
+    torch.accelerator.synchronize()
+    torch.cuda.cudart().cudaProfilerStop()
+
+
 # Mock classes for vLLM attention infrastructure
 
 
@@ -182,18 +208,37 @@ class ParameterSweep:
 
 @dataclass
 class ModelParameterSweep:
-    """Configuration for sweeping a model configuration parameter."""
+    """Configuration for sweeping model configuration parameter(s).
 
-    param_name: str  # Name of the model config parameter to sweep (e.g., "num_q_heads")
-    values: list[Any]  # List of values to test
-    label_format: str = "{backend}_{param_name}_{value}"  # Result label template
+    Supports two modes:
+    - Single param: param_name="head_dim", values=[128, 256, 512]
+    - Multi param: values=[{head_dim: 192, v_head_dim: 128}, {head_dim: 256}]
+      When values are dicts, each dict's keys are applied as config overrides.
+    """
+
+    param_name: str | None = None
+    values: list[Any] | None = None
+    label_format: str = "{backend}_{param_name}_{value}"
 
     def get_label(self, backend: str, value: Any) -> str:
         """Generate a label for a specific parameter value."""
+        if isinstance(value, dict):
+            return self.label_format.format(
+                backend=backend, param_name=self.param_name, value=value, **value
+            )
         return self.label_format.format(
             backend=backend, param_name=self.param_name, value=value
         )
 
+    def apply(self, config_args: dict, value: Any) -> None:
+        """Apply a sweep value to config args."""
+        if isinstance(value, dict):
+            config_args.update(value)
+        elif self.param_name is not None:
+            config_args[self.param_name] = value
+        else:
+            raise ValueError("param_name must be set if sweep values are not dicts")
+
 
 @dataclass
 class BenchmarkConfig:
@@ -208,10 +253,10 @@ class BenchmarkConfig:
     block_size: int
     device: str
     dtype: torch.dtype = torch.float16
-    repeats: int = 1
-    warmup_iters: int = 3
     profile_memory: bool = False
     use_cuda_graphs: bool = False
+    ncu_profile: bool = False
+    warmup_ms: int | None = None
 
     # "auto" or "fp8"
     kv_cache_dtype: str = "auto"
@@ -226,6 +271,7 @@ class BenchmarkConfig:
     # Backend-specific tuning
     num_kv_splits: int | None = None  # CUTLASS MLA
     reorder_batch_threshold: int | None = None  # FlashAttn MLA, FlashMLA
+    num_splits: int | None = None  # FlashAttention split-K (0=auto, 1=disabled)
 
 
 @dataclass
@@ -234,6 +280,7 @@ class BenchmarkResult:
 
     config: BenchmarkConfig
     mean_time: float  # seconds
+    median_time: float  # seconds
     std_time: float  # seconds
     min_time: float  # seconds
     max_time: float  # seconds
@@ -252,6 +299,7 @@ class BenchmarkResult:
         return {
             "config": asdict(self.config),
             "mean_time": self.mean_time,
+            "median_time": self.median_time,
             "std_time": self.std_time,
             "min_time": self.min_time,
             "max_time": self.max_time,

benchmarks/attention_benchmarks/configs/mla_decode.yaml

@@ -56,8 +56,6 @@ backends:
   - TOKENSPEED_MLA  # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)
 
 device: "cuda:0"
-repeats: 100
-warmup_iters: 10
 profile_memory: true
 
 # Backend-specific tuning

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/configs/mla_mixed_batch.yaml

@@ -51,8 +51,6 @@ backends:
   - FLASHMLA         # Hopper only
 
 device: "cuda:0"
-repeats: 5
-warmup_iters: 3
 profile_memory: true
 
 # Analyze chunked prefill workspace size impact

benchmarks/attention_benchmarks/configs/mla_prefill.yaml

@@ -124,5 +124,3 @@ prefill_backends:
   - tokenspeed
 
 device: "cuda:0"
-repeats: 20
-warmup_iters: 5

benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml

@@ -53,6 +53,4 @@ backends:
   - FLASHINFER_MLA_SPARSE
 
 device: "cuda:0"
-repeats: 100
-warmup_iters: 10
 profile_memory: true

benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml

@@ -57,6 +57,4 @@ backends:
   - FLASHINFER_MLA_SPARSE
 
 device: "cuda:0"
-repeats: 10
-warmup_iters: 3
 profile_memory: true

benchmarks/attention_benchmarks/configs/reorder_threshold.yaml

@@ -63,8 +63,6 @@ model:
 
 # Benchmark settings
 device: "cuda:0"
-repeats: 15          # More repeats for spec decode variance
-warmup_iters: 5
 profile_memory: false
 
 # Output

benchmarks/attention_benchmarks/configs/speculative_decode.yaml

@@ -49,8 +49,6 @@ backends:
 
 # Benchmark settings
 device: "cuda:0"
-repeats: 10  # More repeats for statistical significance
-warmup_iters: 5
 profile_memory: false
 
 # Test these threshold values for optimization

benchmarks/attention_benchmarks/configs/standard_attention.yaml

@@ -43,6 +43,4 @@ backends:
   - FLASHINFER
 
 device: "cuda:0"
-repeats: 5
-warmup_iters: 3
 profile_memory: false

benchmarks/attention_benchmarks/configs/standard_decode.yaml

@@ -0,0 +1,142 @@
+# Standard attention decode benchmark configuration
+# Sweeps num_q_heads and num_kv_heads to isolate effects of:
+#   1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)
+#   2. Absolute head count (fixed 4:1 ratio, vary scale)
+
+model:
+  num_layers: 32
+  num_q_heads: 32 # Base value, overridden by sweep
+  num_kv_heads: 8 # Base value, overridden by sweep
+  head_dim: 128
+  block_size: 16
+
+# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and
+# head_dim where it differs from the base (128). Head counts are per-GPU
+# (i.e. after TP sharding).
+#
+# Group A — vary GQA ratio (fixed q=32, head_dim=128):
+#   32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)
+#
+# Groups B-E — real model configs at various TP degrees:
+#   Model              head_dim  Full    TP2      TP4      TP8
+#   Llama 3 8B           128    32:8   16:4      8:2      4:1
+#   Llama 3 70B          128    64:8   32:4     16:2      8:1
+#   GPT-OSS 120B          64    64:8   32:4     16:2      8:1
+#   Llama 3 405B         128   128:8   64:4     32:2     16:1
+model_parameter_sweep:
+  values:
+    # --- head_dim=128 (Llama 3 family) ---
+    - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1
+    - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1
+    - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8
+    - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4
+    - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2
+    - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1
+    - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8
+    - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4
+    - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2
+    - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1
+    # --- head_dim=64 (GPT-OSS 120B) ---
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1
+  label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"
+
+batch_specs:
+  # ---- batch_size x seq_len grid (decode: q_len=1) ----
+  # Small grid for quick iteration. Uncomment for full sweep.
+
+  # Batch size 1
+  - "q1s1k"
+  - "q1s512"
+  - "q1s2k"
+  - "q1s4k"
+  - "q1s8k"
+  - "q1s16k"
+  - "q1s32k"
+
+  # Batch size 2
+  - "2q1s512"
+  - "2q1s1k"
+  - "2q1s2k"
+  - "2q1s4k"
+  - "2q1s8k"
+  - "2q1s16k"
+  - "2q1s32k"
+
+  # Batch size 4
+  - "4q1s512"
+  - "4q1s1k"
+  - "4q1s2k"
+  - "4q1s4k"
+  - "4q1s8k"
+  - "4q1s16k"
+  - "4q1s32k"
+
+  # Batch size 8
+  - "8q1s1k"
+  - "8q1s512"
+  - "8q1s2k"
+  - "8q1s4k"
+  - "8q1s8k"
+  - "8q1s16k"
+  - "8q1s32k"
+
+  # Batch size 16
+  - "16q1s512"
+  - "16q1s1k"
+  - "16q1s2k"
+  - "16q1s4k"
+  - "16q1s8k"
+  - "16q1s16k"
+  - "16q1s32k"
+
+  # Batch size 32
+  - "32q1s512"
+  - "32q1s1k"
+  - "32q1s2k"
+  - "32q1s4k"
+  - "32q1s8k"
+  - "32q1s16k"
+  - "32q1s32k"
+
+  # Batch size 64
+  - "64q1s1k"
+  - "64q1s512"
+  - "64q1s2k"
+  - "64q1s4k"
+  - "64q1s8k"
+  - "64q1s16k"
+  - "64q1s32k"
+
+  # Batch size 128
+  - "128q1s512"
+  - "128q1s1k"
+  - "128q1s2k"
+  - "128q1s4k"
+  - "128q1s8k"
+  - "128q1s16k"
+  - "128q1s32k"
+
+  # Batch size 256
+  - "256q1s1k"
+  - "256q1s512"
+  - "256q1s2k"
+  - "256q1s4k"
+  - "256q1s8k"
+  - "256q1s16k"
+  - "256q1s32k"
+
+# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
+backends:
+  - FLASH_ATTN
+  - TRITON_ATTN
+  - FLASHINFER
+
+device: "cuda:0"
+profile_memory: false

benchmarks/attention_benchmarks/configs/standard_prefill.yaml

@@ -0,0 +1,108 @@
+# Standard attention prefill benchmark configuration
+# Sweeps num_q_heads and num_kv_heads to isolate effects of:
+#   1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)
+#   2. Absolute head count (fixed 4:1 ratio, vary scale)
+
+model:
+  num_layers: 32
+  num_q_heads: 32 # Base value, overridden by sweep
+  num_kv_heads: 8 # Base value, overridden by sweep
+  head_dim: 128
+  block_size: 16
+
+# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and
+# head_dim where it differs from the base (128). Head counts are per-GPU
+# (i.e. after TP sharding).
+#
+# Group A — vary GQA ratio (fixed q=32, head_dim=128):
+#   32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)
+#
+# Groups B-E — real model configs at various TP degrees:
+#   Model              head_dim  Full    TP2      TP4      TP8
+#   Llama 3 8B           128    32:8   16:4      8:2      4:1
+#   Llama 3 70B          128    64:8   32:4     16:2      8:1
+#   GPT-OSS 120B          64    64:8   32:4     16:2      8:1
+#   Llama 3 405B         128   128:8   64:4     32:2     16:1
+model_parameter_sweep:
+  values:
+    # --- head_dim=128 (Llama 3 family) ---
+    - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1
+    - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1
+    - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8
+    - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4
+    - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2
+    - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1
+    - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8
+    - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4
+    - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2
+    - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1
+    # --- head_dim=64 (GPT-OSS 120B) ---
+    - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8
+    - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4
+    - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2
+    - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1
+  label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"
+
+batch_specs:
+  # ---- batch_size x prefill_len grid (prefill: q_len == seq_len) ----
+  # Total tokens = batch_size * prefill_len, and prefill compute scales with
+  # prefill_len^2, so the largest cells are expensive. Trim batch sizes or
+  # lengths for quick iteration.
+
+  # Batch size 1
+  - "q512"
+  - "q1k"
+  - "q2k"
+  - "q4k"
+  - "q8k"
+  - "q16k"
+  - "q32k"
+
+  # Batch size 2
+  - "2q512"
+  - "2q1k"
+  - "2q2k"
+  - "2q4k"
+  - "2q8k"
+  - "2q16k"
+  - "2q32k"
+
+  # Batch size 4
+  - "4q512"
+  - "4q1k"
+  - "4q2k"
+  - "4q4k"
+  - "4q8k"
+  - "4q16k"
+  - "4q32k"
+
+  # Batch size 8
+  - "8q512"
+  - "8q1k"
+  - "8q2k"
+  - "8q4k"
+  - "8q8k"
+  - "8q16k"
+  - "8q32k"
+
+  # Batch size 16
+  - "16q512"
+  - "16q1k"
+  - "16q2k"
+  - "16q4k"
+  - "16q8k"
+  - "16q16k"
+  - "16q32k"
+
+# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER
+backends:
+  - FLASH_ATTN
+  - TRITON_ATTN
+  - FLASHINFER
+
+device: "cuda:0"
+profile_memory: false

benchmarks/attention_benchmarks/mla_runner.py

@@ -8,6 +8,8 @@ This module provides helpers for running MLA backends without
 needing full VllmConfig integration.
 """
 
+import statistics
+
 import numpy as np
 import torch
 from batch_spec import parse_batch_spec
@@ -17,6 +19,8 @@ from common import (
     MockIndexer,
     MockKVBProj,
     MockLayer,
+    run_do_bench,
+    run_ncu_profile,
     setup_mla_dims,
 )
 
@@ -704,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.
@@ -717,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
@@ -820,63 +831,86 @@ def _run_single_benchmark(
             num_prefill, mla_dims, query_fmt, device, torch.bfloat16
         )
 
-    # Build forward function
+    # 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)
 
-    # Warmup
-    for _ in range(config.warmup_iters):
-        forward_fn()
-    torch.accelerator.synchronize()
-
-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            forward_fn()
-        benchmark_fn = graph.replay
-    else:
-        benchmark_fn = forward_fn
-
-    # Benchmark
-    times = []
-    for _ in range(config.repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
+    def benchmark_fn():
         for _ in range(config.num_layers):
-            benchmark_fn()
-        end.record()
+            forward_fn()
 
-        torch.accelerator.synchronize()
-        elapsed_ms = start.elapsed_time(end)
-        times.append(elapsed_ms / 1000.0 / config.num_layers)
+    if config.ncu_profile:
+        run_ncu_profile(benchmark_fn)
+        return BenchmarkResult(
+            config=config,
+            mean_time=0.0,
+            median_time=0.0,
+            std_time=0.0,
+            min_time=0.0,
+            max_time=0.0,
+            throughput_tokens_per_sec=0.0,
+        )
+
+    all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
+
+    # Convert ms to seconds per layer
+    times = [t / 1000.0 / config.num_layers for t in all_ms]
+    mean_time = statistics.mean(times)
 
-    mean_time = float(np.mean(times))
     return BenchmarkResult(
         config=config,
         mean_time=mean_time,
-        std_time=float(np.std(times)),
-        min_time=float(np.min(times)),
-        max_time=float(np.max(times)),
+        median_time=statistics.median(times),
+        std_time=statistics.stdev(times) if len(times) > 1 else 0.0,
+        min_time=min(times),
+        max_time=max(times),
         throughput_tokens_per_sec=total_q / mean_time if mean_time > 0 else 0,
     )
 
@@ -886,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.
@@ -1025,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)
 
@@ -1052,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.
@@ -1071,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)
@@ -1095,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/attention_benchmarks/runner.py

@@ -9,13 +9,20 @@ This module provides helpers for running standard attention backends
 """
 
 import logging
+import statistics
 import types
 from contextlib import contextmanager
 
-import numpy as np
 import torch
 from batch_spec import parse_batch_spec, reorder_for_flashinfer
-from common import BenchmarkConfig, BenchmarkResult, MockLayer, get_attention_scale
+from common import (
+    BenchmarkConfig,
+    BenchmarkResult,
+    MockLayer,
+    get_attention_scale,
+    run_do_bench,
+    run_ncu_profile,
+)
 
 from vllm.config import (
     CacheConfig,
@@ -208,6 +215,13 @@ def _create_backend_impl(
 
     scale = get_attention_scale(config.head_dim)
 
+    # Set v_head_dim for diff-headdim backends. Always reset (defaulting to
+    # head_dim) so a prior run's value doesn't leak into this one via the
+    # backend's class-level state.
+    if hasattr(backend_class, "set_head_size_v"):
+        v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
+        backend_class.set_head_size_v(v_dim)
+
     impl = backend_class.get_impl_cls()(
         num_heads=config.num_q_heads,
         head_size=config.head_dim,
@@ -300,6 +314,7 @@ def _create_input_tensors(
         from vllm.platforms import current_platform
 
         q_dtype = current_platform.fp8_dtype()
+    v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
     q_list = [
         torch.randn(
             total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
@@ -313,9 +328,7 @@ def _create_input_tensors(
         for _ in range(config.num_layers)
     ]
     v_list = [
-        torch.randn(
-            total_q, config.num_kv_heads, config.head_dim, device=device, dtype=dtype
-        )
+        torch.randn(total_q, config.num_kv_heads, v_dim, device=device, dtype=dtype)
         for _ in range(config.num_layers)
     ]
     return q_list, k_list, v_list
@@ -389,14 +402,17 @@ def _run_single_benchmark(
     device: torch.device,
     dtype: torch.dtype,
 ) -> tuple:
-    """Run single benchmark iteration with warmup and timing loop."""
-    total_q = q_list[0].shape[0]
-    out = torch.empty(
-        total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype
-    )
+    """Run single benchmark using triton's do_bench_cudagraph/do_bench.
 
-    # Warmup
-    for _ in range(config.warmup_iters):
+    Returns:
+        (timing_stats, mem_stats) where timing_stats is a dict with
+        mean/std/min/max in seconds per layer.
+    """
+    total_q = q_list[0].shape[0]
+    v_dim = config.v_head_dim if config.v_head_dim is not None else config.head_dim
+    out = torch.empty(total_q, config.num_q_heads, v_dim, device=device, dtype=dtype)
+
+    def benchmark_fn():
         for i in range(config.num_layers):
             impl.forward(
                 layer,
@@ -407,52 +423,22 @@ def _run_single_benchmark(
                 attn_metadata,
                 output=out,
             )
-    torch.accelerator.synchronize()
 
-    # Optionally capture a CUDA graph after warmup.
-    # Graph replay eliminates CPU launch overhead so timings reflect pure
-    # kernel time.
-    if config.use_cuda_graphs:
-        graph = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(graph):
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-        benchmark_fn = graph.replay
+    if config.ncu_profile:
+        run_ncu_profile(benchmark_fn)
+        timing_stats = dict.fromkeys(("mean", "median", "std", "min", "max"), 0.0)
     else:
+        all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
 
-        def benchmark_fn():
-            for i in range(config.num_layers):
-                impl.forward(
-                    layer,
-                    q_list[i],
-                    k_list[i],
-                    v_list[i],
-                    cache_list[i],
-                    attn_metadata,
-                    output=out,
-                )
-
-    # Benchmark
-    times = []
-    for _ in range(config.repeats):
-        start = torch.cuda.Event(enable_timing=True)
-        end = torch.cuda.Event(enable_timing=True)
-
-        start.record()
-        benchmark_fn()
-        end.record()
-
-        torch.accelerator.synchronize()
-        elapsed_ms = start.elapsed_time(end)
-        times.append(elapsed_ms / 1000.0 / config.num_layers)  # seconds per layer
+        # Convert ms to seconds per layer
+        times = [t / 1000.0 / config.num_layers for t in all_ms]
+        timing_stats = {
+            "mean": statistics.mean(times),
+            "std": statistics.stdev(times) if len(times) > 1 else 0.0,
+            "min": min(times),
+            "max": max(times),
+            "median": statistics.median(times),
+        }
 
     mem_stats = {}
     if config.profile_memory:
@@ -461,7 +447,7 @@ def _run_single_benchmark(
             "reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
         }
 
-    return times, mem_stats
+    return timing_stats, mem_stats
 
 
 # ============================================================================
@@ -541,6 +527,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 common_attn_metadata=common_metadata,
             )
 
+            # Override num_splits for split-K testing (FlashAttention only)
+            if config.num_splits is not None and hasattr(
+                attn_metadata, "max_num_splits"
+            ):
+                attn_metadata.max_num_splits = config.num_splits
+
             # Only quantize queries when the impl supports it
             quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
                 impl, "supports_quant_query_input", False
@@ -553,7 +545,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 config, max_num_blocks, backend_class, device, dtype
             )
 
-            times, mem_stats = _run_single_benchmark(
+            timing_stats, mem_stats = _run_single_benchmark(
                 config,
                 impl,
                 layer,
@@ -566,15 +558,16 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 dtype,
             )
 
-    mean_time = np.mean(times)
+    mean_time = timing_stats["mean"]
     throughput = total_q / mean_time if mean_time > 0 else 0
 
     return BenchmarkResult(
         config=config,
         mean_time=mean_time,
-        std_time=np.std(times),
-        min_time=np.min(times),
-        max_time=np.max(times),
+        median_time=timing_stats["median"],
+        std_time=timing_stats["std"],
+        min_time=timing_stats["min"],
+        max_time=timing_stats["max"],
         throughput_tokens_per_sec=throughput,
         memory_allocated_mb=mem_stats.get("allocated_mb"),
         memory_reserved_mb=mem_stats.get("reserved_mb"),

benchmarks/benchmark_hidden_state_extraction.py

@@ -92,7 +92,6 @@ def run_baseline(
     llm = LLM(
         model=model,
         enable_prefix_caching=False,
-        enable_chunked_prefill=False,
         **extra_args,
     )
     sampling_params = SamplingParams(max_tokens=1)
@@ -194,7 +193,6 @@ async def _run_extraction_async(
     engine_args = AsyncEngineArgs(
         model=model,
         enable_prefix_caching=False,
-        enable_chunked_prefill=False,
         max_num_batched_tokens=40960,
         max_model_len=40960,
         speculative_config={

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))

benchmarks/kernels/benchmark_flydsl_moe_w4a16.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2025 FlyDSL Project Contributors
+
+import json
+import os
+
+import torch
+from aiter.test_common import run_perftest
+
+from vllm.model_executor.layers.fused_moe import fused_experts
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    int4_w4a16_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa: E501
+    compressed_tensors_moe_w4a16_flydsl,
+)
+from vllm.platforms import current_platform
+
+RoutingBuffers = tuple[
+    torch.Tensor,  # sorted_token_ids
+    torch.Tensor,  # sorted_weights
+    torch.Tensor,  # sorted_expert_ids
+    torch.Tensor,  # num_valid_ids (shape [1], i32)
+    int,  # sorted_size
+    int,  # blocks
+]
+
+MODEL_PARAMS_TO_TUNE = [
+    # (num_experts, inter_dim, hidden_size, topk)
+    (384, 256, 7168, 8),  # Kimi K2.5 TP=8
+    (384, 512, 7168, 8),  # Kimi K2.5 TP=4
+]
+
+NUM_TOKENS_TO_TUNE = [
+    1,
+    2,
+    4,
+    8,
+    16,
+    24,
+    32,
+    48,
+    64,
+    128,
+    256,
+    512,
+    1024,
+    2048,
+    4096,
+    8192,
+]
+
+TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256]
+TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256]
+TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
+TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256]
+TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512]
+
+TILE_CONFIGS = []
+for tile_m in TILE_M_SEARCH_SPACE:
+    for tile_n in TILE_N_SEARCH_SPACE:
+        for tile_k in TILE_K_SEARCH_SPACE:
+            for tile_n2 in TILE_N2_SEARCH_SPACE:
+                for tile_k2 in TILE_K2_SEARCH_SPACE:
+                    TILE_CONFIGS.append(
+                        {
+                            "tile_m": tile_m,
+                            "tile_n": tile_n,
+                            "tile_k": tile_k,
+                            "tile_n2": tile_n2,
+                            "tile_k2": tile_k2,
+                        }
+                    )
+
+
+def tune_flydsl_moe_w4a16(
+    device: str = "cuda", num_iters: int = 100, num_warmup: int = 10
+):
+    packed_factor = 8
+    w13_num_shards = 2
+    params_dtype = torch.bfloat16
+    group_size = 32
+    scale_factor = 0.01
+
+    for model_params in MODEL_PARAMS_TO_TUNE:
+        num_experts = model_params[0]
+        inter_dim = model_params[1]
+        hidden_size = model_params[2]
+        topk = model_params[3]
+        print(
+            f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, "
+            f"hidden_size={hidden_size}, topk={topk}...\n"
+        )
+
+        w2_scales_size = inter_dim
+        num_groups_w2 = w2_scales_size // group_size
+        num_groups_w13 = hidden_size // group_size
+
+        w13_weight = torch.randint(
+            0,
+            255,
+            (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim),
+            dtype=torch.int32,
+            device=device,
+        )
+
+        w2_weight = torch.randint(
+            0,
+            255,
+            (num_experts, inter_dim // packed_factor, hidden_size),
+            dtype=torch.int32,
+            device=device,
+        )
+        w13_scale = scale_factor * torch.randn(
+            num_experts,
+            num_groups_w13,
+            w13_num_shards * inter_dim,
+            dtype=params_dtype,
+            device=device,
+        )
+        w2_scale = scale_factor * torch.randn(
+            num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device
+        )
+
+        w13 = w13_weight
+        w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13)
+        w13 = w13.view(-1).contiguous()
+
+        w2 = w2_weight
+        w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2)
+        w2 = w2.view(-1).contiguous()
+
+        w13_scale_flydsl = w13_scale
+        w2_scale_flydsl = w2_scale
+
+        if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1:
+            E, G, N = w13_scale.shape
+            w13_scale_flydsl = (
+                w13_scale_flydsl.view(E, G // 2, 2, N)
+                .permute(0, 1, 3, 2)
+                .contiguous()
+                .view(-1)
+                .contiguous()
+            )
+        elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1:
+            w13_scale_flydsl = w13_scale_flydsl.squeeze(1)
+
+        if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1:
+            E, G, N = w2_scale.shape
+            w2_scale_flydsl = (
+                w2_scale_flydsl.view(E, G // 2, 2, N)
+                .permute(0, 1, 3, 2)
+                .contiguous()
+                .view(-1)
+                .contiguous()
+            )
+        elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1:
+            w2_scale_flydsl = w2_scale_flydsl.squeeze(1)
+
+        w13_scale_flydsl = w13_scale_flydsl.contiguous()
+        w2_scale_flydsl = w2_scale_flydsl.contiguous()
+
+        w13.is_shuffled = True
+        w2.is_shuffled = True
+
+        w13_weight_scale = w13_scale.transpose(1, 2).contiguous()
+        w2_weight_scale = w2_scale.transpose(1, 2).contiguous()
+        w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8)
+        w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8)
+
+        moe_quant_config = int4_w4a16_moe_quant_config(
+            w1_scale=w13_weight_scale,
+            w2_scale=w2_weight_scale,
+            w1_zp=None,
+            w2_zp=None,
+            block_shape=[0, group_size],
+        )
+
+        tuned_config = {}
+
+        for num_tokens in NUM_TOKENS_TO_TUNE:
+            score = torch.rand(
+                (num_tokens, num_experts), device=device, dtype=torch.float32
+            )
+            topk_vals, topk_ids = torch.topk(score, k=topk, dim=1)
+            topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32)
+            x = torch.randn(
+                (num_tokens, hidden_size), dtype=torch.bfloat16, device=device
+            )
+            us_best = float("inf")
+            for tile_config in TILE_CONFIGS:
+                try:
+                    tile_m = tile_config["tile_m"]
+                    tile_n = tile_config["tile_n"]
+                    tile_k = tile_config["tile_k"]
+                    tile_n2 = tile_config["tile_n2"]
+                    tile_k2 = tile_config["tile_k2"]
+
+                    model_dim = x.shape[1]
+                    assert model_dim % 64 == 0
+                    assert model_dim % tile_k == 0
+                    assert inter_dim % tile_n == 0
+                    assert model_dim % tile_n2 == 0
+                    assert inter_dim % tile_k2 == 0
+                    assert ((tile_m * tile_k2) % 256) == 0
+                    bytes_per_thread_x = (tile_m * tile_k2) // 256
+                    assert (bytes_per_thread_x % 4) == 0
+
+                    out, _us = run_perftest(
+                        fused_flydsl_moe,
+                        x,
+                        w13,
+                        w2,
+                        num_experts,
+                        inter_dim,
+                        topk_weights,
+                        topk_ids,
+                        num_iters=num_iters,
+                        num_warmup=num_warmup,
+                        w1_scale=w13_scale_flydsl,
+                        w2_scale=w2_scale_flydsl,
+                        topk=topk_weights.shape[-1],
+                        group_size=group_size,
+                        doweight_stage1=False,
+                        scale_is_bf16=True,
+                        config=tile_config,
+                    )
+                    torch.accelerator.synchronize()
+                except Exception:
+                    torch.accelerator.synchronize()
+                    continue
+                else:
+                    us = _us.item()
+                    if us < us_best:
+                        out_ref = fused_experts(
+                            x,
+                            w13_weight_packed,
+                            w2_weight_packed,
+                            topk_weights=topk_weights,
+                            topk_ids=topk_ids,
+                            activation=MoEActivation.SILU,
+                            apply_router_weight_on_input=False,
+                            global_num_experts=num_experts,
+                            expert_map=None,
+                            quant_config=moe_quant_config,
+                        )
+                        try:
+                            assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1)
+                        except Exception:
+                            continue
+                        else:
+                            print(
+                                f"For [num_tokens={num_tokens}, num_experts={num_experts}, "  # noqa: E501
+                                f"inter_dim={inter_dim}] found new best "  # noqa: E501
+                                f"config={tile_config}, us={us:0.3f}"
+                            )
+                            us_best = us
+                            tuned_config[str(num_tokens)] = tile_config
+            device_name = current_platform.get_device_name().replace(" ", "_")
+            tuned_config_file_name = (
+                f"E={num_experts},N={inter_dim},device_name={device_name},"
+                f"dtype=int4_w4a16,backend=flydsl.json"
+            )
+            tuner_dir_path = os.path.dirname(os.path.realpath(__file__))
+            store_path = os.path.join(tuner_dir_path, tuned_config_file_name)
+            with open(store_path, "w") as f:
+                json.dump(tuned_config, f, indent=4)
+            print(
+                f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n"  # noqa: E501
+            )
+
+
+if __name__ == "__main__":
+    tune_flydsl_moe_w4a16(device="cuda")

benchmarks/kernels/benchmark_moe.py

@@ -792,6 +792,12 @@ def get_model_params(config):
         topk = text_config.num_experts_per_tok
         intermediate_size = text_config.moe_intermediate_size
         hidden_size = text_config.hidden_size
+    elif architecture == "DiffusionGemmaForBlockDiffusion":
+        text_config = config.get_text_config()
+        E = text_config.num_experts
+        topk = text_config.top_k_experts
+        intermediate_size = text_config.moe_intermediate_size
+        hidden_size = text_config.hidden_size
     elif architecture == "HunYuanMoEV1ForCausalLM":
         E = config.num_experts
         topk = config.moe_topk[0]

benchmarks/kv_cache_watermark.sh

@@ -0,0 +1,248 @@
+#!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Reproducible demonstration of the KV cache watermark (`--watermark`) for
+# reducing preemption thrashing.
+#
+# The watermark is the fraction of total KV cache blocks the scheduler keeps
+# free when admitting a waiting/preempted request into the running queue.
+#
+# Why this workload triggers thrashing:
+#   Requests are admitted based on the KV cache they need *at admission time*.
+#   With `--scheduler-reserve-full-isl` (default) the input length is reserved up
+#   front, but the *output* length is unknown and unreserved. A decode-heavy
+#   workload (output >> input) at high concurrency therefore over-admits while
+#   requests are short, then runs out of KV cache as they all grow during decode
+#   -> the scheduler preempts (recompute) recently-admitted requests, re-prefills
+#   them later, and repeats. The watermark keeps a block of KV cache free so
+#   running requests can grow into it instead of triggering this churn.
+#
+# This script launches `vllm serve` under a deliberately KV-constrained config
+# and a decode-heavy workload, sweeping the watermark across several values, and
+# reports the preemption count (scraped from /metrics), throughput, and latency
+# percentiles for each. It then plots the results.
+#
+# Default workload: concurrency 200, input ~300 tokens, output ~4000 tokens
+# (+/- 20% variance), sized to run each config for ~5 minutes.
+#
+# Usage:
+#   benchmarks/kv_cache_watermark.sh
+#   MODEL=Qwen/Qwen2.5-14B-Instruct TP=2 benchmarks/kv_cache_watermark.sh
+#
+# Run inside the vLLM virtualenv (so `vllm` and `python` resolve to it).
+set -euo pipefail
+
+# ---- Config (override via environment) -------------------------------------
+MODEL=${MODEL:-Qwen/Qwen2.5-7B-Instruct}
+TP=${TP:-1}
+PORT=${PORT:-8000}
+URL="http://127.0.0.1:${PORT}"
+# Constrain the KV cache to a *near-critical* size: large enough that the engine
+# can run stably, but small enough that greedy over-admission tips it into
+# preemption thrashing. (Independent of GPU size, so the demo is reproducible.)
+# At the default workload this fits ~1.5x the mean concurrent KV demand.
+KV_CACHE_MEMORY_GB=${KV_CACHE_MEMORY_GB:-16}
+MAX_MODEL_LEN=${MAX_MODEL_LEN:-8192}
+MAX_NUM_SEQS=${MAX_NUM_SEQS:-256}
+# Optional weight loader (e.g. fastsafetensors on the GCP cluster).
+LOAD_FORMAT=${LOAD_FORMAT:-auto}
+# Decode-heavy workload: moderate input, long output, with length variance. The
+# long output means preempted requests have generated a lot before eviction, so
+# resuming them re-prefills a long sequence (high recomputation cost).
+INPUT_LEN=${INPUT_LEN:-1000}
+OUTPUT_LEN=${OUTPUT_LEN:-5000}
+RANGE_RATIO=${RANGE_RATIO:-0.2}
+CONCURRENCY=${CONCURRENCY:-128}
+# Enough prompts to keep each config saturated for ~5+ minutes.
+NUM_PROMPTS=${NUM_PROMPTS:-450}
+OUTDIR=${OUTDIR:-./watermark_bench_results}
+# Watermark fractions compared. "label value" per line; value=0 disables it.
+CONFIGS=${CONFIGS:-"off    0
+w0.02  0.02
+w0.05  0.05
+w0.10  0.10
+w0.15  0.15"}
+
+KV_CACHE_MEMORY_BYTES=$((KV_CACHE_MEMORY_GB * 1024 * 1024 * 1024))
+mkdir -p "$OUTDIR"
+
+SERVER_PID=""
+cleanup() { [[ -n "$SERVER_PID" ]] && kill "$SERVER_PID" 2>/dev/null || true; }
+trap cleanup EXIT
+
+scrape_preemptions() {
+  # Sum the vllm:num_preemptions_total counter across engines.
+  python - "${URL}/metrics" <<'PY'
+import sys, urllib.request
+total = 0.0
+try:
+    body = urllib.request.urlopen(sys.argv[1], timeout=10).read().decode("utf-8", "replace")
+    for line in body.splitlines():
+        if line.startswith("vllm:num_preemptions_total"):
+            total += float(line.rsplit(" ", 1)[-1])
+except Exception as e:  # noqa: BLE001
+    print(f"scrape error: {e}", file=sys.stderr)
+print(int(total))
+PY
+}
+
+wait_for_server() {
+  for _ in $(seq 1 300); do
+    if curl -s "${URL}/health" >/dev/null 2>&1; then return 0; fi
+    if ! kill -0 "$SERVER_PID" 2>/dev/null; then
+      echo "ERROR: server process exited during startup" >&2; return 1
+    fi
+    sleep 5
+  done
+  echo "ERROR: server did not become ready" >&2; return 1
+}
+
+run_one() {
+  local label=$1 watermark=$2
+  echo
+  echo "==================== watermark: ${label} (${watermark}) ===================="
+  vllm serve "$MODEL" \
+    --tensor-parallel-size "$TP" \
+    --load-format "$LOAD_FORMAT" \
+    --kv-cache-memory-bytes "$KV_CACHE_MEMORY_BYTES" \
+    --max-model-len "$MAX_MODEL_LEN" \
+    --max-num-seqs "$MAX_NUM_SEQS" \
+    --no-enable-prefix-caching \
+    --watermark "$watermark" \
+    --port "$PORT" >"${OUTDIR}/serve_${label}.log" 2>&1 &
+  SERVER_PID=$!
+  wait_for_server
+  sleep 5
+
+  local pre post
+  pre=$(scrape_preemptions)
+  vllm bench serve \
+    --backend vllm \
+    --base-url "$URL" \
+    --model "$MODEL" \
+    --dataset-name random \
+    --random-input-len "$INPUT_LEN" \
+    --random-output-len "$OUTPUT_LEN" \
+    --random-range-ratio "$RANGE_RATIO" \
+    --ignore-eos \
+    --num-prompts "$NUM_PROMPTS" \
+    --max-concurrency "$CONCURRENCY" \
+    --percentile-metrics "ttft,tpot,itl,e2el" \
+    --metric-percentiles "50,90,99" \
+    --save-result \
+    --result-dir "$OUTDIR" \
+    --result-filename "bench_${label}.json"
+  post=$(scrape_preemptions)
+  echo "${label} ${watermark} $((post - pre))" >>"${OUTDIR}/preemptions.txt"
+
+  kill "$SERVER_PID" 2>/dev/null || true
+  for _ in $(seq 1 60); do curl -s "${URL}/health" >/dev/null 2>&1 || break; sleep 2; done
+  SERVER_PID=""
+  sleep 10
+}
+
+: >"${OUTDIR}/preemptions.txt"
+while read -r label watermark; do
+  [[ -z "${label:-}" ]] && continue
+  run_one "$label" "$watermark"
+done <<<"$CONFIGS"
+
+echo
+echo "==================== summary ===================="
+python - "$OUTDIR" <<'PY'
+import json, os, sys
+outdir = sys.argv[1]
+pre = {}
+order = []
+for line in open(os.path.join(outdir, "preemptions.txt")):
+    label, watermark, n = line.split()
+    pre[label] = (float(watermark), int(n))
+    order.append(label)
+
+def g(d, *names):
+    for n in names:
+        if d.get(n) is not None:
+            return d[n]
+    return float("nan")
+
+cols = ["watermark", "frac", "preempt", "out_tok/s", "req/s",
+        "TTFT_p50", "TTFT_p99", "ITL_p99", "E2EL_p50"]
+print("  ".join(f"{c:>10}" for c in cols))
+rows = []
+for label in order:
+    watermark, n = pre[label]
+    d = json.load(open(os.path.join(outdir, f"bench_{label}.json")))
+    rows.append(dict(
+        label=label, watermark=watermark, preempt=n,
+        out_tok_s=g(d, "output_throughput"),
+        req_s=g(d, "request_throughput"),
+        ttft_p50=g(d, "p50_ttft_ms", "median_ttft_ms"),
+        ttft_p99=g(d, "p99_ttft_ms"),
+        itl_p99=g(d, "p99_itl_ms"),
+        e2el_p50=g(d, "p50_e2el_ms", "median_e2el_ms"),
+    ))
+    print("  ".join(f"{str(v):>10}" for v in [
+        label, watermark, n,
+        f"{rows[-1]['out_tok_s']:.0f}",
+        f"{rows[-1]['req_s']:.3f}",
+        f"{rows[-1]['ttft_p50']/1000:.2f}",
+        f"{rows[-1]['ttft_p99']/1000:.2f}",
+        f"{rows[-1]['itl_p99']:.2f}",
+        f"{rows[-1]['e2el_p50']/1000:.1f}",
+    ]))
+print("\n(TTFT/E2EL in seconds; ITL in ms. Lower preempt is better.)")
+
+# ---- Plot -------------------------------------------------------------------
+try:
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+except Exception as e:  # noqa: BLE001
+    print(f"\n(skip plot: matplotlib unavailable: {e})")
+    sys.exit(0)
+
+x = [r["watermark"] for r in rows]
+xt = [f"{r['watermark']:g}\n({r['label']})" for r in rows]
+idx = list(range(len(rows)))
+
+fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+fig.suptitle(
+    f"KV cache watermark sweep — {os.path.basename(os.path.abspath(outdir))}",
+    fontsize=12,
+)
+
+ax = axes[0][0]
+ax.bar(idx, [r["preempt"] for r in rows], color="tab:red")
+ax.set_title("Preemptions (lower is better)")
+ax.set_ylabel("preemptions")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+
+ax = axes[0][1]
+ax.plot(idx, [r["out_tok_s"] for r in rows], "o-", color="tab:green")
+ax.set_title("Output throughput (higher is better)")
+ax.set_ylabel("tokens/s")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+
+ax = axes[1][0]
+ax.plot(idx, [r["itl_p99"] for r in rows], "o-", color="tab:blue")
+ax.set_title("Inter-token latency p99 (lower is better)")
+ax.set_ylabel("ITL p99 (ms)")
+ax.set_xlabel("watermark fraction")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+
+ax = axes[1][1]
+ax.plot(idx, [r["ttft_p50"] / 1000 for r in rows], "o-", label="TTFT p50")
+ax.plot(idx, [r["ttft_p99"] / 1000 for r in rows], "o-", label="TTFT p99")
+ax.plot(idx, [r["e2el_p50"] / 1000 for r in rows], "o-", label="E2EL p50")
+ax.set_title("Latency (lower is better)")
+ax.set_ylabel("seconds")
+ax.set_xlabel("watermark fraction")
+ax.set_xticks(idx); ax.set_xticklabels(xt)
+ax.legend()
+
+fig.tight_layout(rect=(0, 0, 1, 0.95))
+out_png = os.path.join(outdir, "watermark_results.png")
+fig.savefig(out_png, dpi=120)
+print(f"\nWrote plot: {out_png}")
+PY

cmake/cpu_extension.cmake

@@ -166,6 +166,10 @@ elseif (S390_FOUND)
         "-mtune=native")
 elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
     message(STATUS "RISC-V detected")
+    if(DEFINED VLLM_RVV_VLEN AND NOT VLLM_RVV_VLEN GREATER 0)
+        message(FATAL_ERROR
+            "VLLM_RVV_VLEN must be a positive integer; got '${VLLM_RVV_VLEN}'")
+    endif()
     # VLLM_RVV_VLEN selects the target VLEN. Auto-detected from /proc/cpuinfo
     # by default; override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256.
     if(NOT DEFINED VLLM_RVV_VLEN)
@@ -189,8 +193,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
                 "RISC-V RVV is available but VLEN could not be auto-detected. "
                 "Please specify VLEN explicitly:\n"
                 "  -DVLLM_RVV_VLEN=128   (for VLEN=128 hardware)\n"
-                "  -DVLLM_RVV_VLEN=256   (for VLEN=256 hardware, e.g. Spacemit X100)\n"
-                "  -DVLLM_RVV_VLEN=0     (force scalar, no RVV)")
+                "  -DVLLM_RVV_VLEN=256   (for VLEN=256 hardware, e.g. Spacemit X100)")
         endif()
     endif()
     if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
@@ -219,7 +222,7 @@ endif()
 
 
 # Build oneDNN for GEMM kernels
-if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
+if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND OR RVV_FP16_FOUND OR RVV_BF16_FOUND)
     # Fetch and build Arm Compute Library (ACL) as oneDNN's backend for AArch64
     # TODO [fadara01]: remove this once ACL can be fetched and built automatically as a dependency of oneDNN
     set(ONEDNN_AARCH64_USE_ACL OFF CACHE BOOL "")
@@ -435,6 +438,12 @@ if(USE_ONEDNN)
         ${VLLM_EXT_SRC})
 endif()
 
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
+    set(VLLM_EXT_SRC
+        "csrc/cpu/sgl-kernels/gemm_int4.cpp"
+        ${VLLM_EXT_SRC})
+endif()
+
 if (ENABLE_X86_ISA)
     set(VLLM_EXT_SRC_SGL
         "csrc/cpu/sgl-kernels/conv.cpp"

cmake/external_projects/deepgemm.cmake

@@ -8,43 +8,73 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR})
   set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
 endif()
 
+# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download).
+# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW
+# disallows one-argument Populate(dep) after Declare; MakeAvailable would run
+# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load).
 if(DEEPGEMM_SRC_DIR)
-  FetchContent_Declare(
-    deepgemm
-    SOURCE_DIR ${DEEPGEMM_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  # cmake_path(ABSOLUTE_PATH <var> ...) reads the path from <var>; NORMALIZE is a
+  # flag (no trailing path argument). Resolve relative paths against vLLM root.
+  set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _deepgemm_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}")
+  if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}")
+    message(FATAL_ERROR
+      "DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'")
+  endif()
+  set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}")
+  message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
-  # This ref should be kept in sync with tools/install_deepgemm.sh
-  FetchContent_Declare(
-    deepgemm
-    GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
-    GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  # Keep in sync with tools/install_deepgemm.sh
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
+
+  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _deepgemm_fc_root)
+    set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src")
+  set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build")
+  set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild")
+
+  if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp")
+    set(deepgemm_SOURCE_DIR "${_deepgemm_src}")
+    set(deepgemm_BINARY_DIR "${_deepgemm_bin}")
+  else()
+    FetchContent_Populate(
+      deepgemm
+      SUBBUILD_DIR "${_deepgemm_sub}"
+      SOURCE_DIR "${_deepgemm_src}"
+      BINARY_DIR "${_deepgemm_bin}"
+      GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
+      GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
+      GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+  message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
 endif()
 
-# Use FetchContent_Populate (not MakeAvailable) to avoid processing
-# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls.
-FetchContent_GetProperties(deepgemm)
-if(NOT deepgemm_POPULATED)
-  FetchContent_Populate(deepgemm)
-endif()
-message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
-
-# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100
+# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream),
+# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x
+# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants.
 set(DEEPGEMM_SUPPORT_ARCHS)
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
   list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
 endif()
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
-  list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
-elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-  list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
+  else()
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f")
+  else()
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a")
+  endif()
 endif()
 
 cuda_archs_loose_intersection(DEEPGEMM_ARCHS

cmake/external_projects/fmha_sm100.cmake

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

cmake/external_projects/qutlass.cmake

@@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR})
   set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
 endif()
 
+# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid.
+# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees.
 if(QUTLASS_SRC_DIR)
-  FetchContent_Declare(
-    qutlass
-    SOURCE_DIR ${QUTLASS_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  set(_qutlass_user_src "${QUTLASS_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _qutlass_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(QUTLASS_SRC_DIR "${_qutlass_user_src}")
+  if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}")
+    message(FATAL_ERROR
+      "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'")
+  endif()
+  set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}")
+  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
-  FetchContent_Declare(
-    qutlass
-    GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git
-    GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
-endif()
+  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
 
-FetchContent_Populate(qutlass)
+  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _qutlass_fc_root)
+    set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_qutlass_src "${_qutlass_fc_root}/qutlass-src")
+  set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build")
+  set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild")
+
+  if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp")
+    set(qutlass_SOURCE_DIR "${_qutlass_src}")
+    set(qutlass_BINARY_DIR "${_qutlass_bin}")
+  else()
+    FetchContent_Populate(
+      qutlass
+      SUBBUILD_DIR "${_qutlass_sub}"
+      SOURCE_DIR "${_qutlass_src}"
+      BINARY_DIR "${_qutlass_bin}"
+      GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}"
+      GIT_TAG "${_QUTLASS_UPSTREAM_TAG}"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+endif()
 
 if(NOT qutlass_SOURCE_DIR)
   message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")
@@ -32,22 +54,35 @@ endif()
 message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-  cuda_archs_loose_intersection(QUTLASS_ARCHS "10.0f;12.0f" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0f" "${CUDA_ARCHS}")
 else()
-  cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;12.1a;10.0a;10.3a" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM120_ARCHS "12.0a;12.1a" "${CUDA_ARCHS}")
+  cuda_archs_loose_intersection(QUTLASS_SM100_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
+endif()
+
+# QUTLASS uses TARGET_CUDA_ARCH as a single preprocessor selector for all its
+# sources. Do not compile a mixed SM100/SM120 arch list with one selector; prefer
+# SM100 when both families are requested because that is the primary deployed
+# target for this extension today.
+if(QUTLASS_SM100_ARCHS)
+  set(QUTLASS_ARCHS "${QUTLASS_SM100_ARCHS}")
+  set(QUTLASS_TARGET_CC 100)
+  if(QUTLASS_SM120_ARCHS)
+    message(WARNING
+      "[QUTLASS] Both SM100 and SM120 archs were requested; selecting SM100 "
+      "because TARGET_CUDA_ARCH is a single compile-time selector.")
+  endif()
+elseif(QUTLASS_SM120_ARCHS)
+  set(QUTLASS_ARCHS "${QUTLASS_SM120_ARCHS}")
+  set(QUTLASS_TARGET_CC 120)
+else()
+  set(QUTLASS_ARCHS)
 endif()
 
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
-
-  if(QUTLASS_ARCHS MATCHES "10\\.(0a|3a|0f)")
-    set(QUTLASS_TARGET_CC 100)
-  elseif(QUTLASS_ARCHS MATCHES "12\\.[01][af]?")
-    set(QUTLASS_TARGET_CC 120)
-  else()
-    message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.")
-  endif()
-
   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
@@ -66,8 +101,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. "
@@ -79,12 +125,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>
@@ -99,4 +156,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()

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG dd62dac706b1cf7895bd99b18c6cb7e7e117ee25
+          GIT_TAG 803020a8fa15407871341d41eba4919ade2ee1ee
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

cmake/utils.cmake

@@ -487,9 +487,9 @@ endfunction()
 
 function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f;12.0f" "${TGT_CUDA_ARCHS}")
   else()
-    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
+    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${TGT_CUDA_ARCHS}")
   endif()
   set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
 endfunction()

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;
@@ -822,8 +852,8 @@ struct AttentionInput {
       logits_buffer_t *__restrict__ logits_buffer,                          \
       float *__restrict__ partial_q_buffer, float *__restrict__ max_buffer, \
       float *__restrict__ sum_buffer, int32_t *__restrict__ block_table,    \
-      const int32_t kv_tile_start_pos, const int32_t kv_tile_end_pos,       \
-      const int32_t kv_tile_token_num,                                      \
+      const int32_t kv_end_pos, const int32_t kv_tile_start_pos,            \
+      const int32_t kv_tile_end_pos, const int32_t kv_tile_token_num,       \
       const int64_t kv_cache_num_blocks_stride, const int32_t q_head_num,   \
       const int32_t q_token_num, const int32_t q_tile_start_pos,            \
       const int32_t q_heads_per_kv, const int32_t block_size,               \
@@ -834,7 +864,7 @@ struct AttentionInput {
 
 #define CPU_ATTENTION_PARAMS                                                  \
   q_heads_buffer, k_head_cache_ptr, v_head_cache_ptr, logits_buffer,          \
-      partial_q_buffer, max_buffer, sum_buffer, block_table,                  \
+      partial_q_buffer, max_buffer, sum_buffer, block_table, kv_end_pos,      \
       kv_tile_start_pos, kv_tile_end_pos, kv_tile_token_num,                  \
       kv_cache_num_blocks_stride, q_head_num, q_token_num, q_tile_start_pos,  \
       q_heads_per_kv, block_size, left_window_size, right_window_size, scale, \
@@ -917,6 +947,7 @@ class AttentionMainLoop {
     //  - max_buffer: [MaxQHeadNumPerIteration, 1], store max logits
     //  - sum_buffer: [MaxQHeadNumPerIteration, 1], store sum of exp
     //  - block_table
+    //  - kv_end_pos: un-aligned end position of KV cache
     //  - kv_tile_start_pos: start position of KV cache, aligned to
     //  BlockSizeAlignment
     //  - kv_tile_end_pos: end position of KV cache, aligned to
@@ -1043,7 +1074,7 @@ class AttentionMainLoop {
         }
 
         apply_mask(logits_buffer, kv_tile_token_num, q_tile_start_pos,
-                   kv_tile_start_pos, kv_tile_end_pos, q_token_num,
+                   kv_end_pos, kv_tile_start_pos, kv_tile_end_pos, q_token_num,
                    q_heads_per_kv, left_window_size, right_window_size);
 
         // if (debug_info){
@@ -1126,7 +1157,7 @@ class AttentionMainLoop {
 
     void apply_mask(logits_buffer_t* __restrict__ logits_buffer,
                     const int64_t logits_buffer_stride,
-                    const int32_t q_tile_start_pos,
+                    const int32_t q_tile_start_pos, const int32_t kv_end_pos,
                     const int32_t kv_tile_start_pos,
                     const int32_t kv_tile_end_pos, const int32_t q_token_num,
                     const int32_t q_heads_per_kv,
@@ -1154,7 +1185,7 @@ class AttentionMainLoop {
                            std::max(kv_tile_start_pos,
                                     curr_token_pos + sliding_window_right + 1));
           }
-          return pos;
+          return std::min(pos, kv_end_pos);
         }();
 
         int32_t left_invalid_token_num = left_kv_pos - kv_tile_start_pos;
@@ -1441,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;
 
@@ -1532,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;
@@ -1559,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 &&
@@ -1610,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);
@@ -1724,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;
 
@@ -1789,7 +1825,7 @@ class AttentionMainLoop {
                   attn_impl.template execute_attention<Attention>(
                       curr_q_heads_buffer, curr_k_cache, curr_v_cache,
                       logits_buffer, curr_partial_q_buffer, curr_max_buffer,
-                      curr_sum_buffer, curr_block_table,
+                      curr_sum_buffer, curr_block_table, kv_end_pos,
                       aligned_actual_kv_tile_pos_left,
                       aligned_actual_kv_tile_pos_right, actual_kv_token_num,
                       kv_cache_block_num_stride, q_tile_head_num,

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

@@ -57,6 +57,10 @@ typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
 typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
     __attribute__((riscv_rvv_vector_bits(1024)));
 
+// int8
+typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t
+    __attribute__((riscv_rvv_vector_bits(128)));
+
 // int32
 typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
     __attribute__((riscv_rvv_vector_bits(256)));

csrc/cpu/cpu_types_riscv_impl.hpp

@@ -9,10 +9,14 @@
 
 #include <algorithm>
 #include <cmath>
+#include <cstdint>
 #include <cstring>
 #include <iostream>
 #include <limits>
 #include <torch/all.h>
+
+#include "float_convert.hpp"
+
 namespace vec_op {
 
 // FP8 KV cache is not supported on RISC-V. These tag types and the
@@ -245,8 +249,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
     float tmp[8];
     for (int i = 0; i < 8; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
+      tmp[i] = bf16_to_float(u16[i]);
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
   }
@@ -256,9 +259,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save(void* ptr, int elem_num) const {
@@ -266,9 +267,7 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     RVVI(__riscv_vse32_v_f32, LMUL_256)(tmp, reg_fp32, 8);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -277,10 +276,8 @@ struct BF16Vec8 : public Vec<BF16Vec8> {
     uint8_t* u8 = static_cast<uint8_t*>(ptr);
     ptrdiff_t byte_stride = stride * sizeof(uint16_t);
     for (int i = 0; i < 8; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
+          float_to_bf16(tmp[i]);
     }
   }
 };
@@ -292,8 +289,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
     float tmp[16];
     for (int i = 0; i < 16; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
+      tmp[i] = bf16_to_float(u16[i]);
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_512)(tmp, 16);
   }
@@ -306,9 +302,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save(void* ptr, int elem_num) const {
@@ -316,9 +310,7 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     RVVI(__riscv_vse32_v_f32, LMUL_512)(tmp, reg_fp32, 16);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
   void save_strided(void* ptr, ptrdiff_t stride) const {
@@ -327,10 +319,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
     uint8_t* u8 = static_cast<uint8_t*>(ptr);
     ptrdiff_t byte_stride = stride * sizeof(uint16_t);
     for (int i = 0; i < 16; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
+          float_to_bf16(tmp[i]);
     }
   }
 };
@@ -343,8 +333,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     const uint16_t* u16 = static_cast<const uint16_t*>(ptr);
     float tmp[32];
     for (int i = 0; i < 32; ++i) {
-      uint32_t v = static_cast<uint32_t>(u16[i]) << 16;
-      std::memcpy(&tmp[i], &v, 4);
+      tmp[i] = bf16_to_float(u16[i]);
     }
     reg_fp32 = RVVI(__riscv_vle32_v_f32, LMUL_1024)(tmp, 32);
   }
@@ -371,9 +360,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
 
@@ -382,9 +369,7 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     RVVI(__riscv_vse32_v_f32, LMUL_1024)(tmp, reg_fp32, 32);
     uint16_t* u16 = static_cast<uint16_t*>(ptr);
     for (int i = 0; i < elem_num; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      u16[i] = static_cast<uint16_t>(v >> 16);
+      u16[i] = float_to_bf16(tmp[i]);
     }
   }
 
@@ -394,10 +379,8 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
     uint8_t* u8 = static_cast<uint8_t*>(ptr);
     ptrdiff_t byte_stride = stride * sizeof(uint16_t);
     for (int i = 0; i < 32; ++i) {
-      uint32_t v;
-      std::memcpy(&v, &tmp[i], 4);
-      uint16_t val = static_cast<uint16_t>(v >> 16);
-      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) = val;
+      *reinterpret_cast<uint16_t*>(u8 + i * byte_stride) =
+          float_to_bf16(tmp[i]);
     }
   }
 };
@@ -734,10 +717,18 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
     return FP32Vec16(
         RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
   }
+  FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmax_vv_f32, LMUL_512)(reg, b.reg, elem_num));
+  }
   FP32Vec16 min(const FP32Vec16& b) const {
     return FP32Vec16(
         RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, VEC_ELEM_NUM));
   }
+  FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
+    return FP32Vec16(
+        RVVI(__riscv_vfmin_vv_f32, LMUL_512)(reg, b.reg, elem_num));
+  }
   FP32Vec16 abs() const {
     return FP32Vec16(RVVI(__riscv_vfabs_v_f32, LMUL_512)(reg, VEC_ELEM_NUM));
   }
@@ -867,6 +858,27 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
   }
 };
 
+struct INT8Vec16 : public Vec<INT8Vec16> {
+  constexpr static int VEC_ELEM_NUM = 16;
+  fixed_i8x16_t reg;
+
+  explicit INT8Vec16(const FP32Vec16& vec) {
+    auto i32_vec =
+        RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(vec.reg, VEC_ELEM_NUM);
+    auto i16_vec = RVVI(__riscv_vnclip_wx_i16, LMUL_256)(
+        i32_vec, 0, __RISCV_VXRM_RNU, VEC_ELEM_NUM);
+    reg = RVVI(__riscv_vnclip_wx_i8, LMUL_128)(i16_vec, 0, __RISCV_VXRM_RNU,
+                                               VEC_ELEM_NUM);
+  }
+
+  void save(int8_t* ptr) const {
+    RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, VEC_ELEM_NUM);
+  }
+  void save(int8_t* ptr, int elem_num) const {
+    RVVI(__riscv_vse8_v_i8, LMUL_128)(ptr, reg, elem_num);
+  }
+};
+
 // ============================================================================
 // Type Traits & Global Helpers
 // ============================================================================
@@ -956,9 +968,7 @@ inline BF16Vec16::BF16Vec16(const FP32Vec16& v)
 #else
 template <>
 inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
-  uint32_t val;
-  std::memcpy(&val, &v, 4);
-  *reinterpret_cast<uint16_t*>(ptr) = static_cast<uint16_t>(val >> 16);
+  *reinterpret_cast<uint16_t*>(ptr) = float_to_bf16(v);
 }
 inline BF16Vec8::BF16Vec8(const FP32Vec8& v) : reg_fp32(v.reg) {}
 inline BF16Vec16::BF16Vec16(const FP32Vec16& v) : reg_fp32(v.reg) {}

csrc/cpu/cpu_types_vxe.hpp

@@ -3,7 +3,9 @@
 #define CPU_TYPES_VXE_HPP
 
 #include <vecintrin.h>
+#include <bit>
 #include <cmath>
+#include <cstdint>
 #include <limits>
 #include <torch/all.h>
 namespace vec_op {
@@ -817,8 +819,7 @@ inline void storeFP32<::c10::Half>(float v, ::c10::Half* ptr) {
   // intrinsics for FP32 to FP16 conversion does not use IEEE rounding and can
   // produce incorrect results for some inputs. Process each of the 4 vectors
   // separately.
-  uint32_t in;
-  std::memcpy(&in, &v, sizeof(in));
+  uint32_t in = std::bit_cast<uint32_t>(v);
 
   uint32_t s = (in & 0x80000000) >> 16;  // Sign
   uint32_t e = (in & 0x7F800000) >> 23;  // Exponent

csrc/cpu/cpu_wna16.cpp

@@ -4,6 +4,9 @@
 #ifdef CPU_CAPABILITY_AMXBF16
   #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
 #endif
+#if defined(__riscv_v)
+  #include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
+#endif
 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
 
 #define VLLM_DISPATCH_CASE_16B_TYPES(...)                 \
@@ -319,6 +322,8 @@ void cpu_gemm_wna16(
       return ISA::AMX;
     } else if (isa_hint == "vec") {
       return ISA::VEC;
+    } else if (isa_hint == "rvv") {
+      return ISA::RVV;
     } else {
       TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
     }
@@ -397,6 +402,40 @@ void cpu_gemm_wna16(
             pack_factor);
         return;
       }
+    } else if (isa == ISA::RVV) {
+      using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
+      if (has_zp) {
+        using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
+        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
+            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
+            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
+            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
+            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+            scales_group_stride, zeros_group_stride, group_num, group_size,
+            pack_factor);
+        return;
+      }
+      if (use_desc_act) {
+        using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
+        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
+            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
+            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
+            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
+            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+            scales_group_stride, zeros_group_stride, group_num, group_size,
+            pack_factor);
+        return;
+      } else {
+        using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
+        cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
+            input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
+            output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
+            g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
+            a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
+            scales_group_stride, zeros_group_stride, group_num, group_size,
+            pack_factor);
+        return;
+      }
     }
   });
 }

csrc/cpu/float_convert.hpp

@@ -1,14 +1,15 @@
+#pragma once
 
-static float bf16_to_float(uint16_t bf16) {
+#include <bit>
+#include <cstdint>
+
+inline float bf16_to_float(uint16_t bf16) {
   uint32_t bits = static_cast<uint32_t>(bf16) << 16;
-  float fp32;
-  std::memcpy(&fp32, &bits, sizeof(fp32));
-  return fp32;
+  return std::bit_cast<float>(bits);
 }
 
-static uint16_t float_to_bf16(float fp32) {
-  uint32_t bits;
-  std::memcpy(&bits, &fp32, sizeof(fp32));
+inline uint16_t float_to_bf16(float fp32) {
+  uint32_t bits = std::bit_cast<uint32_t>(fp32);
   return static_cast<uint16_t>(bits >> 16);
 }
 
@@ -18,14 +19,13 @@ static uint16_t float_to_bf16(float fp32) {
  * Codes below copied from
  * https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
  *************************************************/
-static uint16_t float_to_fp16(float fp32) {
+inline uint16_t float_to_fp16(float fp32) {
   uint16_t fp16;
 
-  unsigned x;
   unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
   unsigned sign, exponent, mantissa;
 
-  std::memcpy(&x, &fp32, sizeof(fp32));
+  uint32_t x = std::bit_cast<uint32_t>(fp32);
   u = (x & 0x7fffffff);
 
   // Get rid of +NaN/-NaN case first.
@@ -77,12 +77,11 @@ static uint16_t float_to_fp16(float fp32) {
   return fp16;
 }
 
-static float fp16_to_float(uint16_t fp16) {
+inline float fp16_to_float(uint16_t fp16) {
   unsigned sign = ((fp16 >> 15) & 1);
   unsigned exponent = ((fp16 >> 10) & 0x1f);
   unsigned mantissa = ((fp16 & 0x3ff) << 13);
-  int temp;
-  float fp32;
+  uint32_t temp;
   if (exponent == 0x1f) { /* NaN or Inf */
     mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
     exponent = 0xff;
@@ -101,6 +100,5 @@ static float fp16_to_float(uint16_t fp16) {
     exponent += 0x70;
   }
   temp = ((sign << 31) | (exponent << 23) | mantissa);
-  std::memcpy(&fp32, &temp, sizeof(temp));
-  return fp32;
+  return std::bit_cast<float>(temp);
 }

csrc/cpu/generate_cpu_attn_dispatch.py

@@ -11,7 +11,7 @@ import os
 HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
 
 # Head dimensions divisible by 16 but not 32 (VEC16 only)
-HEAD_DIMS_16 = [80, 112]
+HEAD_DIMS_16 = [48, 80, 112]
 
 # ISA types
 ISA_TYPES = {

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/micro_gemm/cpu_micro_gemm_rvv.hpp

@@ -0,0 +1,228 @@
+#ifndef CPU_MICRO_GEMM_RVV_HPP
+#define CPU_MICRO_GEMM_RVV_HPP
+
+#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
+
+#if defined(__riscv_v)
+
+namespace cpu_micro_gemm {
+namespace {
+
+constexpr int32_t RVV_MGEMM_N8 = 8;
+constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
+
+template <typename scalar_t>
+FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
+
+template <>
+FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
+  return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
+}
+
+template <>
+FORCE_INLINE fixed_fp32x8_t
+load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
+  #if defined(__riscv_zvfh)
+  fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
+      reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
+  return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
+  #else
+  alignas(32) float values[RVV_MGEMM_N8];
+  for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
+    values[i] = static_cast<float>(ptr[i]);
+  }
+  return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
+  #endif
+}
+
+template <>
+FORCE_INLINE fixed_fp32x8_t
+load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
+  #if defined(__riscv_zvfbfmin)
+  fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
+      reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
+  fixed_bf16x8_t vec =
+      RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
+  return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
+  #else
+  fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
+      reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
+  auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
+  auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
+  return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
+  #endif
+}
+
+// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
+// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
+template <int32_t M, typename scalar_t>
+FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
+                                             const scalar_t* __restrict__ b_ptr,
+                                             float* __restrict__ c_ptr,
+                                             const int64_t lda,
+                                             const int64_t ldc, const int32_t k,
+                                             const bool accum_c) {
+  static_assert(0 < M && M <= 8);
+
+  #define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
+  #define RVV_IF_M(i) if constexpr (M > (i))
+
+  #define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
+  RVV_ROWS_APPLY(RVV_DECL_A)
+  #undef RVV_DECL_A
+
+  #define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
+  RVV_ROWS_APPLY(RVV_DECL_ACC)
+  #undef RVV_DECL_ACC
+
+  #define RVV_INIT_ACC(i)                                                  \
+    RVV_IF_M(i) {                                                          \
+      if (accum_c) {                                                       \
+        acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc,    \
+                                                     RVV_MGEMM_N8);        \
+      } else {                                                             \
+        acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
+      }                                                                    \
+    }
+  RVV_ROWS_APPLY(RVV_INIT_ACC)
+  #undef RVV_INIT_ACC
+
+  int32_t k_idx = 0;
+  for (; k_idx + 3 < k; k_idx += 4) {
+  #define RVV_FMA_ROW(i, K_OFFSET)                                     \
+    RVV_IF_M(i) {                                                      \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(                  \
+          acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
+          RVV_MGEMM_N8);                                               \
+    }
+
+  #define RVV_STEP_K(K_OFFSET)                                      \
+    {                                                               \
+      fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>(              \
+          b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
+      RVV_FMA_ROW(0, K_OFFSET)                                      \
+      RVV_FMA_ROW(1, K_OFFSET)                                      \
+      RVV_FMA_ROW(2, K_OFFSET)                                      \
+      RVV_FMA_ROW(3, K_OFFSET)                                      \
+      RVV_FMA_ROW(4, K_OFFSET)                                      \
+      RVV_FMA_ROW(5, K_OFFSET)                                      \
+      RVV_FMA_ROW(6, K_OFFSET)                                      \
+      RVV_FMA_ROW(7, K_OFFSET)                                      \
+    }
+
+    RVV_STEP_K(0)
+    RVV_STEP_K(1)
+    RVV_STEP_K(2)
+    RVV_STEP_K(3)
+  #undef RVV_STEP_K
+  #undef RVV_FMA_ROW
+  }
+
+  for (; k_idx < k; ++k_idx) {
+    fixed_fp32x8_t b =
+        load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
+  #define RVV_TAIL_ROW(i)                                                \
+    RVV_IF_M(i) {                                                        \
+      acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(                    \
+          acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
+    }
+    RVV_ROWS_APPLY(RVV_TAIL_ROW)
+  #undef RVV_TAIL_ROW
+  }
+
+  #define RVV_STORE_ROW(i)                                           \
+    RVV_IF_M(i) {                                                    \
+      RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
+                                          RVV_MGEMM_N8);             \
+    }
+  RVV_ROWS_APPLY(RVV_STORE_ROW)
+  #undef RVV_STORE_ROW
+
+  #undef RVV_ROWS_APPLY
+  #undef RVV_IF_M
+}
+
+template <int32_t M, typename scalar_t>
+FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+  static_assert(0 < M && M <= 8);
+  scalar_t* __restrict__ curr_b_0 = b_ptr;
+  scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
+
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
+                                c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
+                                accum_c);
+  gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
+                                ldc, k, accum_c);
+}
+
+class TileGemmRVV {
+ public:
+  template <typename scalar_t>
+  FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+    switch (m) {
+      case 1:
+        gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 2:
+        gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 3:
+        gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 4:
+        gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 5:
+        gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 6:
+        gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 7:
+        gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
+        break;
+      case 8:
+        gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
+        break;
+    }
+  }
+};
+
+}  // namespace
+
+template <typename scalar_t>
+class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
+ public:
+  static constexpr int32_t MaxMSize = 8;
+  static constexpr int32_t NSize = 32;
+
+ public:
+  void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
+    TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
+  }
+
+  static void pack_weight(const scalar_t* __restrict__ weight,
+                          scalar_t* __restrict__ packed_weight,
+                          const int32_t output_size, const int32_t input_size) {
+    TORCH_CHECK_EQ(output_size % 16, 0);
+    for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
+      const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
+      scalar_t* __restrict__ curr_packed_weight =
+          packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
+      for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
+        *curr_packed_weight = *curr_weight;
+
+        curr_packed_weight += 16;
+        ++curr_weight;
+      }
+    }
+  }
+};
+
+}  // namespace cpu_micro_gemm
+
+#endif  // defined(__riscv_v)
+
+#endif  // CPU_MICRO_GEMM_RVV_HPP

csrc/cpu/sgl-kernels/gemm_int4.cpp

@@ -268,6 +268,23 @@ void _dequant_gemm_accum_small_M(
   _dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
 #endif
 
+template <int64_t N, int64_t ldb>
+inline int32_t load_uint4_vnni(const uint8_t* __restrict__ B, int64_t k, int64_t n) {
+  // B is packed as [_block_k / 4, N / 2, 4] for VNNI4. Each byte stores two
+  // columns from adjacent 8-column groups for one K lane.
+  constexpr int64_t n_group_size = 8;
+  constexpr int64_t vnni_size = 4;
+  static_assert(N % (2 * n_group_size) == 0);
+
+  int64_t n_group = n / n_group_size;
+  int64_t ni = n % n_group_size;
+  int64_t ki = k % vnni_size;
+  int64_t k_base = k - ki;
+  int64_t packed_n = (n_group / 2) * n_group_size + ni;
+  uint8_t packed = B[k_base * ldb + packed_n * vnni_size + ki];
+  return (n_group % 2 == 0) ? (packed & 0x0f) : ((packed >> 4) & 0x0f);
+}
+
 template <int64_t N, int64_t ldb, bool sym_quant_act>
 void _dequant_gemm_accum(
     float* C,
@@ -321,7 +338,24 @@ void _dequant_gemm_accum(
   } else
 #endif
   {
-    TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!");
+    for (int64_t m = 0; m < M; ++m) {
+      for (int64_t n = 0; n < N; ++n) {
+        int32_t acc = 0;
+        for (int64_t k = 0; k < K; ++k) {
+          int32_t b = load_uint4_vnni<N, ldb>(B, k, n) - qzeros_b[n];
+          if constexpr (sym_quant_act) {
+            const int8_t* A_s8 = reinterpret_cast<const int8_t*>(A);
+            acc += static_cast<int32_t>(A_s8[m * lda + k]) * b;
+          } else {
+            acc += static_cast<int32_t>(A[m * lda + k]) * b;
+          }
+        }
+        if constexpr (!sym_quant_act) {
+          acc -= qzeros_a[m] * compensation[n];
+        }
+        C[m * ldc + n] += static_cast<float>(acc) * scales_a[m] * scales_b[n];
+      }
+    }
   }
 }
 
@@ -496,9 +530,11 @@ void _da8w4_linear_impl(
         store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
       }
     }
+#if defined(CPU_CAPABILITY_AVX512)
     if (use_brgemm) {
       at::native::cpublas::brgemm_release();
     }
+#endif
   });
 }
 

csrc/cpu/sgl-kernels/vec.h

@@ -245,7 +245,7 @@ quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restric
 
   for (int64_t k = 0; k < K; ++k) {
     const float val = static_cast<float>(A[k]) * inv_scale;
-    Aq[k] = (uint8_t)(std::round(val)) + 128;
+    Aq[k] = static_cast<uint8_t>(static_cast<int32_t>(std::round(val)) + 128);
   }
   As = scale;
 }

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);
 
@@ -329,8 +332,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
 
   // Quantization
-#if defined(__AVX512F__) || defined(__AVX2__) || \
-    (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
+#if defined(__AVX512F__) || defined(__AVX2__) ||                               \
+    (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
+    defined(__riscv_v)
   // Helper function to release oneDNN handlers
   ops.def("release_dnnl_matmul_handler(int handler) -> ()",
           &release_dnnl_matmul_handler);
@@ -428,19 +432,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
            &int8_scaled_mm_with_quant);
 
-  // Adapted from sglang: INT4 W4A8 kernels
-  ops.def(
-      "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
-      "scales, int quant_method_4bit) -> (Tensor, "
-      "Tensor, Tensor)");
-  ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
-           &convert_weight_packed_scale_zp);
-
-  ops.def(
-      "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
-      "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
-  ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
-
   // Adapted from sglang: FP8 W8A16 kernel
   ops.def(
       "fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
@@ -467,6 +458,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
 #endif
 
+#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
+     defined(__AVX512VNNI__)) ||                        \
+    defined(__riscv)
+  // Adapted from sglang: INT4 W4A8 kernels
+  ops.def(
+      "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
+      "scales, int quant_method_4bit) -> (Tensor, "
+      "Tensor, Tensor)");
+  ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
+           &convert_weight_packed_scale_zp);
+
+  ops.def(
+      "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
+      "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
+  ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
+#endif
+
   // Adapted from sglang: GDN kernels
   ops.def(
       "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
@@ -491,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!) "
@@ -507,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/cpu/utils.hpp

@@ -8,13 +8,15 @@
 #include "cpu/cpu_types.hpp"
 
 namespace cpu_utils {
-enum class ISA { AMX, VEC };
+enum class ISA { AMX, VEC, RVV };
 
 inline ISA get_isa(const std::string& isa) {
   if (isa == "amx") {
     return ISA::AMX;
   } else if (isa == "vec") {
     return ISA::VEC;
+  } else if (isa == "rvv") {
+    return ISA::RVV;
   } else {
     TORCH_CHECK(false, "Invalid isa type: " + isa);
   }

csrc/cutlass_extensions/vllm_cutlass_library_extension.py

@@ -57,13 +57,13 @@ VLLMDataTypeVLLMScalarTypeTag: dict[VLLMDataType | DataType, str] = {
 }
 
 VLLMDataTypeTorchDataTypeTag: dict[VLLMDataType | DataType, str] = {
-    DataType.u8: "at::ScalarType::Byte",
-    DataType.s8: "at::ScalarType::Char",
-    DataType.e4m3: "at::ScalarType::Float8_e4m3fn",
-    DataType.s32: "at::ScalarType::Int",
-    DataType.f16: "at::ScalarType::Half",
-    DataType.bf16: "at::ScalarType::BFloat16",
-    DataType.f32: "at::ScalarType::Float",
+    DataType.u8: "torch::headeronly::ScalarType::Byte",
+    DataType.s8: "torch::headeronly::ScalarType::Char",
+    DataType.e4m3: "torch::headeronly::ScalarType::Float8_e4m3fn",
+    DataType.s32: "torch::headeronly::ScalarType::Int",
+    DataType.f16: "torch::headeronly::ScalarType::Half",
+    DataType.bf16: "torch::headeronly::ScalarType::BFloat16",
+    DataType.f32: "torch::headeronly::ScalarType::Float",
 }
 
 VLLMKernelScheduleTag: dict[MixedInputKernelScheduleType | KernelScheduleType, str] = {

csrc/libtorch_stable/activation_kernels.cu

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

csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu

@@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options(
   StrideQ stride_Q_pe = cute::make_tuple(
       static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));
 
+  // Read the token and page strides from the cache tensor instead of assuming
+  // packed pages, so strided views (e.g. per-layer views into a cross-layer
+  // block-major cache) are addressed correctly.
   StrideK stride_C = cute::make_tuple(
-      static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
+      static_cast<int64_t>(kv_c_and_k_pe_cache.stride(1)), _1{},
+      static_cast<int64_t>(kv_c_and_k_pe_cache.stride(0)));
   StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
   StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
   StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
@@ -268,9 +272,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/cache_kernels.cu

@@ -549,7 +549,7 @@ __global__ void indexer_k_quant_and_cache_kernel(
     const int head_dim,                        // dimension of each head
     const int quant_block_size,                // quantization block size
     const int cache_block_size,                // cache block size
-    const int cache_stride,  // stride for each token in kv_cache
+    const int64_t cache_block_stride,  // stride for each block in kv_cache
 
     const bool use_ue8m0  // use ue8m0 scale format
 ) {
@@ -590,16 +590,15 @@ __global__ void indexer_k_quant_and_cache_kernel(
     scale = exp2f(ceilf(log2f(scale)));
   }
 
-  const int64_t dst_offset = block_idx * cache_block_size * cache_stride +
-                             block_offset * head_dim + head_dim_idx;
+  const int64_t dst_offset =
+      block_idx * cache_block_stride + block_offset * head_dim + head_dim_idx;
   for (int i = 0; i < VEC_SIZE; i++) {
     kv_cache[dst_offset + i] =
         fp8::scaled_convert<cache_t, scalar_t, kv_dt>(k_val_ptr[i], scale);
   }
   if (threadIdx.x == 0) {
     const int64_t dst_scale_idx =
-        block_idx * cache_block_size * cache_stride +
-        cache_block_size * head_dim +
+        block_idx * cache_block_stride + cache_block_size * head_dim +
         (block_offset * head_dim + head_dim_idx) * 4 / quant_block_size;
     reinterpret_cast<float*>(kv_cache)[dst_scale_idx / 4] = scale;
   }
@@ -1452,7 +1451,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
           reinterpret_cast<KV_T*>(k.data_ptr()),                              \
           reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),                    \
           slot_mapping.const_data_ptr<int64_t>(), head_dim, quant_block_size, \
-          cache_block_size, cache_stride, use_ue8m0);
+          cache_block_size, cache_block_stride, use_ue8m0);
 
 void indexer_k_quant_and_cache(
     torch::stable::Tensor& k,         // [num_tokens, head_dim]
@@ -1463,7 +1462,7 @@ void indexer_k_quant_and_cache(
   int num_tokens = k.size(0);
   int head_dim = k.size(1);
   int cache_block_size = kv_cache.size(1);
-  int cache_stride = kv_cache.size(2);
+  int64_t cache_block_stride = kv_cache.stride(0);
   bool use_ue8m0 = scale_fmt == "ue8m0";
 
   STD_TORCH_CHECK(k.device() == kv_cache.device(),

csrc/libtorch_stable/cooperative_topk.cu

@@ -0,0 +1,146 @@
+// Cooperative cluster TopK for DeepSeek V3 sparse attention indexer.
+// See cooperative_topk.cuh for kernel implementation.
+
+#include <cuda_runtime.h>
+
+#include "torch_utils.h"
+
+#ifndef USE_ROCM
+  #include "cooperative_topk.cuh"
+namespace ct = vllm::cooperative;
+namespace hist4096 = vllm::topk_histogram_4096;
+#endif
+
+#ifndef USE_ROCM
+template <uint32_t TopK, uint32_t CS>
+void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
+                                size_t smem, cudaStream_t stream) {
+  auto kernel = []() {
+    if constexpr (CS == 16) {
+      return &ct::cooperative_topk_cs16<TopK>;
+    } else if constexpr (CS == 8) {
+      return &ct::cooperative_topk_cs8<TopK>;
+    } else {
+      static_assert(CS == 4, "unsupported cooperative_topk cluster size");
+      return &ct::cooperative_topk_cs4<TopK>;
+    }
+  }();
+  if constexpr (CS > 8) {
+    cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed,
+                         1);
+  }
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                       smem);
+
+  cudaLaunchConfig_t cfg = {};
+  cfg.gridDim = dim3(params.num_rows, CS);
+  cfg.blockDim = dim3(hist4096::kBlockSize);
+  cfg.dynamicSmemBytes = smem;
+  cfg.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeClusterDimension;
+  attrs[0].val.clusterDim = {1, CS, 1};
+  cfg.numAttrs = 1;
+  cfg.attrs = attrs;
+  cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params);
+  STD_TORCH_CHECK(err == cudaSuccess,
+                  "cooperative_topk launch failed: ", cudaGetErrorString(err));
+}
+
+template <uint32_t TopK>
+void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
+                                  const torch::stable::Tensor& lengths,
+                                  torch::stable::Tensor& output,
+                                  torch::stable::Tensor& workspace,
+                                  int64_t max_seq_len) {
+  (void)max_seq_len;  // Kept for signature parity with persistent_topk.
+  const int64_t num_rows = logits.size(0);
+  const cudaStream_t stream = get_current_cuda_stream();
+
+  const uint32_t stride = static_cast<uint32_t>(logits.stride(0));
+  // 32 = max clusters for CS=4 (32 x 4 = 128 CTAs = 66% of SMs, leaves
+  // headroom)
+  STD_TORCH_CHECK(
+      num_rows <= 32,
+      "cooperative_topk supports <=32 rows; use persistent_topk for "
+      "larger batches");
+
+  STD_TORCH_CHECK(stride % 4 == 0,
+                  "cooperative_topk: stride must be multiple of 4 for TMA "
+                  "alignment, got stride (max_model_len)=",
+                  stride);
+
+  STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
+  STD_TORCH_CHECK(
+      workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
+      "workspace must be uint8");
+
+  ct::CooperativeTopKParams<TopK> params;
+  params.input = logits.const_data_ptr<float>();
+  params.output = output.mutable_data_ptr<int32_t>();
+  params.lengths = lengths.const_data_ptr<int32_t>();
+  params.num_rows = static_cast<uint32_t>(num_rows);
+  params.stride = stride;
+  params.tie_ws =
+      reinterpret_cast<hist4096::Tie*>(workspace.mutable_data_ptr<uint8_t>());
+
+  constexpr uint32_t kTieWsPerRow =
+      TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
+  STD_TORCH_CHECK(
+      workspace.size(0) >=
+          static_cast<int64_t>(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)),
+      "workspace too small");
+
+  const bool supports_cluster16 = get_device_prop()->major >= 10;
+  if (num_rows <= 4 && supports_cluster16) {
+    launch_cooperative_cluster<TopK, 16>(params, ct::kSmemSize8, stream);
+  } else if (num_rows <= 8) {
+    launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
+  } else {
+    launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
+  }
+}
+#endif  // USE_ROCM
+
+void cooperative_topk(const torch::stable::Tensor& logits,
+                      const torch::stable::Tensor& lengths,
+                      torch::stable::Tensor& output,
+                      torch::stable::Tensor& workspace, int64_t k,
+                      int64_t max_seq_len) {
+#ifndef USE_ROCM
+  STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
+  STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
+  STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
+  STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
+                  "Only float32 supported");
+  STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "lengths must be int32");
+  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
+                  "output must be int32");
+  STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
+  STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
+                  "lengths must be 1D or 2D");
+  STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
+  STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
+  const int64_t num_rows = logits.size(0);
+  STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
+  STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
+                  "output size mismatch");
+  STD_TORCH_CHECK(
+      k == 512 || k == 1024 || k == 2048,
+      "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k);
+
+  if (k == 512) {
+    launch_cooperative_topk_impl<512>(logits, lengths, output, workspace,
+                                      max_seq_len);
+  } else if (k == 1024) {
+    launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace,
+                                       max_seq_len);
+  } else {
+    launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace,
+                                       max_seq_len);
+  }
+#else
+  STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm");
+#endif
+}

csrc/libtorch_stable/cooperative_topk.cuh

@@ -0,0 +1,593 @@
+/*
+ * Cooperative TopK kernel for DSA Indexer
+ */
+
+#ifndef COOPERATIVE_TOPK_CUH_
+#define COOPERATIVE_TOPK_CUH_
+
+#include <cooperative_groups.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cuda/ptx>
+#include <algorithm>
+#include <cstdint>
+
+#include "topk_histogram_4096.cuh"
+
+namespace vllm {
+namespace cooperative {
+
+namespace hist4096 = topk_histogram_4096;
+
+constexpr uint32_t kHistBits = 10;
+constexpr uint32_t kHistBins = 1 << kHistBits;
+constexpr uint32_t kMaxTopK = 2048;
+
+constexpr uint32_t kElemPerStage = 16;
+constexpr uint32_t kSizePerStage =
+    kElemPerStage * hist4096::kBlockSize;  // 16384
+
+// CS=4 two-pass path uses two TMA stages as a double buffer.
+constexpr uint32_t kStreamingStagesCS4 = 2;
+// CS=8/16 fused paths keep all loaded TMA stages resident in smem.
+constexpr uint32_t kFusedStagesCS8 = 2;
+constexpr uint32_t kFusedStagesCS16 = 2;
+
+// CS=4 single-pass path
+constexpr uint32_t kMaxSinglePassStages = 3;
+constexpr uint32_t kMaxSinglePassPerBlock =
+    kMaxSinglePassStages * kSizePerStage;  // 49152
+
+template <uint32_t TopK = 1024>
+struct CooperativeTopKParams {
+  const float* __restrict__ input;
+  int32_t* __restrict__ output;
+  const int32_t* __restrict__ lengths;
+  hist4096::Tie* __restrict__ tie_ws;  // per-row tie workspace, see
+                                       // kTieWsPerRow
+  uint32_t num_rows, stride;
+};
+
+// ============================================================================
+// Cooperative helpers
+// ============================================================================
+
+// only CS adjacent lanes participate (sub-warp reduce), in opposite to
+// warp_reduce_sum_full
+template <uint32_t N>
+__device__ __forceinline__ uint32_t warp_reduce_sum_subN(uint32_t v) {
+#pragma unroll
+  for (uint32_t m = N >> 1; m > 0; m >>= 1)
+    v += __shfl_xor_sync(0xFFFFFFFF, v, m, 32);
+  return v;
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+__device__ __forceinline__ uint32_t extract_coarse_bin(float x) {
+  return hist4096::extract_coarse_bin_N<kHistBits>(x);
+}
+
+__device__ __forceinline__ void mbarrier_init(uint64_t* a, uint32_t n) {
+  cuda::ptx::mbarrier_init(a, n);
+}
+__device__ __forceinline__ void mbarrier_wait(uint64_t* a, uint32_t p) {
+  while (!cuda::ptx::mbarrier_try_wait_parity(cuda::ptx::sem_relaxed,
+                                              cuda::ptx::scope_cta, a, p));
+}
+__device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t* a,
+                                                          uint32_t t) {
+  cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_relaxed,
+                                       cuda::ptx::scope_cta,
+                                       cuda::ptx::space_shared, a, t);
+}
+__device__ __forceinline__ void tma_load(void* d, const void* s, uint32_t n,
+                                         uint64_t* m) {
+  cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, d,
+                           s, n, m);
+}
+
+// ============================================================================
+// DSMEM histogram reduce
+// ============================================================================
+
+template <uint32_t CS>
+__device__ __forceinline__ void dsmem_hist_reduce(uint32_t* histogram) {
+  static_assert(kHistBins <= hist4096::kBlockSize);
+  auto cluster = cooperative_groups::this_cluster();
+  cluster.sync();
+  const auto tx = threadIdx.x;
+  const auto rank = blockIdx.y;
+  constexpr auto kLocal = kHistBins / CS;
+  const auto off = kLocal * rank;
+  if (tx < kHistBins) {
+    const auto addr = &histogram[off + tx / CS];
+    const auto src = cluster.map_shared_rank(addr, tx % CS);
+    *src = warp_reduce_sum_subN<CS>(*src);
+  }
+  cluster.sync();
+}
+
+// ============================================================================
+// Find threshold from reduced histogram
+// ============================================================================
+
+// NOTE: caller must ensure a cluster.sync() or __syncthreads() happened
+// before calling this, so warp_sum writes are visible across warps.
+// The first internal __syncthreads() is still needed for the warp_sum exchange.
+template <uint32_t TopK>
+__device__ __forceinline__ void find_threshold(uint32_t* histogram,
+                                               uint32_t* warp_sum,
+                                               uint32_t* counter_gt,
+                                               uint32_t* counter_eq,
+                                               hist4096::MatchBin* match) {
+  const auto tx = threadIdx.x;
+  const auto li = tx % hist4096::kWarpSize, wi = tx / hist4096::kWarpSize;
+  const auto value = tx < kHistBins ? histogram[tx] : 0;
+  const auto winc = hist4096::warp_inclusive_sum(li, value);
+  if (li == hist4096::kWarpSize - 1) warp_sum[wi] = winc;
+  __syncthreads();
+  const auto tmp = warp_sum[li];
+  const auto total = hist4096::warp_reduce_sum_full(tmp);
+  auto pfx = hist4096::warp_reduce_sum_full(li < wi ? tmp : 0) + winc;
+  const auto above = total - pfx;
+  if (tx < kHistBins && above < TopK && above + value >= TopK) {
+    *counter_gt = *counter_eq = 0;
+    *match = {.bin = tx, .above_count = above, .equal_count = value};
+  }
+  __syncthreads();
+}
+
+// Streams data through shared memory in chunks, processing each chunk before
+// loading the next overwrites each buffer after processing it (the epilogue
+// prefetch loads the next chunk into the same slot)
+template <typename SmemType, uint32_t kStages, uint32_t kBinBits,
+          bool kIsScatter>
+__device__ void tma_stream_pass(const float* scores, uint32_t length,
+                                uint32_t thr_bin, int32_t* indices,
+                                uint32_t* phases, SmemType* smem) {
+  const auto tx = threadIdx.x;
+  const auto lane = tx % hist4096::kWarpSize;
+  const auto ni =
+      (length + kSizePerStage - 1) / kSizePerStage;  // total stages needed
+  const auto la =
+      (length + 3u) & ~3u;  // length rounded up to float4 (TMA alignment)
+  const auto pass =
+      kIsScatter ? 1 : 0;  // barrier dim: [0] for histogram, [1] for scatter
+
+  // Prologue: issue initial TMA loads - prefill the pipeline
+  if (tx == 0) {
+#pragma unroll
+    for (uint32_t i = 0; i < kStages; i++) {
+      if (i >= ni) {
+        break;
+      }
+      const auto o = i * kSizePerStage;
+      const auto sz = min(kSizePerStage, la - o) * sizeof(float);
+      tma_load(smem->score_buffer[i], scores + o, sz,
+               &smem->barrier[pass][i]);  // cp.async.bulk is non-blocking
+      mbarrier_arrive_expect_tx(&smem->barrier[pass][i], sz);
+    }
+  }
+
+  // Main loop: process stages
+  for (uint32_t it = 0; it < ni; it++) {
+    const auto b = it % kStages;  // which buffer slot (0 or 1)
+    const auto o = it * kSizePerStage;
+    const auto sz = min(kSizePerStage, length - o);
+
+    if (lane == 0) {
+      mbarrier_wait(&smem->barrier[pass][b],
+                    phases[b] & 1);  // wait for the data
+    }
+    phases[b]++;  // advances the phase for next time this slot is reused
+    __syncwarp();
+
+#pragma unroll
+    for (uint32_t i = 0; i < kElemPerStage; i++) {
+      const auto li = tx + i * hist4096::kBlockSize;
+      if (li >= sz) {
+        break;
+      }
+      const auto sc = smem->score_buffer[b][li];
+      const auto bn = hist4096::extract_coarse_bin_N<kBinBits>(sc);
+      if constexpr (kIsScatter) {  // compile-time branch
+        // Scatter pass: place above-threshold and collect ties
+        const auto gi = o + li;
+        if (bn > thr_bin) {
+          indices[atomicAdd(&smem->counter_gt, 1)] = gi;
+        } else if (bn == thr_bin) {
+          const auto p = atomicAdd(&smem->counter_eq, 1);
+          if (p < hist4096::kMaxTies) {
+            smem->tie_buffer[p] = {gi, sc};
+          }
+        }
+      } else {
+        // Histogram pass: just count
+        atomicAdd(&smem->histogram[bn], 1);
+      }
+    }
+    __syncthreads();  // ensures all threads finished processing their buffer
+                      // before next TMA load
+
+    // Epilogue: issue next TMA load
+    if (tx == 0 && it + kStages < ni) {
+      const auto no = (it + kStages) * kSizePerStage;
+      const auto nsz = min(kSizePerStage, la - no) * sizeof(float);
+      tma_load(smem->score_buffer[b], scores + no, nsz,
+               &smem->barrier[pass][b]);
+      mbarrier_arrive_expect_tx(&smem->barrier[pass][b], nsz);
+    }
+  }
+}
+
+// ============================================================================
+// Fused path: single TMA pass, rescan smem for scatter
+// ============================================================================
+
+// Fused shared memory layout for cluster cooperative paths.
+// kPasses=1 for single-pass (CS=8, CS=4 singlepass), kPasses=2 for two-pass
+// (CS=4).
+template <uint32_t kStages, uint32_t kPasses = 1>
+struct SmemFused {
+  uint64_t barrier[kPasses][kStages];
+  alignas(128) uint32_t counter_gt;
+  alignas(128) uint32_t counter_eq;
+  alignas(128) hist4096::MatchBin match;
+  uint32_t warp_sum[hist4096::kNumWarps];
+  union {
+    uint32_t histogram[kHistBins];
+    hist4096::Tie tie_buffer[kMaxTopK];
+  };
+  alignas(128) float score_buffer[kStages][kSizePerStage];
+};
+
+using Smem8 = SmemFused<kFusedStagesCS8>;
+using Smem16 = SmemFused<kFusedStagesCS16>;
+using Smem4 = SmemFused<kStreamingStagesCS4, 2>;
+using SmemSinglePass = SmemFused<kMaxSinglePassStages>;
+
+// Cluster-cooperative large path.
+// kFused=true: all TMA stages resident, single-pass histogram + scatter (rescan
+// from smem). kFused=false: TMA double-buffer streaming, two passes (histogram
+// then scatter).
+template <uint32_t TopK, uint32_t CS, typename SmemType, bool kFused>
+__device__ void large_topk(const float* __restrict__ row_input,
+                           int32_t* __restrict__ row_output, uint32_t seq_len,
+                           uint32_t* phases, hist4096::Tie* tie_ws) {
+  const auto rank = blockIdx.y;  // this block's position in cluster
+  const auto tx = threadIdx.x;
+  const auto lane = tx % hist4096::kWarpSize;
+
+  extern __shared__ uint8_t smem_raw[];
+  auto* smem = reinterpret_cast<SmemType*>(smem_raw);
+  int32_t* s_topk = reinterpret_cast<int32_t*>(smem_raw + sizeof(SmemType));
+
+  // Partition row across cluster ranks
+  constexpr uint32_t kAlign = 4;
+  const auto units =
+      (seq_len + kAlign - 1) / kAlign;  // float4-aligned element count
+  const auto base = units / CS, extra = units % CS;  // elements per block
+  const auto lu = base + (rank < extra ? 1u : 0u);   // remainder blocks
+  const auto ou =
+      rank * base + min(rank, extra);  // this block's count (load-balanced)
+  const auto my_start = ou * kAlign;   // global start offset
+  const auto my_len = min(my_start + lu * kAlign, seq_len) -
+                      my_start;  // actual length of this block
+  const auto num_iters =
+      (my_len + kSizePerStage - 1) / kSizePerStage;  // TMA stages needed
+  const auto len_aligned = (my_len + 3u) & ~3u;
+
+  if constexpr (kFused) {
+    // Fused init + TMA prologue
+    if (tx < kHistBins) {
+      smem->histogram[tx] = 0;  // all threads zero histogram
+    }
+    if (tx == 0) {  // thread 0 issues TMA - then all threads continue working
+                    // until mbarrier sync
+      smem->counter_gt = 0;
+      smem->counter_eq = 0;
+      for (uint32_t i = 0; i < num_iters; i++) {
+        const auto off = i * kSizePerStage;
+        const auto sz = min(kSizePerStage, len_aligned - off) * sizeof(float);
+        tma_load(smem->score_buffer[i], row_input + my_start + off, sz,
+                 &smem->barrier[0][i]);  // cp.async.bulk of size kSizePerStage
+                                         // × sizeof(float)
+        mbarrier_arrive_expect_tx(&smem->barrier[0][i], sz);
+      }
+    }
+    __syncthreads();
+
+    // Histogram build. ILP unroll-by-2, no inter-stage sync
+    for (uint32_t iter = 0; iter < num_iters; iter++) {
+      const auto off = iter * kSizePerStage;
+      const auto sz = min(kSizePerStage, my_len - off);
+      if (lane == 0) {
+        mbarrier_wait(&smem->barrier[0][iter],
+                      phases[iter] & 1);  // wait for TMA
+      }
+      phases[iter]++;
+      __syncwarp();
+#pragma unroll
+      for (uint32_t i = 0; i < kElemPerStage; i += 2) {
+        const auto li0 = tx + i * hist4096::kBlockSize;
+        const auto li1 = tx + (i + 1) * hist4096::kBlockSize;
+        if (li0 >= sz) {
+          break;
+        }
+        const auto b0 = extract_coarse_bin(smem->score_buffer[iter][li0]);
+        if (li1 < sz) {
+          const auto b1 = extract_coarse_bin(smem->score_buffer[iter][li1]);
+          atomicAdd(&smem->histogram[b0], 1);
+          atomicAdd(&smem->histogram[b1], 1);
+        } else {
+          atomicAdd(&smem->histogram[b0], 1);
+        }
+      }
+    }
+  } else {
+    // Twopass: init then stream histogram pass
+    if (tx < kHistBins) {
+      smem->histogram[tx] = 0;
+    }
+    if (tx == 0) {
+      smem->counter_gt = 0;
+      smem->counter_eq = 0;
+    }
+    __syncthreads();
+    tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, false>(
+        row_input + my_start, my_len, 0, nullptr, phases, smem);
+  }
+
+  // DSMEM all-reduce + find threshold
+  dsmem_hist_reduce<CS>(
+      smem->histogram);  // each block histogram is summed across all CS blocks
+  find_threshold<TopK>(smem->histogram, smem->warp_sum, &smem->counter_gt,
+                       &smem->counter_eq, &smem->match);
+
+  const auto thr = smem->match.bin;
+
+  if constexpr (kFused) {
+    // Fused scatter: rescan score_buffer (still in smem)
+    for (uint32_t iter = 0; iter < num_iters; iter++) {
+      const auto off = iter * kSizePerStage;
+      const auto sz = min(kSizePerStage, my_len - off);
+#pragma unroll
+      for (uint32_t i = 0; i < kElemPerStage; i++) {
+        const auto li = tx + i * hist4096::kBlockSize;
+        if (li >= sz) {
+          break;
+        }
+        const auto score = smem->score_buffer[iter][li];  // still in smem
+        const auto bin = extract_coarse_bin(score);
+        const auto gidx = off + li;
+        if (bin > thr) {
+          s_topk[atomicAdd(&smem->counter_gt, 1)] = gidx;  // above -> s_topk
+        } else if (bin == thr) {
+          const auto p = atomicAdd(&smem->counter_eq,
+                                   1);  // equal -> ties (later refinement)
+          if (p < hist4096::kMaxTies) {
+            smem->tie_buffer[p] = {gidx, score};
+          }
+        }
+      }
+    }
+    __syncthreads();
+  } else {
+    // Twopass scatter: re-stream data via TMA
+    uint32_t scatter_phases[kStreamingStagesCS4] = {0, 0};
+    tma_stream_pass<SmemType, kStreamingStagesCS4, kHistBits, true>(
+        row_input + my_start, my_len, thr, s_topk, scatter_phases, smem);
+  }
+
+  // Output collection via DSMEM prefix sum
+  constexpr uint32_t kAboveBits = 16;
+  constexpr uint32_t kAboveMask = (1 << kAboveBits) - 1;
+  static_assert(kAboveMask >= TopK);
+  static_assert(kAboveMask >= kMaxSinglePassPerBlock,
+                "kAboveBits must cover max per-block element count");
+
+  const uint32_t la = smem->counter_gt;
+  const uint32_t le_full = smem->counter_eq;
+  const uint32_t le =
+      min(le_full, hist4096::kMaxTies);  // written smem tie_buffer entries
+
+  __shared__ uint32_t s_local_counts[CS];
+  __shared__ uint32_t s_prefix_packed;
+  __shared__ uint32_t s_total_above, s_total_equal;
+
+  auto cluster = cooperative_groups::this_cluster();
+  if (tx < CS) {
+    // Pack written tie counts into 32-bit: (equal << 16) | above.
+    // `le_full` may exceed the per-block tie buffer cap; using it here creates
+    // holes in tie_ws and can make TopK=2048 refine unwritten workspace slots.
+    const uint32_t packed = (le << kAboveBits) | la;
+    const auto dst = cluster.map_shared_rank(s_local_counts, tx);
+    dst[rank] = packed;  // write my count to every block's s_local_counts[rank]
+  }
+  cluster.sync();
+
+  // Thread 0 computes serial prefix sum
+  if (tx == 0) {
+    uint32_t prefix = 0, ta = 0, te = 0;
+    for (uint32_t i = 0; i < CS; i++) {
+      if (i == rank) {
+        s_prefix_packed = prefix;  // my prefix
+      }
+      ta += s_local_counts[i] & kAboveMask;   // total above
+      te += s_local_counts[i] >> kAboveBits;  // total equal
+      prefix += s_local_counts[i];
+    }
+    s_total_above = ta;
+    s_total_equal = te;
+  }
+  __syncthreads();
+
+  const uint32_t prefix_above = s_prefix_packed & kAboveMask;
+  const uint32_t prefix_equal = s_prefix_packed >> kAboveBits;
+
+  // Write to global output
+  for (uint32_t i = tx; i < la; i += hist4096::kBlockSize) {
+    // indices are placed contiguously starting at prefix_above
+    row_output[prefix_above + i] =
+        s_topk[i] + my_start;  // my_start: block-local -> row-global index
+  }
+  for (uint32_t i = tx; i < le; i += hist4096::kBlockSize) {
+    const auto t = smem->tie_buffer[i];
+    uint32_t p = s_total_above + prefix_equal + i;
+    if (p < TopK) {
+      row_output[p] = t.idx + my_start;
+    }
+    uint32_t tp = prefix_equal + i;
+    if (tp < (TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK)) {
+      tie_ws[tp] = hist4096::Tie{t.idx + my_start, t.score};
+    }
+  }
+
+  // Tie refinement
+  cooperative_groups::this_cluster().sync();
+  if (rank != 0) {  // only rank 0 does tie refinement
+    return;
+  }
+  if (s_total_above + s_total_equal <= TopK) {  // no ties to refine
+    return;
+  }
+
+  // Tie-breaking uses FP32 (4-round radix sort)
+  if constexpr (TopK <= hist4096::kBlockSize) {
+    // copy ties from tie_ws back to smem, then refine
+    const uint32_t num_ties = min(s_total_equal, hist4096::kMaxTies);
+    // TODO (roberto): could vectorize with uint2 (8 bytes = exactly one Tie)
+    for (uint32_t i = tx; i < num_ties; i += hist4096::kBlockSize) {
+      smem->tie_buffer[i] = hist4096::Tie{tie_ws[i].idx, tie_ws[i].score};
+    }
+    __syncthreads();
+    hist4096::tie_handle<TopK>(smem->tie_buffer, num_ties, s_total_above,
+                               row_output, smem);
+  } else {
+    // TopK=2048: process directly from tie_ws (GMEM)
+    const uint32_t num_ties = min(s_total_equal, static_cast<uint32_t>(TopK));
+    hist4096::tie_handle_large<TopK>(tie_ws, num_ties, s_total_above,
+                                     row_output, smem);
+  }
+}
+
+// ============================================================================
+// Adapted from https://github.com/sgl-project/sglang/pull/23600
+// sgl-project/sglang
+// (python/sglang/jit_kernel/include/sgl_kernel/deepseek_v4/topk/)
+// ============================================================================
+
+template <uint32_t TopK, uint32_t CS>
+__device__ void cooperative_topk_body(CooperativeTopKParams<TopK> params) {
+  const auto rank = blockIdx.y, row = blockIdx.x, tx = threadIdx.x;
+  const auto sl = params.lengths[row];
+  int32_t* out = params.output + row * TopK;
+  const float* in = params.input + row * params.stride;
+
+  // Trivial: seq_len <= TopK
+  if (sl <= static_cast<int32_t>(TopK)) {
+    if (rank == 0) {
+      for (uint32_t i = tx; i < TopK; i += hist4096::kBlockSize) {
+        out[i] = (i < static_cast<uint32_t>(sl)) ? static_cast<int32_t>(i) : -1;
+      }
+    }
+    return;
+  }
+
+  // Short-Medium path: histogram_4096_topk on rank 0 only - all data fits in RF
+  if (sl <= static_cast<int32_t>(hist4096::kHist4096MaxLen)) {
+    if (rank == 0) {
+      extern __shared__ uint8_t sr[];
+      hist4096::histogram_4096_topk<TopK, 12>(
+          in, out, sl, sr);  // 4096-bin (12-bit) histogram
+    }
+    return;
+  }
+
+  // Large path: init mbarriers + state, then dispatch fused or twopass
+  const uint32_t per_block =
+      (params.stride + CS - 1) / CS;  // how many elements per block
+  constexpr uint32_t kFusedMax = ((CS == 16)  ? kFusedStagesCS16
+                                  : (CS == 8) ? kFusedStagesCS8
+                                              : kMaxSinglePassStages) *
+                                 kSizePerStage;
+  const bool use_singlepass =
+      per_block <=
+      kFusedMax;  // single pass or TMA streaming: histogram+scatter
+
+  // Select smem type and stage count at compile time based on CS
+  constexpr uint32_t kFusedStages = (CS == 16)  ? kFusedStagesCS16
+                                    : (CS == 8) ? kFusedStagesCS8
+                                                : kMaxSinglePassStages;
+  using FusedSmem = SmemFused<kFusedStages>;
+
+  extern __shared__ uint8_t sr[];
+
+  constexpr uint32_t kTieWsPerRow =
+      TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
+  hist4096::Tie* row_tie_ws = params.tie_ws + row * kTieWsPerRow;
+
+  if (use_singlepass) {
+    auto* smem = reinterpret_cast<FusedSmem*>(sr);
+    const uint32_t sp_stages = (per_block + kSizePerStage - 1) / kSizePerStage;
+    if (tx < sp_stages) {
+      mbarrier_init(&smem->barrier[0][tx],
+                    1);  // init 1 barrier per TMA stage -
+                         // signal when async copies complete
+    }
+    __syncthreads();
+    uint32_t phases[kFusedStages] =
+        {};  // tracks the parity for mbarrier wait/arrive protocol
+    large_topk<TopK, CS, FusedSmem, true>(in, out, sl, phases, row_tie_ws);
+  } else {
+    // Two-pass: only CS=4 in practice (CS=8 always fits in singlepass)
+    auto* smem = reinterpret_cast<Smem4*>(sr);
+    if (tx < 2 * kStreamingStagesCS4) {
+      mbarrier_init(&smem->barrier[0][tx],
+                    1);  // init 2×2=4 barriers (2 passes × 2 stages)
+    }
+    __syncthreads();
+    uint32_t hp[kStreamingStagesCS4] = {0,
+                                        0};  // histogram+scatter pass counters
+    large_topk<TopK, CS, Smem4, false>(in, out, sl, hp, row_tie_ws);
+  }
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 4, 1)
+        cooperative_topk_cs4(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 4>(params);
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 8, 1)
+        cooperative_topk_cs8(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 8>(params);
+}
+
+template <uint32_t TopK>
+__global__ void __launch_bounds__(hist4096::kBlockSize, 1)
+    __cluster_dims__(1, 16, 1)
+        cooperative_topk_cs16(CooperativeTopKParams<TopK> params) {
+  cooperative_topk_body<TopK, 16>(params);
+}
+
+constexpr size_t kSmemSize4_base = sizeof(Smem4);
+constexpr size_t kSmemSize4_sp = sizeof(SmemSinglePass);
+constexpr size_t kSmemSize4 =
+    (kSmemSize4_base > kSmemSize4_sp ? kSmemSize4_base : kSmemSize4_sp) +
+    sizeof(int32_t) * 2048 + 128;
+constexpr size_t kSmemSize8 =
+    sizeof(SmemFused<kFusedStagesCS8>) + sizeof(int32_t) * 2048 + 128;
+
+}  // namespace cooperative
+
+}  // namespace vllm
+
+#endif  // COOPERATIVE_TOPK_CUH_

csrc/libtorch_stable/cuda_vec_utils.cuh

@@ -21,7 +21,7 @@
 // together enable 256-bit (v8.u32) PTX load/store instructions.
 // Use for PTX instruction selection with architecture fallback paths.
 #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
-    defined(CUDA_VERSION) && CUDA_VERSION >= 12090
+    defined(CUDART_VERSION) && CUDART_VERSION >= 12090
   #define VLLM_256B_PTX_ENABLED 1
 #else
   #define VLLM_256B_PTX_ENABLED 0

csrc/libtorch_stable/dispatch_utils.h

@@ -30,6 +30,28 @@
   THO_DISPATCH_SWITCH(TYPE, NAME,                            \
                       VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
 
+#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(...)                  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Byte, __VA_ARGS__)  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Char, __VA_ARGS__)  \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Short, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Int, __VA_ARGS__)   \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::Long, __VA_ARGS__)
+
+#define VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...)      \
+  VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)                 \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt16, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt32, __VA_ARGS__) \
+  THO_DISPATCH_CASE(torch::headeronly::ScalarType::UInt64, __VA_ARGS__)
+
+#define VLLM_STABLE_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH(TYPE, NAME,                            \
+                      VLLM_STABLE_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
+
+#define VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
+  THO_DISPATCH_SWITCH(                                                    \
+      TYPE, NAME,                                                         \
+      VLLM_STABLE_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
+
 // FP8 type dispatch - ROCm uses FNUZ format, CUDA uses OCP format
 #ifdef USE_ROCM
   #define VLLM_STABLE_DISPATCH_CASE_FP8_TYPES(...)                 \

csrc/libtorch_stable/fp32_router_gemm.cu

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