[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)
[CI] Raise gsm8k startup timeout for MoE Refactor Qwen3 NVFP4 configs (#46882 )
2026-06-28 00:04:08 -07:00 · 2026-06-26 16:33:47 -07:00 · 2026-06-26 16:15:00 -07:00 · 2026-06-25 13:15:20 -07:00 · 2026-06-25 13:15:20 -07:00 · 2026-06-25 13:15:02 -07:00
213 changed files with 12341 additions and 1807 deletions
@@ -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: []
@@ -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 \
@@ -1,5 +1,5 @@
 group: Expert Parallelism
-depends_on: 
+depends_on:
  - image-build-xpu
 steps:
 - label: EPLB Algorithm
@@ -1,5 +1,5 @@
 group: Models - Multimodal
-depends_on: 
+depends_on:
  - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
@@ -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]
@@ -432,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
@@ -500,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]
@@ -780,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]
@@ -807,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
@@ -980,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]
@@ -1057,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]
@@ -1410,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]
@@ -1742,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
@@ -1770,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
@@ -2336,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]
@@ -3411,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
@@ -12,6 +12,7 @@ steps:
  - tests/basic_correctness/test_cpu_offload
  - tests/basic_correctness/test_mem.py
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_mem.py
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_basic_correctness.py
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_cpu_offload.py
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s basic_correctness/test_mem.py
+  - pytest -v -s basic_correctness/test_basic_correctness.py
+  - pytest -v -s basic_correctness/test_cpu_offload.py
@@ -14,7 +14,8 @@ steps:
  - vllm/v1/cudagraph_dispatcher.py
  - tests/compile/correctness_e2e/test_sequence_parallel.py
  commands:
-  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py

 - label: Sequence Parallel Correctness Tests (2xH100)
  key: sequence-parallel-correctness-tests-2xh100
@@ -24,7 +25,8 @@ steps:
  optional: true
  num_devices: 2
  commands:
-  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py

 - label: AsyncTP Correctness Tests (2xH100)
  key: asynctp-correctness-tests-2xh100
@@ -34,7 +36,8 @@ steps:
  optional: true
  num_devices: 2
  commands:
-  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py

 - label: AsyncTP Correctness Tests (B200)
  key: asynctp-correctness-tests-b200
@@ -44,7 +47,8 @@ steps:
  optional: true
  num_devices: 2
  commands:
-  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py

 - label: Distributed Compile Unit Tests (2xH100)
  key: distributed-compile-unit-tests-2xh100
@@ -57,7 +61,8 @@ steps:
  - vllm/model_executor/layers
  - tests/compile/passes/distributed/
  commands:
-  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -s -v tests/compile/passes/distributed
+  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
+  - pytest -s -v tests/compile/passes/distributed

 - label: Fusion and Compile Unit Tests (2xB200)
  key: fusion-and-compile-unit-tests-2xb200
@@ -32,10 +32,11 @@ steps:
  - tests/entrypoints/openai/test_multi_api_servers.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - 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: Distributed Compile + RPC Tests (2 GPUs)
  key: distributed-compile-rpc-tests-2-gpus
@@ -55,9 +56,10 @@ steps:
  - tests/entrypoints/llm/test_collective_rpc.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/test_wrapper.py
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s ./compile/test_wrapper.py

 - label: Distributed Torchrun + Shutdown Tests (2 GPUs)
  key: distributed-torchrun-shutdown-tests-2-gpus
@@ -76,10 +78,11 @@ steps:
  - tests/v1/worker/test_worker_memory_snapshot.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - NCCL_CUMEM_HOST_ENABLE=0 VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
-  - NCCL_CUMEM_HOST_ENABLE=0 CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/worker/test_worker_memory_snapshot.py
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - VLLM_TEST_SAME_HOST=1 VLLM_TEST_WITH_DEFAULT_DEVICE_SET=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+  - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
+  - pytest -v -s v1/worker/test_worker_memory_snapshot.py

 - label: Distributed Torchrun + Examples (4 GPUs)
  key: distributed-torchrun-examples-4-gpus
@@ -94,23 +97,24 @@ steps:
  - tests/examples/features/data_parallel/data_parallel_offline.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
+  - export NCCL_CUMEM_HOST_ENABLE=0
  # test with torchrun tp=2 and external_dp=2
-  - NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+  - torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
  # test with torchrun tp=2 and pp=2
-  - NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
+  - PP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example.py
  # test with torchrun tp=4 and dp=1
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - TP_SIZE=4 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
  # test with torchrun tp=2, pp=2 and dp=1
-  - NCCL_CUMEM_HOST_ENABLE=0 PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
  # test with torchrun tp=1 and dp=4 with ep
-  - NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
  # test with torchrun tp=2 and dp=2 with ep
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
+  - TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
  # test with internal dp
-  - NCCL_CUMEM_HOST_ENABLE=0 python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
+  - python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
  # rlhf examples
-  - NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
-  - NCCL_CUMEM_HOST_ENABLE=0 VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
+  - VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py

 - label: Distributed DP Tests (4 GPUs)
  key: distributed-dp-tests-4-gpus
@@ -124,13 +128,14 @@ steps:
  - tests/distributed/test_utils
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_utils.py
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - TP_SIZE=2 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
+  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py
+  - TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py
+  - pytest -v -s v1/engine/test_engine_core_client.py::test_kv_cache_events_dp
+  - pytest -v -s distributed/test_utils.py

 - label: Distributed Compile + Comm (4 GPUs)
  key: distributed-compile-comm-4-gpus
@@ -146,12 +151,13 @@ steps:
  - tests/distributed/test_multiproc_executor.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s compile/fullgraph/test_basic_correctness.py
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pynccl.py
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_events.py
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_symm_mem_allreduce.py
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - pytest -v -s compile/fullgraph/test_basic_correctness.py
+  - pytest -v -s distributed/test_pynccl.py
+  - pytest -v -s distributed/test_events.py
+  - pytest -v -s distributed/test_symm_mem_allreduce.py
  # test multi-node TP with multiproc executor (simulated on single node)
-  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
+  - pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node

 - label: Distributed Tests (8 GPUs)(H100)
  key: distributed-tests-8-gpus-h100
@@ -170,8 +176,9 @@ steps:

  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
+  - export NCCL_CUMEM_HOST_ENABLE=0
  # test with torchrun tp=2 and dp=4 with ep
-  - NCCL_CUMEM_HOST_ENABLE=0 torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
+  - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep

 - label: Distributed Tests (4 GPUs)(A100)
  key: distributed-tests-4-gpus-a100
@@ -264,7 +271,9 @@ steps:
  - tests/distributed/test_pipeline_parallel.py
  - tests/basic_correctness/test_basic_correctness.py
  commands:
-  - VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor.py
-  - VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_ray_v2_executor_e2e.py
-  - VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
-  - VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1 NCCL_CUMEM_HOST_ENABLE=0 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
+  - export VLLM_USE_RAY_V2_EXECUTOR_BACKEND=1
+  - export NCCL_CUMEM_HOST_ENABLE=0
+  - pytest -v -s distributed/test_ray_v2_executor.py
+  - pytest -v -s distributed/test_ray_v2_executor_e2e.py
+  - pytest -v -s distributed/test_pipeline_parallel.py -k "ray"
+  - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -k "ray"
@@ -22,9 +22,10 @@ steps:
  - vllm/
  - tests/entrypoints/llm
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py --ignore=entrypoints/llm/offline_mode
+  - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
+  - pytest -v -s entrypoints/llm/offline_mode # Needs to avoid interference with other tests
  mirror:
    amd:
      device: mi325_1
@@ -40,8 +41,9 @@ steps:
  - vllm/
  - tests/entrypoints/serve
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/serve --ignore=entrypoints/serve/dev/rpc
+  - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc
  mirror:
    amd:
      device: mi325_1
@@ -57,7 +59,8 @@ steps:
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/correctness
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - 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: mi325_1
@@ -74,8 +77,9 @@ steps:
  - tests/entrypoints/openai
  - tests/entrypoints/test_chat_utils
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/openai/chat_completion
+  - pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
  mirror:
    amd:
      device: mi325_1
@@ -124,7 +128,8 @@ steps:
  - vllm/
  - tests/entrypoints/speech_to_text
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/speech_to_text
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/speech_to_text

 - label: Entrypoints Integration (Multimodal)
  device: h200_35gb
@@ -135,7 +140,8 @@ steps:
  - vllm/
  - tests/entrypoints/multimodal
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/multimodal
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/multimodal

 - label: Entrypoints Integration (Pooling)
  key: entrypoints-integration-pooling
@@ -145,7 +151,8 @@ steps:
  - vllm/
  - tests/entrypoints/pooling
  commands:
-  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/pooling
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s entrypoints/pooling

 - label: OpenAI API Correctness
  key: openai-api-correctness
@@ -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
@@ -50,7 +50,8 @@ steps:
  - csrc/
  - vllm/model_executor/layers/quantization
  commands:
-    - VLLM_USE_DEEP_GEMM=0 pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4  # Triton is faster than DeepGEMM for H100
+    - export VLLM_USE_DEEP_GEMM=0  # We found Triton is faster than DeepGEMM for H100
+    - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4

 - label: LM Eval Small Models (B200)
  key: lm-eval-small-models-b200
@@ -107,7 +108,9 @@ steps:
      depends_on:
      - image-build-amd
      commands:
-      - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_ROCM_ARCH=gfx942 pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt  # Limit Quark compilation to save time
+      - 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
@@ -133,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
@@ -174,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
@@ -36,14 +36,14 @@ steps:
  commands:
    # FIXIT: find out which code initialize cuda before running the test
    # before the fix, we need to use spawn to test it
-    #
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # Alot of these tests are on the edge of OOMing
-    #
+    - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
    # There is some Tensor Parallelism related processing logic in LoRA that
    # requires multi-GPU testing for validation.
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_chatglm3_tp.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_llama_tp.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen3_with_multi_loras.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_olmoe_tp.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_gptoss_tp.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+    - pytest -v -s -x lora/test_chatglm3_tp.py
+    - pytest -v -s -x lora/test_llama_tp.py
+    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
+    - pytest -v -s -x lora/test_olmoe_tp.py
+    - pytest -v -s -x lora/test_gptoss_tp.py
+    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -18,8 +18,9 @@ steps:
    - vllm/v1/
    - tests/v1/spec_decode
  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # TODO: create another `optional` test group for slow tests
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not slow_test' v1/spec_decode
+    - pytest -v -s -m 'not slow_test' v1/spec_decode
  mirror:
    amd:
      device: mi300_1
@@ -49,11 +50,12 @@ steps:
    - tests/v1/test_request.py
    - tests/v1/test_outputs.py
  commands:
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/logits_processors
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_oracle.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_request.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_outputs.py
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s v1/sample
+    - pytest -v -s v1/logits_processors
+    - pytest -v -s v1/test_oracle.py
+    - pytest -v -s v1/test_request.py
+    - pytest -v -s v1/test_outputs.py
  mirror:
    amd:
      device: mi325_1
@@ -91,17 +93,18 @@ steps:
    - tests/entrypoints/openai/correctness/test_lmeval.py
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # split the test to avoid interference
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/core
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/executor
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_offload
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/simple_kv_offload
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/worker
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/metrics
+    - pytest -v -s -m 'not cpu_test' v1/core
+    - pytest -v -s v1/executor
+    - pytest -v -s v1/kv_offload
+    - pytest -v -s v1/simple_kv_offload
+    - pytest -v -s v1/worker
+    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
+    - pytest -v -s -m 'not cpu_test' v1/metrics
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
      device: mi325_1
@@ -150,7 +153,8 @@ steps:
    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
    - tests/v1/kv_connector/extract_hidden_states_integration
  commands:
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_connector/extract_hidden_states_integration
+    - 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
@@ -163,7 +167,8 @@ steps:
    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
    - tests/v1/kv_connector/extract_hidden_states_integration
  commands:
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration

 - label: Regression
  key: regression
@@ -355,9 +360,10 @@ steps:
    - vllm/model_executor/layers
    - tests/v1/determinism/
  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pip install pytest-timeout pytest-forked
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
+    - pytest -v -s v1/determinism/test_batch_invariance.py
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]

 - label: Batch Invariance (H100)
  key: batch-invariance-h100
@@ -368,11 +374,12 @@ steps:
    - vllm/model_executor/layers
    - tests/v1/determinism/
  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pip install pytest-timeout pytest-forked
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - pytest -v -s v1/determinism/test_batch_invariance.py
+    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]

 - label: Batch Invariance (B200)
  key: batch-invariance-b200
@@ -383,13 +390,14 @@ steps:
    - vllm/model_executor/layers
    - tests/v1/determinism/
  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pip install pytest-timeout pytest-forked
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_batch_invariance.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
+    - pytest -v -s v1/determinism/test_batch_invariance.py
+    - pytest -v -s v1/determinism/test_rms_norm_batch_invariant.py
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
+    - VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
+    - pytest -v -s v1/determinism/test_nvfp4_batch_invariant_scaled_mm.py
  
 - label: Acceptance Length Test (Large Models) # optional
  device: h200_35gb
@@ -404,4 +412,5 @@ steps:
  - vllm/model_executor/models/mlp_speculator.py
  - tests/v1/spec_decode/test_acceptance_length.py
  commands:
-    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
+    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
+    - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test
@@ -13,12 +13,13 @@ steps:
  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
+    # init surfaces a stack instead of silently stalling.
+    - export PYTHONFAULTHANDLER=1
    # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
    # init) fails fast with a traceback instead of running until the global
    # build timeout. The `thread` method also handles hangs inside C/CUDA
    # calls that the signal method cannot interrupt.
-    #
-    # Env vars are inlined because CONTINUE_ON_FAILURE wraps each command
-    # in a subshell, so a standalone `export` would be lost.
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
-    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTHONFAULTHANDLER=1 pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
+    - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
@@ -16,14 +16,15 @@ steps:
  - tests/entrypoints/llm/test_struct_output_generate.py
  commands:
  - set -x
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s v1/engine/test_llm_engine.py -k "not test_engine_metrics"
  # This requires eager until we sort out CG correctness issues.
  # TODO: remove ENFORCE_EAGER here after https://github.com/vllm-project/vllm/pull/32936 is merged.
-  - VLLM_USE_V2_MODEL_RUNNER=1 ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_context_length.py
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/general/test_min_tokens.py
+  - ENFORCE_EAGER=1 pytest -v -s v1/e2e/general/test_async_scheduling.py -k "not ngram"
+  - pytest -v -s v1/e2e/general/test_context_length.py
+  - pytest -v -s v1/e2e/general/test_min_tokens.py
  # Temporary hack filter to exclude ngram spec decoding based tests.
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
+  - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"

 - label: Model Runner V2 Examples
  device: h200_35gb
@@ -41,25 +42,26 @@ steps:
    - examples/features/tensorize_vllm_model.py
  commands:
    - set -x
+    - export VLLM_USE_V2_MODEL_RUNNER=1
    - pip install tensorizer # for tensorizer test
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/chat.py # for basic
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    #- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
-    #- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/embed.py   # TODO
+    - python3 basic/offline_inference/chat.py # for basic
+    - 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  # TODO
+    #- python3 basic/offline_inference/embed.py   # TODO
    # for multi-modal models
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/audio_language_offline.py --seed 0
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_offline.py --seed 0
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+    - 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
    # for pooling models
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 pooling/embed/vision_embedding_offline.py --seed 0
+    - python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 deployment/llm_engine_example.py
-    - VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-    - VLLM_USE_V2_MODEL_RUNNER=1 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/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
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
-    - VLLM_USE_V2_MODEL_RUNNER=1 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
+    - 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

 - label: Model Runner V2 Distributed (2 GPUs)
  key: model-runner-v2-distributed-2-gpus
@@ -74,11 +76,13 @@ steps:
    - tests/v1/distributed/test_eagle_dp.py
  commands:
    - set -x
+    - export VLLM_USE_V2_MODEL_RUNNER=1
    # The "and not True" here is a hacky way to exclude the prompt_embeds cases which aren't yet supported.
-    - VLLM_USE_V2_MODEL_RUNNER=1 TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
+    - TARGET_TEST_SUITE=L4 pytest -v -s basic_correctness/test_basic_correctness.py -m 'distributed(num_gpus=2)' -k "not ray and not True"
    # https://github.com/NVIDIA/nccl/issues/1838
-    - VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
-    - VLLM_USE_V2_MODEL_RUNNER=1 NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+    - export NCCL_CUMEM_HOST_ENABLE=0
+    - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py -k "not ray"
+    - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py

 - label: Model Runner V2 Pipeline Parallelism (4 GPUs)
  key: model-runner-v2-pipeline-parallelism-4-gpus
@@ -93,9 +97,10 @@ steps:
    - tests/v1/distributed/test_pp_dp_v2.py
  commands:
    - set -x
-    - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
-    - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
-    - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/distributed/test_pp_dp_v2.py
+    - export VLLM_USE_V2_MODEL_RUNNER=1
+    - pytest -v -s distributed/test_pipeline_parallel.py -k "not ray and not Jamba"
+    - pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
+    - pytest -v -s v1/distributed/test_pp_dp_v2.py

 - label: Model Runner V2 Spec Decode
  device: h200_35gb
@@ -110,7 +115,8 @@ steps:
  - tests/v1/e2e/spec_decode/test_spec_decode.py
  commands:
  - set -x
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
-  - VLLM_USE_V2_MODEL_RUNNER=1 pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
+  - export VLLM_USE_V2_MODEL_RUNNER=1
+  - pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
+  - pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
+  - pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
+  - pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
@@ -23,15 +23,17 @@ steps:
  # - tests/entrypoints/openai/test_uds.py
  - tests/v1/sample/test_logprobs_e2e.py
  commands:
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
-  # - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
+  - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
+  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"

-  # - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
-  # - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_return_token_ids.py
-  # - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/test_uds.py
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
+  # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
+  - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
+  # - pytest -v -s entrypoints/openai/test_return_token_ids.py
+  # - pytest -v -s entrypoints/openai/test_uds.py
+  - pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"

 - label: Rust Frontend Serve/Admin Coverage
  timeout_in_minutes: 60
@@ -49,11 +51,13 @@ steps:
  - tests/entrypoints/serve/instrumentator/test_metrics.py
  # - tests/entrypoints/serve/dev/test_sleep.py
  commands:
-  # - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
-  # - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/serve/dev/test_sleep.py
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  # - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
+  - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
+  - pytest -v -s entrypoints/serve/disagg/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
+  - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
+  # - pytest -v -s entrypoints/serve/dev/test_sleep.py

 - label: Rust Frontend Core Correctness
  timeout_in_minutes: 30
@@ -65,7 +69,9 @@ steps:
  - tests/utils.py
  - tests/entrypoints/openai/correctness/test_lmeval.py
  commands:
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine

 - label: Rust Frontend Tool Use
  timeout_in_minutes: 60
@@ -77,7 +83,9 @@ steps:
  - tests/utils.py
  - tests/tool_use/
  commands:
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
+  - export VLLM_USE_RUST_FRONTEND=1
+  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  - pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"

 - label: Rust Frontend Distributed
  timeout_in_minutes: 30
@@ -95,6 +103,9 @@ steps:
  - tests/v1/distributed/test_hybrid_lb_dp.py
  - tests/v1/distributed/test_internal_lb_dp.py
  commands:
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn 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"
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
-  - VLLM_USE_RUST_FRONTEND=1 VLLM_WORKER_MULTIPROC_METHOD=spawn NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
+  - 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"
@@ -153,7 +153,8 @@ steps:
    - vllm/model_executor/models/qwen3_dflash.py
    - tests/v1/spec_decode/test_speculators_correctness.py
  commands:
-    - VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
+    - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
+    - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test

 - label: Spec Decode MTP hybrid (B200)
  timeout_in_minutes: 30
@@ -382,6 +382,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/custom_all_reduce.cu"
    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

+  if(VLLM_GPU_LANG STREQUAL "CUDA" AND
+     DEFINED CMAKE_CUDA_COMPILER_VERSION AND
+     CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
+
+    if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
+        "9.0a;10.0f;10.1f;10.3f;11.0f;12.0f;12.1f" "${CUDA_ARCHS}")
+    else()
+      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
+        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
+    endif()
+
+    if(COOPERATIVE_TOPK_ARCHS)
+      list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
+
+    endif()
+  endif()
+
  if(VLLM_GPU_LANG STREQUAL "CUDA")
    SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

@@ -498,6 +516,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${VLLM_STABLE_EXT_SRC}"
      CUDA_ARCHS "${CUDA_ARCHS}")

+    if(COOPERATIVE_TOPK_ARCHS)
+      list(APPEND VLLM_STABLE_EXT_SRC
+        "csrc/libtorch_stable/cooperative_topk.cu")
+      set_gencode_flags_for_srcs(
+        SRCS "csrc/libtorch_stable/cooperative_topk.cu"
+        CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
+    endif()
+
    # Only build Marlin kernels if we are building for at least some compatible archs.
    # Keep building Marlin for 9.0 as there are some group sizes and shapes that
    # are not supported by Machete yet.
@@ -843,9 +869,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  endif()

  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
  else()
-    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.3a" "${CUDA_ARCHS}")
+    cuda_archs_loose_intersection(SCALED_MM_ARCHS "10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
    set(CUTLASS_MOE_SM100_SRCS "csrc/libtorch_stable/quantization/w8a8/cutlass/moe/grouped_mm_c3x_sm100.cu")
@@ -1049,6 +1075,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    target_compile_definitions(_C_stable_libtorch PRIVATE
      TORCH_TARGET_VERSION=0x020B000000000000ULL)
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
+    if(COOPERATIVE_TOPK_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_COOPERATIVE_TOPK=1)
+    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -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
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
+    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -36,13 +36,38 @@ 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)
@@ -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.")
@@ -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
+}
@@ -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_
@@ -67,6 +67,13 @@
  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif

+// Direct float -> E4M3 FP8 conversion for the indexer Q / index-K outputs.
+#ifndef USE_ROCM
+  #include <cuda_fp8.h>
+#else
+  #include <hip/hip_fp8.h>
+#endif
+
 #ifndef FINAL_MASK
  #ifdef USE_ROCM
    #define FINAL_MASK 0xffffffffffffffffULL
@@ -75,6 +82,19 @@
  #endif
 #endif

+#ifdef USE_ROCM
+// ROCm-compatible direct float -> E4M3 FP8 conversion (mirrors the DeepSeek V4
+// fused kernel).
+__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
+  #if defined(HIP_FP8_TYPE_OCP)
+  __hip_fp8_e4m3 fp8_val(val);
+  #else
+  __hip_fp8_e4m3_fnuz fp8_val(val);
+  #endif
+  return reinterpret_cast<uint8_t&>(fp8_val);
+}
+#endif
+
 namespace vllm {
 namespace minimax_m3_fused_ops {

@@ -193,6 +213,8 @@ __device__ __forceinline__ void storeElems(
  *reinterpret_cast<uint2*>(dst) = v;
 }

+// Main K/V cache store. kAuto = unquantized (cache_t == scalar_t); fp8 cache
+// dtypes use the scaled-convert path with identity scale.
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
 __device__ __forceinline__ void storeCacheElems(
    cache_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
@@ -208,6 +230,32 @@ __device__ __forceinline__ void storeCacheElems(
  }
 }

+// Store 4 fp32 registers -> 4 contiguous E4M3 FP8 bytes (direct cast,
+// saturating to ±448). Used for the fp8 indexer-Q / index-K outputs; no scale
+// (RMSNorm outputs are O(1) and the score path only needs relative block
+// ordering).
+__device__ __forceinline__ void storeElemsFp8(
+    uint8_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
+  constexpr float kFp8Max = 448.0f;
+#ifndef USE_ROCM
+  __nv_fp8x2_storage_t out2[kElemsPerLane / 2];
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane / 2; i++) {
+    float2 vv = make_float2(elems[2 * i], elems[2 * i + 1]);
+    vv.x = fminf(fmaxf(vv.x, -kFp8Max), kFp8Max);
+    vv.y = fminf(fmaxf(vv.y, -kFp8Max), kFp8Max);
+    out2[i] = __nv_cvt_float2_to_fp8x2(vv, __NV_SATFINITE, __NV_E4M3);
+  }
+  *reinterpret_cast<uint32_t*>(dst) = *reinterpret_cast<uint32_t const*>(out2);
+#else
+  #pragma unroll
+  for (int i = 0; i < kElemsPerLane; i++) {
+    float vv = fminf(fmaxf(elems[i], -kFp8Max), kFp8Max);
+    dst[i] = rocm_cvt_float_to_fp8_e4m3(vv);
+  }
+#endif
+}
+
 // ────────────────────────────────────────────────────────────────────────────
 // Kernel
 // ────────────────────────────────────────────────────────────────────────────
@@ -224,12 +272,14 @@ __device__ __forceinline__ void storeCacheElems(
 //     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
 //     IQ: niq  only if kIsSparse        (norm+RoPE)
 //     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
+// cache_t/kv_dt: main attention KV-cache dtype (auto/fp8). out_idx_t/kFp8Idx:
+// indexer index-K cache + index-Q output dtype (scalar_t or e4m3 byte).
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt,
-          bool kIsSparse, bool kInsertKV>
+          typename out_idx_t, bool kIsSparse, bool kInsertKV, bool kFp8Idx>
 __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
    scalar_t* __restrict__ qkv,  // [N, qkv_row] in/out (packs index if sparse)
-    scalar_t* __restrict__ q_out,        // [N, nq*128] contiguous, or nullptr
-    scalar_t* __restrict__ index_q_out,  // [N, niq*128] contiguous, or nullptr
+    scalar_t* __restrict__ q_out,         // [N, nq*128] contiguous, or nullptr
+    out_idx_t* __restrict__ index_q_out,  // [N, niq*128]; scalar_t or e4m3 byte
    scalar_t const* __restrict__ q_norm_w,
    scalar_t const* __restrict__ k_norm_w,
    scalar_t const* __restrict__ iq_norm_w,
@@ -238,8 +288,8 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
    int64_t const* __restrict__ positions,       // [N] i64
    int64_t const* __restrict__ slot_mapping,    // main K/V slots or nullptr
    int64_t const* __restrict__ index_slot_mapping,  // index K slots/nullptr
-    cache_t* __restrict__ kv_cache,      // [nb,2,bs,nkv,128] or nullptr
-    scalar_t* __restrict__ index_cache,  // [nb*bs, 128] or nullptr
+    cache_t* __restrict__ kv_cache,       // [nb,2,bs,nkv,128] or nullptr
+    out_idx_t* __restrict__ index_cache,  // [nb*bs, 128]; scalar_t or e4m3 byte
    float const eps, int const rotary_dim, int const num_tokens, int const nq,
    int const nkv, int const niq, int const block_size,
    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
@@ -334,9 +384,12 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
      store_ptr = q_out + static_cast<int64_t>(tokenIdx) * nq * kHeadDim +
                  slot * kHeadDim;
    } else if (isIQ && index_q_out != nullptr) {
-      store_ptr = index_q_out +
-                  static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
-                  (slot - iq_begin) * kHeadDim;
+      // bf16 index_q_out: gather here. fp8: written by the explicit fp8 store.
+      if constexpr (!kFp8Idx) {
+        store_ptr = index_q_out +
+                    static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                    (slot - iq_begin) * kHeadDim;
+      }
    }

    // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
@@ -356,7 +409,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
      scalar_t const* cos_ptr = cos_sin_cache + pos * rotary_dim;
      normAndRope<scalar_t>(elems, laneId, eps, norm_w, do_rope, rotary_dim,
                            cos_ptr, /*apply_norm=*/norm_w != nullptr);
-      storeElems<scalar_t>(store_ptr + dim_base, elems);
+      if constexpr (kFp8Idx) {
+        // index_q is e4m3 bytes; Q/K (and in-place index_k) stay scalar_t.
+        if (isIQ && index_q_out != nullptr) {
+          storeElemsFp8(index_q_out +
+                            static_cast<int64_t>(tokenIdx) * niq * kHeadDim +
+                            (slot - iq_begin) * kHeadDim + dim_base,
+                        elems);
+        } else {
+          storeElems<scalar_t>(store_ptr + dim_base, elems);
+        }
+      } else {
+        storeElems<scalar_t>(store_ptr + dim_base, elems);
+      }
    }

    // ── Cache inserts (sparse serving only). ───────────────────────────────
@@ -367,8 +432,11 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
                             : (isIK ? index_slot_mapping[tokenIdx] : -1);
      if (sm >= 0) {  // skip padded / unscheduled tokens
        if (isIK) {
-          scalar_t* dst = index_cache + sm * kHeadDim + dim_base;
-          storeElems<scalar_t>(dst, elems);
+          if constexpr (kFp8Idx) {
+            storeElemsFp8(index_cache + sm * kHeadDim + dim_base, elems);
+          } else {
+            storeElems<scalar_t>(index_cache + sm * kHeadDim + dim_base, elems);
+          }
        } else if (isK || isV) {
          // kv_cache logical shape [num_blocks, 2, block_size, nkv, head_dim].
          // Paging is logical (block = sm/block_size, token = sm%block_size);
@@ -398,19 +466,19 @@ __global__ void fusedMiniMaxM3QNormRopeKVInsertKernel(
 // Launch wrapper
 // ────────────────────────────────────────────────────────────────────────────
 template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
-void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
-                          scalar_t const* q_norm_w, scalar_t const* k_norm_w,
-                          scalar_t const* iq_norm_w, scalar_t const* ik_norm_w,
-                          scalar_t const* cos_sin_cache,
-                          int64_t const* positions, int64_t const* slot_mapping,
-                          int64_t const* index_slot_mapping, cache_t* kv_cache,
-                          scalar_t* index_cache, float const eps,
-                          int const rotary_dim, int const num_tokens,
-                          int const nq, int const nkv, int const niq,
-                          int const block_size, int64_t const kv_s_block,
-                          int64_t const kv_s_kv, int64_t const kv_s_token,
-                          int64_t const kv_s_head, bool const has_index,
-                          bool const insert_kv, cudaStream_t stream) {
+void launchFusedMiniMaxM3(
+    scalar_t* qkv, scalar_t* q_out, void* index_q_out, scalar_t const* q_norm_w,
+    scalar_t const* k_norm_w, scalar_t const* iq_norm_w,
+    scalar_t const* ik_norm_w, scalar_t const* cos_sin_cache,
+    int64_t const* positions, int64_t const* slot_mapping,
+    int64_t const* index_slot_mapping, cache_t* kv_cache, void* index_cache,
+    float const eps, int const rotary_dim, int const num_tokens, int const nq,
+    int const nkv, int const niq, int const block_size,
+    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
+    int64_t const kv_s_head, bool const has_index, bool const insert_kv,
+    bool const fp8_idx, cudaStream_t stream) {
+  // Index outputs are scalar_t (bf16) or e4m3 bytes (uint8_t); reinterpret the
+  // void* pointers per instantiation in the LAUNCH macro.
  // Slot count must match the kernel's compile-time gating.
  int const v_slots = insert_kv ? nkv : 0;
  int const idx_slots = has_index ? niq + 1 : 0;
@@ -440,25 +508,27 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
  config.attrs = attrs;
  config.numAttrs = (sm_version >= 90) ? 1 : 0;

-  #define LAUNCH(IS_SPARSE, INSERT)                                           \
-    cudaLaunchKernelEx(                                                       \
-        &config,                                                              \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,       \
-                                              IS_SPARSE, INSERT>,             \
-        qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w, ik_norm_w,    \
-        cos_sin_cache, positions, slot_mapping, index_slot_mapping, kv_cache, \
-        index_cache, eps, rotary_dim, num_tokens, nq, nkv, niq, block_size,   \
-        kv_s_block, kv_s_kv, kv_s_token, kv_s_head)
+  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                                \
+    cudaLaunchKernelEx(                                                        \
+        &config,                                                               \
+        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T, \
+                                              IS_SPARSE, INSERT, FP8>,         \
+        qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w, k_norm_w, \
+        iq_norm_w, ik_norm_w, cos_sin_cache, positions, slot_mapping,          \
+        index_slot_mapping, kv_cache, reinterpret_cast<OUT_T*>(index_cache),   \
+        eps, rotary_dim, num_tokens, nq, nkv, niq, block_size, kv_s_block,     \
+        kv_s_kv, kv_s_token, kv_s_head)
 #else
  // ROCm: standard kernel launch syntax (no PDL/stream serialization).
  // clang-format off
-  #define LAUNCH(IS_SPARSE, INSERT)                                          \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt,          \
-                                          IS_SPARSE, INSERT>                 \
+  #define LAUNCH(IS_SPARSE, INSERT, FP8, OUT_T)                              \
+    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, cache_t, kv_dt, OUT_T,   \
+                                          IS_SPARSE, INSERT, FP8>            \
        <<<grid, kBlockSize, 0, stream>>>(                                   \
-            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
-            ik_norm_w, cos_sin_cache, positions, slot_mapping,               \
-            index_slot_mapping, kv_cache, index_cache, eps, rotary_dim,      \
+            qkv, q_out, reinterpret_cast<OUT_T*>(index_q_out), q_norm_w,     \
+            k_norm_w, iq_norm_w, ik_norm_w, cos_sin_cache, positions,        \
+            slot_mapping, index_slot_mapping, kv_cache,                      \
+            reinterpret_cast<OUT_T*>(index_cache), eps, rotary_dim,          \
            num_tokens, nq, nkv, niq, block_size, kv_s_block, kv_s_kv,       \
            kv_s_token, kv_s_head)
  // clang-format on
@@ -466,14 +536,22 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,

  if (has_index) {
    if (insert_kv) {
-      LAUNCH(true, true);  // sparse serving
+      if (fp8_idx) {
+        LAUNCH(true, true, true, uint8_t);  // sparse serving, fp8 index outputs
+      } else {
+        LAUNCH(true, true, false, scalar_t);  // sparse serving, bf16
+      }
    } else {
-      LAUNCH(true, false);  // sparse profiling
+      if (fp8_idx) {
+        LAUNCH(true, false, true, uint8_t);  // sparse profiling, fp8 index_q
+      } else {
+        LAUNCH(true, false, false, scalar_t);  // sparse profiling, bf16
+      }
    }
  } else {
    // Dense layer: never has an index branch and never inserts here (the
    // generic Attention layer owns the KV insert).
-    LAUNCH(false, false);
+    LAUNCH(false, false, false, scalar_t);
  }
 #undef LAUNCH
 }
@@ -485,8 +563,9 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
  vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st, CACHE_T, KV_DTYPE>(     \
      reinterpret_cast<st*>(qkv.data_ptr()),                                   \
      q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr()) : nullptr,  \
-      index_q_out.has_value() ? reinterpret_cast<st*>(index_q_out->data_ptr()) \
-                              : nullptr,                                       \
+      index_q_out.has_value()                                                  \
+          ? reinterpret_cast<void*>(index_q_out->data_ptr())                   \
+          : nullptr,                                                           \
      reinterpret_cast<st const*>(q_norm_weight.data_ptr()),                   \
      reinterpret_cast<st const*>(k_norm_weight.data_ptr()),                   \
      has_index ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr()) \
@@ -502,11 +581,11 @@ void launchFusedMiniMaxM3(scalar_t* qkv, scalar_t* q_out, scalar_t* index_q_out,
                : nullptr,                                                     \
      insert_kv ? reinterpret_cast<CACHE_T*>(kv_cache->data_ptr()) : nullptr,  \
      (insert_kv && has_index)                                                 \
-          ? reinterpret_cast<st*>(index_cache->data_ptr())                     \
+          ? reinterpret_cast<void*>(index_cache->data_ptr())                   \
          : nullptr,                                                           \
      static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens, nq,   \
      nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv, kv_s_token, \
-      kv_s_head, has_index, insert_kv, stream)
+      kv_s_head, has_index, insert_kv, fp8_idx, stream)

 // ────────────────────────────────────────────────────────────────────────────
 // Torch op wrapper
@@ -612,6 +691,7 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
                 torch::headeronly::ScalarType::Long &&
             index_slot_mapping->numel() == slot_mapping->numel()),
        "index_slot_mapping must be int64 CUDA with slot_mapping length");
+    // Main attention KV cache: auto matches qkv, fp8 uses uint8 storage.
    if (kv_dt == vllm::Fp8KVCacheDataType::kAuto) {
      STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
                      "auto kv_cache dtype must match qkv");
@@ -620,9 +700,13 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
          kv_cache->scalar_type() == torch::headeronly::ScalarType::Byte,
          "fp8 kv_cache must use uint8 storage");
    }
-    STD_TORCH_CHECK(index_cache.has_value() &&
-                        index_cache->scalar_type() == qkv.scalar_type(),
-                    "insert mode requires matching index_cache");
+    // Indexer index-K cache: independent dtype -- qkv dtype or fp8 e4m3.
+    STD_TORCH_CHECK(
+        index_cache.has_value() &&
+            (index_cache->scalar_type() == qkv.scalar_type() ||
+             index_cache->scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e4m3fn),
+        "insert mode requires index_cache matching qkv dtype or fp8 e4m3");
    STD_TORCH_CHECK(kv_cache->dim() == 5 && kv_cache->stride(4) == 1,
                    "kv_cache must be [nb,2,bs,nkv,head_dim] with contiguous "
                    "head_dim (stride(4)==1)");
@@ -652,14 +736,31 @@ void fused_minimax_m3_qknorm_rope_kv_insert(
        "index_q_out requires the index branch (num_index_heads > 0)");
    STD_TORCH_CHECK(
        index_q_out->is_cuda() && index_q_out->is_contiguous() &&
-            index_q_out->scalar_type() == qkv.scalar_type(),
-        "index_q_out must be a contiguous CUDA tensor matching qkv dtype");
+            (index_q_out->scalar_type() == qkv.scalar_type() ||
+             index_q_out->scalar_type() ==
+                 torch::headeronly::ScalarType::Float8_e4m3fn),
+        "index_q_out must be contiguous CUDA, qkv dtype or fp8 e4m3");
    STD_TORCH_CHECK(index_q_out->numel() ==
                        static_cast<int64_t>(num_tokens) * niq * kHeadDim,
                    "index_q_out must have num_tokens * num_index_heads * 128 "
                    "elements");
  }

+  // fp8 index path: the index-K cache and index-Q outputs are e4m3 bytes while
+  // q/k/v + q_out stay qkv dtype. Both index outputs must agree.
+  auto const kFp8 = torch::headeronly::ScalarType::Float8_e4m3fn;
+  bool const fp8_idx =
+      (index_cache.has_value() && index_cache->scalar_type() == kFp8) ||
+      (index_q_out.has_value() && index_q_out->scalar_type() == kFp8);
+  if (fp8_idx) {
+    STD_TORCH_CHECK(
+        !index_cache.has_value() || index_cache->scalar_type() == kFp8,
+        "fp8 index path: index_cache must be fp8 e4m3");
+    STD_TORCH_CHECK(
+        !index_q_out.has_value() || index_q_out->scalar_type() == kFp8,
+        "fp8 index path: index_q_out must be fp8 e4m3");
+  }
+
  const torch::stable::accelerator::DeviceGuard device_guard(
      qkv.get_device_index());
  auto stream = get_current_cuda_stream(qkv.get_device_index());
@@ -286,3 +286,52 @@ template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>(

 template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>(
    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 6144>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
@@ -41,6 +41,7 @@ inline int getSMVersion() {
 static constexpr int DEFAULT_NUM_EXPERTS = 256;
 static constexpr int KIMI_K2_NUM_EXPERTS = 384;
 static constexpr int DEFAULT_HIDDEN_DIM = 7168;
+static constexpr int GLM_5_HIDDEN_DIM = 6144;

 template <typename T, int kNumTokens, int kNumExperts, int kHiddenDim>
 void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b,
@@ -121,14 +122,21 @@ void dsv3_router_gemm(

  STD_TORCH_CHECK(mat_a.size(1) == mat_b.size(1),
                  "mat_a and mat_b must have the same hidden_dim");
-  STD_TORCH_CHECK(hidden_dim == DEFAULT_HIDDEN_DIM,
-                  "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
-                  ", but got hidden_dim=", hidden_dim);
+  STD_TORCH_CHECK(
+      hidden_dim == DEFAULT_HIDDEN_DIM || hidden_dim == GLM_5_HIDDEN_DIM,
+      "Expected hidden_dim=", DEFAULT_HIDDEN_DIM,
+      " or hidden_dim=", GLM_5_HIDDEN_DIM, ", but got hidden_dim=", hidden_dim);
  STD_TORCH_CHECK(
      num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS,
      "Expected num_experts=", DEFAULT_NUM_EXPERTS,
      " or num_experts=", KIMI_K2_NUM_EXPERTS,
      ", but got num_experts=", num_experts);
+  // KIMI_K2_NUM_EXPERTS is only instantiated for the default hidden_dim.
+  STD_TORCH_CHECK(
+      hidden_dim == DEFAULT_HIDDEN_DIM || num_experts == DEFAULT_NUM_EXPERTS,
+      "hidden_dim=", GLM_5_HIDDEN_DIM,
+      " only supports num_experts=", DEFAULT_NUM_EXPERTS,
+      ", but got num_experts=", num_experts);
  STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
                  "currently num_tokens must be less than or equal to 16 for "
                  "router_gemm");
@@ -148,35 +156,49 @@ void dsv3_router_gemm(

  const cudaStream_t stream = get_current_cuda_stream(mat_a.get_device_index());

+  __nv_bfloat16 const* a_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
+  __nv_bfloat16 const* b_ptr =
+      reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr());
+
  if (output.scalar_type() == torch::headeronly::ScalarType::Float) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_float_output(
-              num_tokens, reinterpret_cast<float*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
+    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
+      if (num_experts == DEFAULT_NUM_EXPERTS) {
+        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
+                                                              out_ptr, a_ptr,
+                                                              b_ptr, stream);
+      } else {
+        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_float_output(num_tokens,
+                                                              out_ptr, a_ptr,
+                                                              b_ptr, stream);
+      }
+    } else {  // GLM_5_HIDDEN_DIM
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                   GLM_5_HIDDEN_DIM>::unroll_float_output(num_tokens, out_ptr,
+                                                          a_ptr, b_ptr, stream);
    }
  } else if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    if (num_experts == DEFAULT_NUM_EXPERTS) {
-      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
-    } else if (num_experts == KIMI_K2_NUM_EXPERTS) {
-      LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::
-          unroll_bf16_output(
-              num_tokens,
-              reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
-              reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), stream);
+    __nv_bfloat16* out_ptr =
+        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr());
+    if (hidden_dim == DEFAULT_HIDDEN_DIM) {
+      if (num_experts == DEFAULT_NUM_EXPERTS) {
+        LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
+                                                             out_ptr, a_ptr,
+                                                             b_ptr, stream);
+      } else {
+        LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS,
+                     DEFAULT_HIDDEN_DIM>::unroll_bf16_output(num_tokens,
+                                                             out_ptr, a_ptr,
+                                                             b_ptr, stream);
+      }
+    } else {  // GLM_5_HIDDEN_DIM
+      LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS,
+                   GLM_5_HIDDEN_DIM>::unroll_bf16_output(num_tokens, out_ptr,
+                                                         a_ptr, b_ptr, stream);
    }
  }
 }
@@ -286,3 +286,52 @@ template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>(

 template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
@@ -343,6 +343,14 @@ void persistent_topk(const torch::stable::Tensor& logits,
                     torch::stable::Tensor& workspace, int64_t k,
                     int64_t max_seq_len);

+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+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);
+#endif
+
 void selective_scan_fwd(
    const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
    const torch::stable::Tensor& A, const torch::stable::Tensor& B,
@@ -11,6 +11,8 @@
 #include <cub/cub.cuh>
 #include <cstdint>

+#include "topk_histogram_4096.cuh"
+
 namespace vllm {
 namespace persistent {

@@ -935,8 +937,16 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
 }  // namespace persistent

 // ============================================================================
-// FlashInfer FilteredTopK (BS>32 dispatch) — float32 only.
-// Extracted from flashinfer_topk.cuh. Lives in namespace vllm (not persistent).
+// ============================================================================
+// Optimized FilteredTopK — single CTA per row for bs > 32.
+// Kept with persistent_topk so the portable fallback owns the non-cluster path.
+// ============================================================================
+namespace filtered_topk {
+
+namespace hist4096 = topk_histogram_4096;
+
+// ============================================================================
+// FilteredTopK — single CTA per row for bs > 32
 // Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
 // ============================================================================

@@ -963,13 +973,6 @@ struct vec_t {
      data[i] = ptr[i];
    }
  }
-
-  FLASHINFER_INLINE void cast_store(T* ptr) const {
-#pragma unroll
-    for (size_t i = 0; i < N; ++i) {
-      ptr[i] = data[i];
-    }
-  }
 };
 #undef FLASHINFER_INLINE

@@ -1013,7 +1016,8 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
 * \tparam IdType Index type (int32_t)
 * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
 */
-template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
+template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048,
+          bool UsePredicatedShortLoads = false>
 __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
    FilteredTopKUnifiedKernel(const DType* __restrict__ input,
                              IdType* __restrict__ output,
@@ -1042,6 +1046,19 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
    return;
  }

+  // Short path
+  if (length <= 32768) {
+    extern __shared__ uint8_t _smem_reg[];
+    if constexpr (UsePredicatedShortLoads) {
+      hist4096::histogram_4096_topk_predicated<MAX_K, 12, 8>(score, dst, length,
+                                                             _smem_reg);
+    } else {
+      hist4096::histogram_4096_topk<MAX_K, 12, 8>(score, dst, length,
+                                                  _smem_reg);
+    }
+    return;
+  }
+
  // Static shared memory
  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
  alignas(128) __shared__ int s_counter;
@@ -1285,14 +1302,15 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,

  const int vec_size = ComputeFilteredTopKVecSize<DType>(max_len);

-#define DISPATCH_VEC_SIZE(VS)                                               \
-  if (vec_size == VS) {                                                     \
-    auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>;      \
-    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                              \
-        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));   \
-    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
-                                          smem_size, stream));              \
-    return cudaSuccess;                                                     \
+#define DISPATCH_VEC_SIZE(VS)                                                 \
+  if (vec_size == VS) {                                                       \
+    auto kernel =                                                             \
+        FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K, (VS != MAX_VEC)>; \
+    FLASHINFER_CUDA_CALL(cudaFuncSetAttribute(                                \
+        kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));     \
+    FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args,   \
+                                          smem_size, stream));                \
+    return cudaSuccess;                                                       \
  }

  DISPATCH_VEC_SIZE(1)
@@ -1306,6 +1324,19 @@ cudaError_t FilteredTopKRaggedTransform(const DType* input,
  return cudaSuccess;
 }

+}  // namespace filtered_topk
+
+template <typename DType, typename IdType, uint32_t MAX_K = 2048>
+cudaError_t FilteredTopKRaggedTransform(const DType* input,
+                                        IdType* output_indices,
+                                        const IdType* lengths,
+                                        uint32_t num_rows, uint32_t top_k_val,
+                                        uint32_t max_len,
+                                        cudaStream_t stream = 0) {
+  return filtered_topk::FilteredTopKRaggedTransform<DType, IdType, MAX_K>(
+      input, output_indices, lengths, num_rows, top_k_val, max_len, stream);
+}
+
 }  // namespace vllm

 #endif  // PERSISTENT_TOPK_CUH_
@@ -180,7 +180,7 @@ bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {

 #if defined CUDA_VERSION
  #if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
-  if (cuda_device_capability >= 100 && cuda_device_capability < 110) {
+  if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
    return CUDA_VERSION >= 12080;
  }
  #endif
@@ -0,0 +1,563 @@
+/*
+ * Shared 4096-bin single-CTA TopK helpers.
+ */
+
+#ifndef TOPK_HISTOGRAM_4096_CUH_
+#define TOPK_HISTOGRAM_4096_CUH_
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace vllm {
+namespace topk_histogram_4096 {
+
+constexpr uint32_t kBlockSize = 1024;
+constexpr uint32_t RADIX = 256;
+constexpr uint32_t kMaxTies = 1024;
+static_assert(kMaxTies <= kBlockSize,
+              "tie_handle requires kMaxTies <= kBlockSize");
+constexpr uint32_t kWarpSize = 32;
+constexpr uint32_t kNumWarps = kBlockSize / kWarpSize;
+
+// Register path
+constexpr uint32_t kHist4096VecsPerThread = 4;
+constexpr uint32_t kHist4096MaxLen =
+    kHist4096VecsPerThread * 4 * kBlockSize;  // 16384
+
+struct alignas(16) MatchBin {
+  uint32_t bin, above_count, equal_count;
+};
+struct alignas(8) Tie {
+  uint32_t idx;
+  float score;
+};
+
+__device__ __forceinline__ void load_float4_predicated(const float* ptr,
+                                                       int base, int seq_len,
+                                                       float& v0, float& v1,
+                                                       float& v2, float& v3) {
+  uint32_t r0, r1, r2, r3;
+  const int p0 = (base < seq_len);
+  const int p1 = (base + 1 < seq_len);
+  const int p2 = (base + 2 < seq_len);
+  const int p3 = (base + 3 < seq_len);
+  asm volatile(
+      "{\n"
+      "  .reg .pred pr0, pr1, pr2, pr3;\n"
+      "  setp.ne.u32 pr0, %4, 0;\n"
+      "  setp.ne.u32 pr1, %5, 0;\n"
+      "  setp.ne.u32 pr2, %6, 0;\n"
+      "  setp.ne.u32 pr3, %7, 0;\n"
+      "  mov.u32 %0, 0xFF800000;\n"
+      "  mov.u32 %1, 0xFF800000;\n"
+      "  mov.u32 %2, 0xFF800000;\n"
+      "  mov.u32 %3, 0xFF800000;\n"
+      "  @pr0 ld.global.cg.u32 %0, [%8];\n"
+      "  @pr1 ld.global.cg.u32 %1, [%8+4];\n"
+      "  @pr2 ld.global.cg.u32 %2, [%8+8];\n"
+      "  @pr3 ld.global.cg.u32 %3, [%8+12];\n"
+      "}\n"
+      : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3)
+      : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr));
+  v0 = __uint_as_float(r0);
+  v1 = __uint_as_float(r1);
+  v2 = __uint_as_float(r2);
+  v3 = __uint_as_float(r3);
+}
+
+// converts the float32 score to a 32-bit ordered unsigned integer — the full
+// precision key for radix sorting
+__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits ->
+// bin 0-4095)
+template <uint32_t kBits>
+__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits)
+                                 : static_cast<uint16_t>(bits | 0x8000);
+  return key >> (16 - kBits);
+}
+
+// running sum within each warp — thread 0 gets its own value, thread 1 gets
+// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc.
+__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane,
+                                                       uint32_t v) {
+#pragma unroll
+  for (uint32_t o = 1; o < 32; o *= 2) {
+    uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o);
+    if (lane >= o) v += n;
+  }
+  return v;
+}
+
+// Returns the sum of a value across all 32 threads in the warp, and every
+// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX
+// instruction for hardware warp-wide reduction. Older targets use the
+// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32
+// lanes).
+__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
+  uint32_t r;
+  asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v));
+  return r;
+#else
+  #pragma unroll
+  for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) {
+    v += __shfl_xor_sync(0xFFFFFFFF, v, mask);
+  }
+  return v;
+#endif
+}
+
+// ============================================================================
+// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered
+// key Each round narrows by 8 bits until ties are fully resolved
+// ============================================================================
+
+template <uint32_t TopK>
+__device__ void tie_handle(const Tie* ties, uint32_t num_ties,
+                           uint32_t num_above, int32_t* output, void* _smem) {
+  struct TS {
+    alignas(128) uint32_t counter;
+    alignas(128) MatchBin match;
+    uint32_t histogram[RADIX];
+    uint32_t warp_sum[kNumWarps];
+  };
+  auto* s = static_cast<TS*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto li = tx % kWarpSize, wi = tx / kWarpSize;
+
+  // Each thread loads one tie element.
+  const bool has = tx < num_ties;
+  const auto tie = has ? ties[tx] : Tie{0, 0.0f};
+  const uint32_t key = convert_to_uint32_v2(tie.score);
+
+  bool active = has;  // tracks whether this thread's tie is still a candidate.
+  uint32_t remain =
+      TopK - num_above;  // decreases each round as ties are resolved.
+  uint32_t wpos = TopK;  // wpos will hold the final output position.
+  s->counter = 0;
+  __syncthreads();
+
+  // The 4-round radix loop - each round narrows by 8 bits until ties are fully
+  // resolved
+#pragma unroll
+  for (int r = 0; r < 4; r++) {
+    uint32_t sh = 24 - r * 8;  // round 0: bits 31-24, round 1: 23-16, etc.
+    uint32_t bin = (key >> sh) & 0xFF;  // this tie's 8-bit bin for this round
+
+    // Step 1: Build 256-bin histogram.
+    if (tx < RADIX) s->histogram[tx] = 0;
+    __syncthreads();
+    if (active) atomicAdd(&s->histogram[bin], 1);
+    __syncthreads();
+
+    // Step 2: Prefix scan to find threshold
+    uint32_t hv = 0, wi2 = 0;
+    if (tx < RADIX) {
+      hv = s->histogram[tx];
+      wi2 = warp_inclusive_sum(li, hv);
+      if (li == kWarpSize - 1) s->warp_sum[wi] = wi2;
+    }
+    __syncthreads();
+
+    if (tx < RADIX) {
+      auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
+      auto tot = warp_reduce_sum_full(tmp);
+      auto inter = warp_reduce_sum_full(li < wi ? tmp : 0);
+      auto above = tot - (inter + wi2);
+      if (above < remain && above + hv >= remain) {
+        s->match = {tx, above, remain - above};
+      }
+    }
+    __syncthreads();
+
+    // Step 3: Scatter
+    auto [thr, na, _] = s->match;  // threshold bin, num above, unused
+    if (active) {
+      if (bin > thr) {
+        wpos = num_above +
+               atomicAdd(&s->counter, 1);  // above -> place in output directly
+        active = false;
+      } else if (bin < thr)
+        active = false;  // below -> discard
+      else if (r == 3)
+        wpos = TopK - atomicAdd(&s->match.equal_count,
+                                -1u);  // last round: place remaining
+    }
+    remain -= na;
+    if (!remain) break;  // all ties resolved early
+  }
+  // Final write
+  if (wpos < TopK) output[wpos] = tie.idx;
+}
+
+// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048).
+// tie_handle assumes 1 tie per thread (max 1024).
+// This version handles 2 ties per thread via kPerThread=2
+template <uint32_t TopK>
+__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties,
+                                 uint32_t num_above, int32_t* output,
+                                 void* _smem) {
+  static_assert(TopK > kBlockSize);
+  struct TS {
+    alignas(128) uint32_t counter;
+    alignas(128) MatchBin match;
+    uint32_t histogram[RADIX];
+    uint32_t warp_sum[kNumWarps];
+  };
+  auto* s = static_cast<TS*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto li = tx % kWarpSize;
+  const auto wi = tx / kWarpSize;
+
+  constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize;
+  Tie my_ties[kPerThread];
+  uint32_t keys[kPerThread];
+  bool active[kPerThread];
+
+  for (uint32_t e = 0; e < kPerThread; e++) {
+    uint32_t idx = e * kBlockSize + tx;
+    if (idx < num_ties) {
+      my_ties[e] = ties[idx];
+      keys[e] = convert_to_uint32_v2(ties[idx].score);
+      active[e] = true;
+    } else {
+      my_ties[e] = {0, 0.0f};
+      keys[e] = 0;
+      active[e] = false;
+    }
+  }
+
+  uint32_t remain = TopK - num_above;
+  s->counter = 0;
+  __syncthreads();
+
+  for (int r = 0; r < 4; r++) {
+    uint32_t sh = 24 - r * 8;
+    if (tx < RADIX) {
+      s->histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (uint32_t e = 0; e < kPerThread; e++) {
+      if (active[e]) {
+        atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1);
+      }
+    }
+    __syncthreads();
+
+    uint32_t hv = 0;
+    if (tx < RADIX) {
+      hv = s->histogram[tx];
+      auto wi2 = warp_inclusive_sum(li, hv);
+      if (li == kWarpSize - 1) {
+        s->warp_sum[wi] = wi2;
+      }
+    }
+    __syncthreads();
+    if (tx < RADIX) {
+      auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0;
+      auto total = warp_reduce_sum_full(tmp2);
+      auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0);
+      auto wi2 = warp_inclusive_sum(li, hv);
+      auto above = total - (inter + wi2);
+      if (above < remain && above + hv >= remain) {
+        s->match = {
+            .bin = tx, .above_count = above, .equal_count = remain - above};
+      }
+    }
+    __syncthreads();
+
+    auto thr = s->match.bin;
+    auto na = s->match.above_count;
+
+    for (uint32_t e = 0; e < kPerThread; e++) {
+      if (!active[e]) {
+        continue;
+      }
+      uint32_t bin = (keys[e] >> sh) & 0xFF;
+      if (bin > thr) {
+        uint32_t wpos = num_above + atomicAdd(&s->counter, 1);
+        if (wpos < TopK) {
+          output[wpos] = my_ties[e].idx;
+        }
+        active[e] = false;
+      } else if (bin < thr) {
+        active[e] = false;
+      } else if (r == 3) {
+        uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u);
+        if (wpos < TopK) {
+          output[wpos] = my_ties[e].idx;
+        }
+      }
+    }
+
+    num_above += na;
+    remain -= na;
+    __syncthreads();
+    s->counter = 0;
+    __syncthreads();
+  }
+}
+
+// ============================================================================
+// Register-based single-CTA fast path for seq_len <= 16384
+// 4 float4 per thread × 1024 threads = 16384 elements max
+// Uses 4096-bin (12-bit) histogram for better precision
+// ============================================================================
+
+template <uint32_t TopK, uint32_t HIST_BITS>
+struct Histogram4096Smem {
+  static constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
+  static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies;
+  alignas(128) uint32_t counter_gt;
+  alignas(128) uint32_t counter_eq;
+  MatchBin match;
+  uint32_t warp_sum[kNumWarps];
+  union {
+    uint32_t histogram[HIST_BINS];
+    Tie tie_buffer[TIE_CAPACITY];
+  };
+};
+
+template <uint32_t TopK, uint32_t HIST_BITS,
+          uint32_t VECS_PER_THREAD = kHist4096VecsPerThread,
+          bool UsePredicatedLoads = false>
+__device__ void histogram_4096_topk(const float* __restrict__ scores,
+                                    int32_t* __restrict__ output,
+                                    uint32_t length, void* _smem) {
+  constexpr uint32_t HIST_BINS = 1 << HIST_BITS;
+  constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize;
+  static_assert(HIST_BINS >= kBlockSize,
+                "HIST_BITS must give >= kBlockSize bins");
+
+  using Smem = Histogram4096Smem<TopK, HIST_BITS>;
+  auto* smem = static_cast<Smem*>(_smem);
+  const auto tx = threadIdx.x;
+  const auto lane_id = tx % kWarpSize;
+  const auto warp_id = tx / kWarpSize;
+
+  // Phase 1: Load all data into RF + build histogram
+  float4
+      vecs[VECS_PER_THREAD];  // 4 vectors x 4 floats = 16 elements per thread
+  if constexpr (ITEMS_PER_THREAD >= 4) {
+    // Zero the histogram (SMEM writes)
+    for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++)
+      reinterpret_cast<uint4*>(
+          smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] =
+          make_uint4(0, 0, 0, 0);
+  } else {
+    if (tx < HIST_BINS) smem->histogram[tx] = 0;
+  }
+  if (tx == 0) {
+    smem->counter_gt = 0;
+    smem->counter_eq = 0;
+  }
+  if constexpr (UsePredicatedLoads) {
+    const bool row_aligned = (reinterpret_cast<uintptr_t>(scores) & 0xFu) == 0;
+#pragma unroll
+    for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
+      const uint32_t base = (tx + v * kBlockSize) * 4;
+      if (base < length) {
+        if (row_aligned && base + 3 < length) {
+          vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+        } else {
+          load_float4_predicated(scores + base, static_cast<int>(base),
+                                 static_cast<int>(length), vecs[v].x, vecs[v].y,
+                                 vecs[v].z, vecs[v].w);
+        }
+      }
+    }
+  } else {
+#pragma unroll
+    for (uint32_t v = 0; v < VECS_PER_THREAD; v++) {
+      const uint32_t base = (tx + v * kBlockSize) * 4;
+      if (base < length) {
+        vecs[v] = *reinterpret_cast<const float4*>(scores + base);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Build histogram from RF via atomic adds into the shared histogram
+  bool done = false;
+#pragma unroll
+  for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
+    const float* elems = reinterpret_cast<const float*>(&vecs[v]);
+#pragma unroll
+    for (uint32_t e = 0; e < 4 && !done; e++) {
+      const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
+      if (idx >= length) {
+        done = true;
+      } else {
+        atomicAdd(&smem->histogram[extract_coarse_bin_N<HIST_BITS>(elems[e])],
+                  1);
+      }
+    }
+  }
+  __syncthreads();
+
+  // Phase 2: Prefix scan to find threshold bin
+  // Multi-element scan (4096 bins: 4 per thread)
+  uint32_t orig[ITEMS_PER_THREAD];
+  uint32_t local_sum = 0;
+
+  // Step 1: Each thread sums its 4 bins
+#pragma unroll
+  for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
+    orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i];
+    local_sum += orig[i];
+  }
+
+  // Step 2: Warp-level inclusive prefix sum on local_sum
+  const auto warp_inc = warp_inclusive_sum(lane_id, local_sum);
+  if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc;
+  __syncthreads();
+
+  // Step 3: Inter-warp prefix across warp sums.
+  const auto tmp = smem->warp_sum[lane_id];
+  uint32_t prefix = warp_reduce_sum_full(
+      lane_id < warp_id ? tmp : 0);  // sum of all prior warps
+  prefix +=
+      warp_inc - local_sum;  // exclusive prefix within this thread's position
+
+  // Step 4: Find threshold - scan 4 bins, accumulate prefix
+#pragma unroll
+  for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) {
+    prefix += orig[i];
+    const auto above = length - prefix;  // elements in bins ABOVE this one
+    if (above < TopK && above + orig[i] >= TopK) {
+      smem->match = {.bin = tx * ITEMS_PER_THREAD + i,
+                     .above_count = above,
+                     .equal_count = orig[i]};
+    }
+  }
+
+  __syncthreads();
+
+  // Phase 3: Scatter from registers
+  const auto [thr_bin, num_above, num_equal] = smem->match;
+  const bool need_tie = (num_equal + num_above > TopK);
+
+  done = false;
+#pragma unroll
+  for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) {
+    const float* elems = reinterpret_cast<const float*>(&vecs[v]);
+#pragma unroll
+    for (uint32_t e = 0; e < 4 && !done; e++) {
+      const uint32_t idx = (tx + v * kBlockSize) * 4 + e;
+      if (idx >= length) {
+        done = true;
+      } else {
+        const uint32_t bin = extract_coarse_bin_N<HIST_BITS>(elems[e]);
+        if (bin > thr_bin) {
+          output[atomicAdd(&smem->counter_gt, 1)] =
+              idx;  // above -> output directly
+        } else if (bin == thr_bin) {
+          const auto pos = atomicAdd(&smem->counter_eq, 1);
+          if (!need_tie) {
+            if (pos + num_above < TopK) {
+              output[pos + num_above] = idx;  // all fit
+            }
+          } else {
+            if (pos < TopK) {
+              smem->tie_buffer[pos] = {idx, elems[e]};  // store for refirement
+            }
+          }
+        }
+        // else: bin < thr_bin - discard (not in top-k)
+      }
+    }
+  }
+
+  // Phase 4: Tie-breaking
+  if (!need_tie) return;
+  __syncthreads();
+
+  // Fast warp-ballot tie-breaking for small tie counts
+  const uint32_t num_ties = min(num_equal, static_cast<uint32_t>(TopK));
+  const uint32_t topk_remain =
+      TopK - num_above;  // pick exactly remaining elements to fill topK
+
+  auto is_greater = [](const Tie& a, const Tie& b) {
+    return (a.score > b.score) || (a.score == b.score && a.idx < b.idx);
+  };
+
+  if (num_ties <= kWarpSize) {
+    // <=32 ties - Use warp ballot
+    // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024
+    // comparisons in one instruction per warp. O(1) work.
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    if (lane_id >= num_ties || warp_id >= num_ties) return;
+    const uint32_t mask = (1ull << num_ties) - 1u;
+    const auto tie = smem->tie_buffer[lane_id];  // each lane holds one tie
+    const auto target =
+        smem->tie_buffer[warp_id];  // each warp evaluates one candidate
+    const bool pred =
+        is_greater(tie, target);  // compare all ties against target
+    const auto rank = static_cast<uint32_t>(
+        __popc(__ballot_sync(mask, pred)));  // count how many are greater
+    if (lane_id == 0 && rank < topk_remain) {
+      output[num_above + rank] = target.idx;  // place at correct position
+    }
+  } else if (num_ties <=
+             kWarpSize *
+                 2) {  // TODO (roberto): try to refactor this with <=32 case
+    //  Same idea but each thread handles 2 tie elements
+    const auto lane_id = tx % kWarpSize;
+    const auto warp_id = tx / kWarpSize;
+    const auto lane1 = lane_id + kWarpSize;
+    const auto warp1 = warp_id + kWarpSize;
+    const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__};
+    const auto tie0 = smem->tie_buffer[lane_id];
+    const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid;
+    if (warp_id < num_ties) {
+      const auto target = smem->tie_buffer[warp_id];
+      const auto r0 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
+      const auto r1 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
+      if (lane_id == 0 && r0 + r1 < topk_remain)
+        output[num_above + r0 + r1] = target.idx;
+    }
+    if (warp1 < num_ties) {
+      const auto target = smem->tie_buffer[warp1];
+      const auto r0 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target)));
+      const auto r1 =
+          __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target)));
+      if (lane_id == 0 && r0 + r1 < topk_remain)
+        output[num_above + r0 + r1] = target.idx;
+    }
+  } else {
+    // Large tie count: fall back to 4-round radix-256 sort
+    if constexpr (TopK <= kBlockSize) {
+      tie_handle<TopK>(smem->tie_buffer, num_ties, num_above, output, smem);
+    } else {
+      tie_handle_large<TopK>(smem->tie_buffer, num_ties, num_above, output,
+                             smem);
+    }
+  }
+}
+
+template <uint32_t TopK, uint32_t HIST_BITS,
+          uint32_t VECS_PER_THREAD = kHist4096VecsPerThread>
+__device__ __noinline__ void histogram_4096_topk_predicated(
+    const float* __restrict__ scores, int32_t* __restrict__ output,
+    uint32_t length, void* _smem) {
+  histogram_4096_topk<TopK, HIST_BITS, VECS_PER_THREAD, true>(scores, output,
+                                                              length, _smem);
+}
+
+}  // namespace topk_histogram_4096
+}  // namespace vllm
+
+#endif  // TOPK_HISTOGRAM_4096_CUH_
@@ -493,6 +493,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, "
      "Tensor workspace, int k, int max_seq_len) -> ()");

+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+  ops.def(
+      "cooperative_topk(Tensor logits, Tensor lengths, Tensor! output, "
+      "Tensor workspace, int k, int max_seq_len) -> ()");
+#endif
+
  // Activation ops
  ops.def(
      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
@@ -711,6 +717,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill));
  ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode));
  ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));
+#ifdef VLLM_ENABLE_COOPERATIVE_TOPK
+  ops.impl("cooperative_topk", TORCH_BOX(&cooperative_topk));
+#endif

  // Activation kernels (shared CUDA/ROCm)
  ops.impl("persistent_masked_m_silu_mul_quant",
@@ -338,7 +338,7 @@ vllm bench serve \
    --model meta-llama/Meta-Llama-3-8B-Instruct \
    --dataset-name spec_bench \
    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/spec_bench/question.jsonl" \
-    --num-prompts -1
+    --num-prompts -1 \
    --spec-bench-category "summarization"
 ```

@@ -352,7 +352,7 @@ vllm bench serve \
 First, download the dataset to a folder, using this one liner:

 ```bash
-curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
+curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -
 ```

 The command supports also the following arguments:
@@ -388,7 +388,7 @@ vllm bench serve \
    --model meta-llama/Llama-3.3-70B-Instruct \
    --dataset-name speed_bench \
    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
-    --num-prompts -1
+    --num-prompts -1 \
    --speed-bench-category "multilingual"
 ```

@@ -398,7 +398,7 @@ Run all categories in the Throughput split (2k ISL):
 vllm bench serve \
    --model meta-llama/Llama-3.3-70B-Instruct \
    --dataset-name speed_bench \
-    --speed-bench-dataset-subset throughput_2k
+    --speed-bench-dataset-subset throughput_2k \
    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
    --num-prompts -1
 ```
@@ -109,7 +109,7 @@ from vllm import LLM

 # Combine pipeline and tensor parallelism
 llm = LLM(
-    model="meta-llama/Llama-3.3-70B-Instruct,
+    model="meta-llama/Llama-3.3-70B-Instruct",
    tensor_parallel_size=4,
    pipeline_parallel_size=2,
 )
@@ -133,16 +133,6 @@ Priority is **1 = highest** (tried first).
 | 7 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 8 | `FLASHMLA_SPARSE` |

-**Ampere/Hopper (SM 8.x-9.x):**
-
-| Priority | Backend |
-| -------- | ------- |
-| 1 | `FLASH_ATTN_MLA` |
-| 2 | `FLASHMLA` |
-| 3 | `FLASHINFER_MLA` |
-| 4 | `TRITON_MLA` |
-| 5 | `FLASHMLA_SPARSE` |
-
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -170,7 +160,7 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 7.x-9.x |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 8.x-9.x |
 | `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
 | `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
@@ -231,7 +221,8 @@ MLA decode backends are selected using the standard
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
@@ -248,10 +239,11 @@ DeepSeek V4 sparse MLA uses its own decode backends, selected via
 `--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
 `FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
 pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
-default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
+default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and
+`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.

 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
-| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
-| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x |
+| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
@@ -161,11 +161,11 @@ class AttentionCGSupport(enum.Enum):
    ALWAYS = 3
    """CUDA Graphs always supported; supports mixed-prefill-decode"""
    UNIFORM_BATCH = 2
-    """CUDA Graphs supported for batches the only contain query lengths that are
+    """CUDA Graphs supported for batches that only contain query lengths that are
    the same, this can be used for spec-decode 
        i.e. "decodes" are 1 + num_speculative_tokens"""
    UNIFORM_SINGLE_TOKEN_DECODE = 1
-    """CUDA Graphs supported for batches the only contain query_len==1 decodes"""
+    """CUDA Graphs supported for batches that only contain query_len==1 decodes"""
    NEVER = 0
    """NO CUDA Graphs support"""
 ```
@@ -685,7 +685,7 @@ documentation for this option states:
 > use of possibly costly and or blocking operations and hence might
 > have a performance impact.

-The metrics were added by <https://github.com/vllm-project/vllm/pull/7089> and who up in an OpenTelemetry trace
+The metrics were added by <https://github.com/vllm-project/vllm/pull/7089> and show up in an OpenTelemetry trace
 as:

 ```text
@@ -27,7 +27,7 @@ In the example above, the KV cache in the first block can be uniquely identified
    For `vllm serve`, you can control the hashing algorithm via `--prefix-caching-hash-algo`:
    - `sha256` (default): Uses Python's `pickle` for serialization. Hashes may not be reproducible across different Python or vLLM versions.
    - `sha256_cbor`: Uses `cbor2` for serialization, providing a reproducible, cross-language compatible hash. This is recommended for deterministic caching across environments.
-    - `xxhash`: `Uses Pickle serialization with xxHash (128-bit) for faster, non-cryptographic hashing. Requires the optional `xxhash` package. IMPORTANT: Use of a hashing algorithm that is not considered cryptographically secure theoretically increases the risk of hash collisions, which can cause undefined behavior or even leak private information in multi-tenant environments. Even if collisions are still very unlikely, it is important to consider your security risk tolerance against the performance benefits before turning this on.
+    - `xxhash`: Uses Pickle serialization with xxHash (128-bit) for faster, non-cryptographic hashing. Requires the optional `xxhash` package. IMPORTANT: Use of a hashing algorithm that is not considered cryptographically secure theoretically increases the risk of hash collisions, which can cause undefined behavior or even leak private information in multi-tenant environments. Even if collisions are still very unlikely, it is important to consider your security risk tolerance against the performance benefits before turning this on.
    - `xxhash_cbor` combines canonical CBOR serialization with xxHash for reproducible hashing. Requires the optional `xxhash` package.    

 **A hashing example with multi-modality inputs**  
@@ -197,7 +197,7 @@ As can be seen, block 3 is a new full block and is cached. However, it is redund

 When a request is finished, we free all its blocks if no other requests are using them (reference count = 0). In this example, we free request 1 and block 2, 3, 4, 8 associated with it. We can see that the freed blocks are added to the tail of the free queue in the *reverse* order. This is because the last block of a request must hash more tokens and is less likely to be reused by other requests. As a result, it should be evicted first.

-![Free queue after a request us freed](../assets/design/prefix_caching/free.png)
+![Free queue after a request is freed](../assets/design/prefix_caching/free.png)

 ### Eviction (LRU)

@@ -55,7 +55,7 @@ Here is an example of how to quantize `meta-llama/Llama-3.2-1B-Instruct`:

 ## Running a quantized model with vLLM

-To run an GPTQModel quantized model with vLLM, you can use [DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2](https://huggingface.co/ModelCloud/DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2) with the following command:
+To run a GPTQModel quantized model with vLLM, you can use [DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2](https://huggingface.co/ModelCloud/DeepSeek-R1-Distill-Qwen-7B-gptqmodel-4bit-vortex-v2) with the following command:

 ```bash
 python examples/deployment/llm_engine_example.py \
@@ -78,8 +78,6 @@ def tokenize(sample):
 ds = ds.map(tokenize, remove_columns=ds.column_names)
 ```

-</details>
-
 ### 3. Applying Quantization

 Now, apply the quantization algorithms:
@@ -190,7 +190,7 @@ For mitigation strategies, please refer to the FAQ entry *Can the output of a pr

 ## Known Feature Incompatibility

-1. Pipeline parallelism is not composible with speculative decoding as of `vllm<=0.15.0`
+1. Pipeline parallelism is not composable with speculative decoding as of `vllm<=0.15.0`
 2. Speculative decoding with a draft models is not supported in `vllm<=0.10.0`

 ## Resources for vLLM contributors
@@ -338,7 +338,7 @@ Supported models:

 Flags: `--tool-call-parser deepseek_v31 --chat-template {see_above}`

-### OpenAI OSS Models ('openai`)
+### OpenAI OSS Models (`openai`)

 Supported models:

@@ -143,7 +143,7 @@ enabling the corresponding APIs.

 The [classify][vllm.LLM.classify] method outputs a probability vector for each prompt.
 It is primarily designed for [classification models](classify.md).
-For more information about `LLM.embed`, see [this page](classify.md#offline-inference).
+For more information about `LLM.classify`, see [this page](classify.md#offline-inference).

 ### `LLM.embed`

@@ -302,7 +302,7 @@ Pooling models now support token-wise task.

 ### Score task

-`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
+`score` task has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
 equal to 1 can it be used as a scoring model and have its scoring API enabled.

 ### Pooling multitask support
@@ -440,7 +440,7 @@ More examples can be found here: [examples/pooling/score](../../../examples/pool

 ## Supported Features

-AS cross-encoder models are a subset of classification models that accept two prompts as input and output num_labels equal to 1, cross-encoder features should be consistent with (sequence) classification. For more information, see [this page](classify.md#supported-features).
+As cross-encoder models are a subset of classification models that accept two prompts as input and output num_labels equal to 1, cross-encoder features should be consistent with (sequence) classification. For more information, see [this page](classify.md#supported-features).

 ### Score Template

@@ -597,6 +597,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
 | `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
+| `MossAudioModel` | MOSS-Audio | T + A<sup>+</sup> | `OpenMOSS-Team/MOSS-Audio-4B-Instruct`, `OpenMOSS-Team/MOSS-Audio-4B-Thinking`, `OpenMOSS-Team/MOSS-Audio-8B-Instruct`, `OpenMOSS-Team/MOSS-Audio-8B-Thinking` | ✅︎ | ✅︎ |
 | `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
 | `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
 | `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
@@ -26,4 +26,4 @@ quack-kernels>=0.3.3
 tokenspeed-mla==0.1.2

 # Humming kernels for quantization gemm
-humming-kernels[cu13]==0.1.4
+humming-kernels[cu13]==0.1.6
@@ -365,6 +365,7 @@ impl MultimodalModelInfo {
        let processor = self.image_processor.raw;
        let images = image_frames.iter().map(|frame| frame.data().clone()).collect::<Vec<_>>();

+        // TODO: is it still necessary given that we've already in a dedicated runtime?
        tokio::task::spawn_blocking(move || {
            processor.preprocess(&images, &config).map_err(|error| multimodal!("{error}"))
        })
@@ -15,6 +15,7 @@ use crate::protocol::handshake::EngineCoreReadyResponse;
 use crate::protocol::lora::LoraRequest;
 use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
 use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
+use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
 use crate::transport::{self, ConnectedEngine};

 pub(crate) mod imp;
@@ -201,6 +202,8 @@ pub struct EngineCoreClient {
    coordinator: Option<CoordinatorHandle>,
    abort_tx: mpsc::UnboundedSender<AbortRequest>,

+    /// Runtime used to send messages to the engine and drive all background tasks.
+    runtime: BackgroundShutdownRuntime,
    // Background tasks
    output_task: AbortOnDropHandle<()>,
    dispatcher_task: AbortOnDropHandle<()>,
@@ -280,21 +283,22 @@ impl EngineCoreClient {
        let (output_tx, output_rx) = mpsc::channel(64);
        let (abort_tx, abort_rx) = mpsc::unbounded_channel();
        let engines = connected.engines;
+        let runtime = build_zmq_runtime();
        let inner = Arc::new(ClientInner::new(
            connected.input_send,
+            runtime.handle().clone(),
            config.model_name.clone(),
            &engines,
        ));
-        let output_task = AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(
+        let output_task = AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(
            connected.output_socket,
            output_tx,
        )));
-        let dispatcher_task = AbortOnDropHandle::new(tokio::spawn(run_output_dispatcher_loop(
-            inner.clone(),
-            output_rx,
-        )));
+        let dispatcher_task = AbortOnDropHandle::new(
+            runtime.spawn(run_output_dispatcher_loop(inner.clone(), output_rx)),
+        );
        let abort_task =
-            AbortOnDropHandle::new(tokio::spawn(run_abort_loop(inner.clone(), abort_rx)));
+            AbortOnDropHandle::new(runtime.spawn(run_abort_loop(inner.clone(), abort_rx)));

        // If any engine reported a dp_stats_address in its ready response, use it
        // as the external coordinator address.
@@ -307,13 +311,13 @@ impl EngineCoreClient {
                    CoordinatorHandle::new_inproc(coordinator_transport.input_socket);
                let (coordinator_output_tx, coordinator_output_rx) = mpsc::channel(64);
                let coordinator_output_task =
-                    AbortOnDropHandle::new(tokio::spawn(transport::run_output_loop(
+                    AbortOnDropHandle::new(runtime.spawn(transport::run_output_loop(
                        coordinator_transport.output_socket,
                        coordinator_output_tx,
                    )));
-                let coordinator_task = AbortOnDropHandle::new(tokio::spawn(
-                    runner.run(coordinator_output_rx, inner.clone()),
-                ));
+                let coordinator_task = AbortOnDropHandle::new(
+                    runtime.spawn(runner.run(coordinator_output_rx, inner.clone())),
+                );
                (
                    Some(handle),
                    Some(coordinator_output_task),
@@ -327,7 +331,7 @@ impl EngineCoreClient {
            {
                let (handle, service) = CoordinatorHandle::connect_external(address).await?;
                let coordinator_task =
-                    AbortOnDropHandle::new(tokio::spawn(service.run(inner.clone())));
+                    AbortOnDropHandle::new(runtime.spawn(service.run(inner.clone())));
                (Some(handle), None, Some(coordinator_task))
            } else {
                (None, None, None)
@@ -341,6 +345,7 @@ impl EngineCoreClient {
            inner,
            coordinator,
            abort_tx,
+            runtime,
            output_task,
            dispatcher_task,
            abort_task,
@@ -737,6 +742,7 @@ impl EngineCoreClient {
        let Self {
            inner,
            abort_tx,
+            runtime,
            output_task,
            dispatcher_task,
            abort_task,
@@ -757,6 +763,8 @@ impl EngineCoreClient {

        tasks.iter().for_each(|t| t.abort());
        join_all(tasks).await;
+        drop(inner);
+        drop(runtime);

        info!("engine-core client shut down");
        Ok(())
@@ -5,6 +5,7 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use arc_swap::ArcSwapOption;
 use parking_lot::Mutex;
 use thiserror_ext::AsReport as _;
+use tokio::runtime::Handle;
 use tokio::sync::mpsc;
 use tracing::{debug, info, trace, warn};
 use vllm_metrics::METRICS;
@@ -26,6 +27,8 @@ use crate::{Error, Result, transport};

 pub(crate) struct ClientInner {
    input_send: RouterSendHalf,
+    /// The runtime handle used for sending messages to the engine.
+    handle: Handle,
    model_name: String,
    request_reg: Mutex<RequestRegistry>,
    utility_reg: Mutex<UtilityRegistry>,
@@ -37,11 +40,13 @@ impl ClientInner {
    /// handshake completes.
    pub fn new(
        input_send: RouterSendHalf,
+        handle: Handle,
        model_name: String,
        engines: &[ConnectedEngine],
    ) -> Self {
        Self {
            input_send,
+            handle,
            model_name,
            request_reg: Mutex::new(RequestRegistry::new(engines)),
            utility_reg: Mutex::new(UtilityRegistry::default()),
@@ -213,9 +218,19 @@ impl ClientInner {
        // frames instead of always producing a single msgpack frame.
        let payload = encode_msgpack(payload)?;
        let mut input_send = self.input_send.clone();
-        transport::send_message(&mut input_send, engine_id, request_type.to_frame(), payload)
-            .await?;
-        Ok(())
+        let engine_id = engine_id.clone();
+
+        self.handle
+            .spawn(async move {
+                transport::send_message(
+                    &mut input_send,
+                    &engine_id,
+                    request_type.to_frame(),
+                    payload,
+                )
+                .await
+            })
+            .await?
    }

    /// Handle an abort request by sending the abort message to the engine.
@@ -434,6 +449,7 @@ mod tests {
        let (send, _) = socket.split();
        ClientInner::new(
            send,
+            Handle::current(),
            "test-model".to_string(),
            &[ConnectedEngine {
                engine_id: EngineId::from(b"engine-0"),
@@ -29,6 +29,8 @@ pub enum Error {
    Io(#[from] std::io::Error),
    #[error("transport error")]
    Transport(#[from] zeromq::ZmqError),
+    #[error("ZMQ runtime task failed")]
+    ZmqRuntimeTask(#[from] tokio::task::JoinError),
    #[error("engine core reported fatal failure")]
    EngineCoreDead,
    #[error("startup handshake timed out while waiting for {stage} after {timeout:?}")]
@@ -4,6 +4,7 @@ mod error;
 mod metrics;
 pub mod mock_engine;
 pub mod protocol;
+pub mod runtime;
 #[cfg(any(test, feature = "test-util"))]
 pub mod test_utils;
 mod transport;
@@ -0,0 +1,75 @@
+use std::mem::ManuallyDrop;
+use std::ops::{Deref, DerefMut};
+use std::sync::OnceLock;
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+use tokio::runtime::Runtime;
+
+/// A wrapper around [`Runtime`] that shuts down the runtime in the background when dropped.
+///
+/// This can be useful in some cases, because sometimes we want to drop the runtime without
+/// blocking the current thread, for example, when it's nested inside another runtime.
+pub struct BackgroundShutdownRuntime(ManuallyDrop<Runtime>);
+
+impl Drop for BackgroundShutdownRuntime {
+    fn drop(&mut self) {
+        // Safety: The runtime is only dropped once here.
+        let runtime = unsafe { ManuallyDrop::take(&mut self.0) };
+        runtime.shutdown_background();
+    }
+}
+
+impl Deref for BackgroundShutdownRuntime {
+    type Target = Runtime;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for BackgroundShutdownRuntime {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl From<Runtime> for BackgroundShutdownRuntime {
+    fn from(runtime: Runtime) -> Self {
+        Self(ManuallyDrop::new(runtime))
+    }
+}
+
+const ZMQ_WORKER_THREADS_ENV: &str = "VLLM_RS_ZMQ_WORKER_THREADS";
+/// The number of tasks running on the ZMQ runtime is fixed and expected to remain
+/// small, and multiple engines share the same ZMQ socket. Therefore, based on
+/// benchmarks, a default value of 4 is generally sufficient.
+const DEFAULT_ZMQ_WORKER_THREADS: usize = 4;
+
+static ZMQ_RUNTIME_SEQUENCE: OnceLock<AtomicUsize> = OnceLock::new();
+
+/// Build a Tokio runtime for ZMQ tasks. Multiple calls to this function will
+/// return multiple runtimes with distinct thread name suffixes.
+pub(crate) fn build_zmq_runtime() -> BackgroundShutdownRuntime {
+    let sequence = ZMQ_RUNTIME_SEQUENCE
+        .get_or_init(|| AtomicUsize::new(0))
+        .fetch_add(1, Ordering::Relaxed);
+
+    tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(zmq_worker_threads())
+        .thread_name_fn(move || format!("vllm-zmq-{sequence}"))
+        .enable_all()
+        .build()
+        .expect("failed to build vLLM ZMQ runtime")
+        .into()
+}
+
+/// Get the number of worker threads to use for the ZMQ runtime. If env var
+/// `VLLM_RS_ZMQ_WORKER_THREADS` is set and a valid positive integer, it will be used.
+/// Otherwise, the default value of `DEFAULT_ZMQ_WORKER_THREADS` will be used.
+fn zmq_worker_threads() -> usize {
+    std::env::var(ZMQ_WORKER_THREADS_ENV)
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| *value > 0)
+        .unwrap_or(DEFAULT_ZMQ_WORKER_THREADS)
+}
@@ -30,6 +30,7 @@ tokio-stream.workspace = true
 tokio-util.workspace = true
 tonic.workspace = true
 tonic-prost.workspace = true
+tower.workspace = true
 tower-http.workspace = true
 tracing.workspace = true
 tracing-futures.workspace = true
@@ -7,6 +7,7 @@ mod listener;
 mod lora;
 mod middleware;
 mod routes;
+mod runtime;
 mod server_info;
 mod state;
 mod utils;
@@ -157,6 +158,9 @@ where
            .with_context(|| format!("failed to bind gRPC listener on {grpc_host}:{grpc_port}"))?;
        let addr = grpc_listener.local_addr()?;
        let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
+        let svc = TonicServer::builder()
+            .layer(middleware::request_runtime_layer(state.clone()))
+            .add_service(svc);
        info!(%addr, "starting gRPC server");
        Some((grpc_listener, svc))
    } else {
@@ -238,7 +242,7 @@ where
                shutdown.cancelled().await;
                return Ok(());
            };
-            let server = TonicServer::builder().add_service(svc).serve_with_incoming_shutdown(
+            let server = svc.serve_with_incoming_shutdown(
                TcpListenerStream::new(grpc_listener),
                shutdown.cancelled_owned(),
            );
@@ -2,10 +2,12 @@ mod auth;
 mod cors;
 mod load;
 mod metrics;
+mod offload;
 mod request_id;

 pub use auth::authenticate_api_key;
 pub use cors::{cors_layer, strip_cors_on_no_origin};
 pub use load::track_server_load;
 pub use metrics::track_http_metrics;
+pub(crate) use offload::request_runtime_layer;
 pub use request_id::set_request_id_header;
@@ -0,0 +1,134 @@
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use axum::http::Request;
+use axum::response::{IntoResponse, Response};
+use futures::future::BoxFuture;
+use tokio_util::task::AbortOnDropHandle;
+use tonic::Status;
+use tower::Service;
+use tower::layer::layer_fn;
+use tracing::error;
+
+use crate::error::{ApiError, server_error};
+use crate::state::AppState;
+
+/// Request paths that are run on the request runtime.
+///
+/// These routes can perform CPU-heavy request preparation, including JSON
+/// extraction, validation, chat-template rendering, tokenization, request
+/// lowering, and engine submission. Lightweight operational routes stay on the
+/// HTTP runtime.
+const OFFLOADED_PATHS: &[&str] = &[
+    // HTTP routes:
+    "/v1/chat/completions",
+    "/v1/completions",
+    "/tokenize",
+    "/detokenize",
+    "/inference/v1/generate",
+    // gRPC routes:
+    "/vllm.Generate/Generate",
+    "/vllm.Generate/GenerateStream",
+];
+
+/// Return a Tower layer that runs selected data-plane requests on the request runtime,
+/// so that we can offset heavy request parsing and preprocessing from the HTTP runtime.
+pub(crate) fn request_runtime_layer<S>(
+    state: Arc<AppState>,
+) -> impl tower::Layer<S, Service = RequestRuntimeService<S>> + Clone {
+    layer_fn(move |inner| RequestRuntimeService {
+        inner,
+        state: state.clone(),
+    })
+}
+
+/// Service produced by [`request_runtime_layer`].
+#[derive(Clone)]
+pub(crate) struct RequestRuntimeService<S> {
+    inner: S,
+    state: Arc<AppState>,
+}
+
+impl<S, B> Service<Request<B>> for RequestRuntimeService<S>
+where
+    S: Service<Request<B>> + Clone + Send + 'static,
+    S::Future: Send + 'static,
+    S::Response: RequestRuntimeErrorResponse + Send + 'static,
+    S::Error: Send + 'static,
+    B: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Request<B>) -> Self::Future {
+        if !should_offload(req.uri().path()) {
+            return Box::pin(self.inner.call(req));
+        }
+
+        // Axum extractors and route handlers execute inside the inner service,
+        // so offloading here moves request parsing and preprocessing off the
+        // HTTP runtime without wrapping each handler manually. For streaming
+        // HTTP responses, the response body is still polled on the HTTP runtime.
+        let clone = self.inner.clone();
+        let mut inner = std::mem::replace(&mut self.inner, clone);
+        let task = AbortOnDropHandle::new(self.state.request_runtime().spawn(inner.call(req)));
+
+        Box::pin(async move {
+            match task.await {
+                Ok(result) => result,
+                Err(error) => {
+                    error!(%error, "request runtime task failed");
+                    Ok(S::Response::request_runtime_error_response())
+                }
+            }
+        })
+    }
+}
+
+trait RequestRuntimeErrorResponse {
+    fn request_runtime_error_response() -> Self;
+}
+
+impl RequestRuntimeErrorResponse for Response {
+    fn request_runtime_error_response() -> Self {
+        server_error!("request runtime task failed").into_response()
+    }
+}
+
+impl RequestRuntimeErrorResponse for axum::http::Response<tonic::body::Body> {
+    fn request_runtime_error_response() -> Self {
+        Status::internal("request runtime task failed").into_http()
+    }
+}
+
+fn should_offload(path: &str) -> bool {
+    OFFLOADED_PATHS.contains(&path)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::should_offload;
+
+    #[test]
+    fn offloads_generation_and_tokenization_paths() {
+        assert!(should_offload("/v1/chat/completions"));
+        assert!(should_offload("/v1/completions"));
+        assert!(should_offload("/tokenize"));
+        assert!(should_offload("/detokenize"));
+        assert!(should_offload("/inference/v1/generate"));
+        assert!(should_offload("/vllm.Generate/Generate"));
+        assert!(should_offload("/vllm.Generate/GenerateStream"));
+    }
+
+    #[test]
+    fn passes_through_lightweight_paths() {
+        assert!(!should_offload("/health"));
+        assert!(!should_offload("/metrics"));
+        assert!(!should_offload("/v1/models"));
+    }
+}
@@ -108,6 +108,7 @@ fn build_router_with_options(
    let enable_api_key_auth = state.has_api_keys();
    let mut router = router
        .with_state(state.clone())
+        .layer(middleware::request_runtime_layer(state.clone()))
        .layer(from_fn_with_state(
            state.clone(),
            middleware::track_server_load,
@@ -0,0 +1,54 @@
+use tokio::runtime::Builder;
+use tracing::{info, warn};
+use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;
+
+const REQUEST_WORKER_THREADS_ENV: &str = "VLLM_RS_REQUEST_WORKER_THREADS";
+const DEFAULT_MAX_REQUEST_WORKER_THREADS: usize = 32;
+
+/// Build a Tokio runtime for heavyweight request paths outside the HTTP runtime.
+///
+/// The server middleware uses this runtime for inference and tokenization
+/// routes so CPU-heavy request preparation does not monopolize the HTTP
+/// runtime's worker queue. Dropping the wrapper shuts the runtime down in the
+/// background.
+pub(crate) fn build_request_runtime() -> BackgroundShutdownRuntime {
+    Builder::new_multi_thread()
+        .enable_all()
+        .thread_name("vllm-request")
+        .worker_threads(request_worker_threads())
+        .build()
+        .expect("failed to build request runtime")
+        .into()
+}
+
+/// Get the number of worker threads to use for the request runtime.
+///
+/// If `VLLM_RS_REQUEST_WORKER_THREADS` is set to a valid positive integer, it is
+/// used directly. Otherwise, the runtime uses available parallelism capped by
+/// `DEFAULT_MAX_REQUEST_WORKER_THREADS`.
+fn request_worker_threads() -> usize {
+    if let Some(value) = std::env::var_os(REQUEST_WORKER_THREADS_ENV) {
+        match value.to_string_lossy().parse::<usize>() {
+            Ok(worker_threads) if worker_threads > 0 => return worker_threads,
+            _ => warn!(
+                value = %value.to_string_lossy(),
+                "ignoring invalid {REQUEST_WORKER_THREADS_ENV}"
+            ),
+        }
+    }
+
+    std::thread::available_parallelism()
+        .map(|parallelism| {
+            let available = parallelism.get();
+            let worker_threads = available.min(DEFAULT_MAX_REQUEST_WORKER_THREADS);
+            if worker_threads < available {
+                info!(
+                    available_parallelism = available,
+                    capped_worker_threads = worker_threads,
+                    "capping request runtime worker threads, set {REQUEST_WORKER_THREADS_ENV} to override"
+                );
+            }
+            worker_threads
+        })
+        .unwrap_or(DEFAULT_MAX_REQUEST_WORKER_THREADS)
+}
@@ -1,16 +1,20 @@
 use std::sync::Arc;
+use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};

 use serde_json::Value;
 use sha2::{Digest, Sha256};
+use tokio::runtime::Runtime;
 use tokio::time::{Duration, Instant, sleep_until};
 use tracing::warn;
 use vllm_chat::ChatLlm;
 use vllm_engine_core_client::EngineCoreClient;
 use vllm_engine_core_client::protocol::lora::LoraRequest;
+use vllm_engine_core_client::runtime::BackgroundShutdownRuntime;

 use crate::config::{ApiServerOptions, CorsConfig};
 use crate::lora::{LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError};
+use crate::runtime::build_request_runtime;
 use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot};

 const SHUTDOWN_REFCOUNT_POLL_INTERVAL: Duration = Duration::from_millis(100);
@@ -42,6 +46,8 @@ pub struct AppState {
    lora_manager: LoraManager,
    /// Backend model path reported as `root` for base-model cards.
    model_path: Option<String>,
+    /// Lazily initialized runtime for heavyweight request paths.
+    request_runtime: OnceLock<BackgroundShutdownRuntime>,
 }

 impl AppState {
@@ -68,6 +74,7 @@ impl AppState {
            server_load: AtomicU64::new(0),
            lora_manager: LoraManager::new(),
            model_path: None,
+            request_runtime: OnceLock::new(),
        }
    }

@@ -185,6 +192,12 @@ impl AppState {
        self.chat.engine_core_client()
    }

+    /// Runtime used by middleware to isolate heavyweight request handlers from
+    /// the HTTP reactor.
+    pub(crate) fn request_runtime(&self) -> &Runtime {
+        self.request_runtime.get_or_init(build_request_runtime)
+    }
+
    /// Return the current in-flight inference request count for the `/load`
    /// endpoint.
    pub fn server_load(&self) -> u64 {
@@ -214,6 +227,7 @@ impl AppState {
            match Arc::try_unwrap(self) {
                Ok(state) => {
                    state.chat.shutdown().await?;
+                    drop(state.request_runtime); // shutdown in background
                    return Ok(());
                }
                Err(state) => self = state,
@@ -1171,7 +1171,15 @@ package_data = {
        "third_party/deep_gemm/include/**/*.h",
        "third_party/deep_gemm/include/**/*.hpp",
        # fmha_sm100 sparse CuTe-DSL helper kernels (vendored via cmake)
+        "third_party/fmha_sm100/csrc/**/*.cu",
+        "third_party/fmha_sm100/csrc/**/*.h",
+        "third_party/fmha_sm100/csrc/**/*.jinja",
+        "third_party/fmha_sm100/csrc/**/*.cu.jinja",
        "third_party/fmha_sm100/cute/**/*.cu",
+        "third_party/fmha_sm100/cutlass/include/**/*.h",
+        "third_party/fmha_sm100/cutlass/include/**/*.hpp",
+        "third_party/fmha_sm100/cutlass/tools/util/include/**/*.h",
+        "third_party/fmha_sm100/cutlass/tools/util/include/**/*.hpp",
    ]
 }

@@ -4,6 +4,13 @@ import subprocess

 import pytest

+from vllm.benchmarks.datasets import SampleRequest
+from vllm.benchmarks.throughput import (
+    _run_vllm_chat_requests,
+    add_cli_args,
+)
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"


@@ -28,3 +35,70 @@ def test_bench_throughput():
    print(result.stderr)

    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"
+
+
+def test_bench_throughput_accepts_custom_audio_args():
+    parser = FlexibleArgumentParser()
+    add_cli_args(parser)
+
+    args = parser.parse_args(
+        [
+            "--dataset-name",
+            "custom_audio",
+            "--dataset-path",
+            "audio.jsonl",
+            "--no-oversample",
+            "--custom-output-len",
+            "32",
+            "--enable-multimodal-chat",
+        ]
+    )
+
+    assert args.dataset_name == "custom_audio"
+    assert args.no_oversample
+    assert args.custom_output_len == 32
+    assert args.enable_multimodal_chat
+
+
+def test_vllm_chat_requests_include_multimodal_content():
+    class FakeLLM:
+        def __init__(self):
+            self.prompts = None
+
+        def chat(self, prompts, sampling_params, use_tqdm):
+            del sampling_params, use_tqdm
+            self.prompts = prompts
+            return []
+
+    llm = FakeLLM()
+    audio_content = {
+        "type": "input_audio",
+        "input_audio": {"data": "abc", "format": "wav"},
+    }
+    request = SampleRequest(
+        prompt="Transcribe this audio.",
+        prompt_len=1,
+        expected_output_len=8,
+        multi_modal_data=audio_content,
+    )
+
+    _run_vllm_chat_requests(
+        llm,
+        [request],
+        n=1,
+        disable_detokenize=False,
+        do_profile=False,
+        prequeue_requests=False,
+    )
+
+    assert llm.prompts == [
+        [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Transcribe this audio."},
+                    audio_content,
+                ],
+            }
+        ]
+    ]
@@ -59,9 +59,8 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
    return accuracy


-@multi_gpu_test(num_gpus=4)
-def test_elastic_ep_scaling():
-    vllm_serve_args = [
+def _base_serve_args(use_async_eplb: bool = False) -> list[str]:
+    args = [
        "--trust-remote-code",
        "--tensor-parallel-size",
        "1",
@@ -79,7 +78,11 @@ def test_elastic_ep_scaling():
        "--eplb-config.num_redundant_experts",
        "0",
        "--eplb-config.use_async",
-        "false",
+        "true" if use_async_eplb else "false",
+        "--eplb-config.step_interval",
+        "10",
+        "--eplb-config.window_size",
+        "5",
        "--data-parallel-backend",
        "ray",
        "--data-parallel-size",
@@ -90,7 +93,23 @@ def test_elastic_ep_scaling():

    leader_address = os.environ.get("LEADER_ADDRESS")
    if leader_address:
-        vllm_serve_args.extend(["--data-parallel-address", leader_address])
+        args.extend(["--data-parallel-address", leader_address])
+
+    return args
+
+
+@pytest.mark.parametrize(
+    "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
+)
+@multi_gpu_test(num_gpus=4)
+def test_elastic_ep_scaling(use_async_eplb: bool):
+    if use_async_eplb:
+        from vllm.distributed.eplb.eplb_communicator import has_nixl
+
+        if not has_nixl():
+            pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
+
+    vllm_serve_args = _base_serve_args(use_async_eplb)

    with RemoteOpenAIServer(
        MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
@@ -128,44 +147,24 @@ def test_elastic_ep_scaling():
        print(f"  Tolerance:  {ACCURACY_TOL:.3f}")


+@pytest.mark.parametrize(
+    "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
+)
@multi_gpu_test(num_gpus=4)
-def test_elastic_ep_scaling_uneven():
+def test_elastic_ep_scaling_uneven(use_async_eplb: bool):
    """Test scale up with uneven worker distribution.

    This tests the case where num_new_workers % old_dp_size != 0,
    specifically 2 -> 3 where remainder = 1 % 2 = 1.
    This exercises the remainder handling in sender-receiver pairing.
    """
-    vllm_serve_args = [
-        "--trust-remote-code",
-        "--tensor-parallel-size",
-        "1",
-        "--gpu-memory-utilization",
-        "0.8",
-        "--max-model-len",
-        "4096",
-        "--max-num-seqs",
-        str(MAX_NUM_SEQS),
-        "--enable-expert-parallel",
-        "--all2all-backend",
-        "allgather_reducescatter",
-        "--enable-elastic-ep",
-        "--enable-eplb",
-        "--eplb-config.num_redundant_experts",
-        "0",
-        "--eplb-config.use_async",
-        "false",
-        "--data-parallel-backend",
-        "ray",
-        "--data-parallel-size",
-        "2",
-        "--api-server-count",
-        "1",
-    ]
+    if use_async_eplb:
+        from vllm.distributed.eplb.eplb_communicator import has_nixl

-    leader_address = os.environ.get("LEADER_ADDRESS")
-    if leader_address:
-        vllm_serve_args.extend(["--data-parallel-address", leader_address])
+        if not has_nixl():
+            pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
+
+    vllm_serve_args = _base_serve_args(use_async_eplb)

    with RemoteOpenAIServer(
        MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
@@ -782,3 +782,125 @@ def test_rearrange_expert_weights_profile_mode(world_size):
        _test_rearrange_expert_weights_profile_mode,
        world_size,
    )
+
+
+def _test_nixl_deferred_init_worker(
+    env,
+    world_size: int,
+    num_layers: int,
+    num_local_experts: int,
+    num_logical_experts: int,
+) -> None:
+    """Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""
+    from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator
+
+    set_env_vars_and_device(env)
+
+    vllm_config = VllmConfig()
+    vllm_config.parallel_config.tensor_parallel_size = world_size
+
+    with set_current_vllm_config(vllm_config):
+        ensure_model_parallel_initialized(
+            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
+        )
+
+        ep_group_coordinator = get_tp_group()
+        ep_group = ep_group_coordinator.cpu_group
+        ep_rank = torch.distributed.get_rank()
+        device = torch.device(f"cuda:{ep_rank}")
+
+        total_physical_experts = world_size * num_local_experts
+        hidden_sizes = [32, 64]
+
+        redundancy_config = create_redundancy_config(
+            num_logical_experts, total_physical_experts
+        )
+        old_indices = create_expert_indices_with_redundancy(
+            num_layers,
+            num_logical_experts,
+            total_physical_experts,
+            redundancy_config,
+        )
+
+        new_redundancy_config = create_redundancy_config(
+            num_logical_experts, total_physical_experts
+        )
+        new_indices = create_expert_indices_with_redundancy(
+            num_layers,
+            num_logical_experts,
+            total_physical_experts,
+            new_redundancy_config,
+        )
+
+        expert_weights = create_expert_weights(
+            num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
+        )
+
+        expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
+
+        communicator = NixlEplbCommunicator(
+            cpu_group=ep_group_coordinator.cpu_group,
+            all_expert_weights=expert_weights,
+            expert_buffer=expert_buffer,
+            defer_remote_setup=True,
+        )
+        assert not communicator._remote_state_initialized
+
+        rearrange_expert_weights_inplace(
+            old_indices,
+            new_indices,
+            expert_weights,
+            expert_buffer,
+            ep_group,
+            communicator,
+        )
+
+        assert communicator._remote_state_initialized
+
+    local_ok = verify_expert_weights_after_shuffle(
+        expert_weights,
+        new_indices,
+        hidden_sizes,
+        ep_rank,
+        num_local_experts,
+    )
+
+    local_ok = (
+        verify_redundant_experts_have_same_weights(
+            expert_weights,
+            new_indices,
+            hidden_sizes,
+            ep_rank,
+            world_size,
+            num_local_experts,
+        )
+        and local_ok
+    )
+    assert_verification_synced(
+        local_ok,
+        "Deferred NIXL init verification failed on at least one rank.",
+    )
+
+
+@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available")
+@pytest.mark.parametrize(
+    "world_size,num_layers,num_local_experts,num_logical_experts",
+    [(2, 2, 3, 4)],
+)
+def test_nixl_deferred_init(
+    world_size,
+    num_layers,
+    num_local_experts,
+    num_logical_experts,
+):
+    """Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""
+
+    if torch.accelerator.device_count() < world_size:
+        pytest.skip(f"Need at least {world_size} GPUs to run the test")
+    distributed_run(
+        _test_nixl_deferred_init_worker,
+        world_size,
+        num_layers,
+        num_local_experts,
+        num_logical_experts,
+    )
@@ -824,7 +824,10 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA
    chat_output = chat_response.json()
    invocation_output = invocation_response.json()

-    assert chat_output.keys() == invocation_output.keys()
+    extra_keys = invocation_output.keys() - chat_output.keys()
+    missing_keys = chat_output.keys() - invocation_output.keys()
+    assert missing_keys == set()
+    assert extra_keys <= {"moderation"}
    assert chat_output["choices"] == invocation_output["choices"]


@@ -540,8 +540,8 @@ async def test_max_tokens_with_tool_choice_required(
        tool_choice=tool_choice,
    )
    # When `tool_choice="required"` and the tokens of `tools` exceed `max_tokens`,
-    # both `tool_calls` and `content` should be empty.
+    # `tool_calls` should be absent and `content` should be empty.
    # This behavior should be consistent with OpenAI.
    choice = chat_completion.choices[0]
    assert choice.finish_reason == "length"
-    assert len(choice.message.tool_calls) == 0
+    assert choice.message.tool_calls is None
@@ -465,7 +465,7 @@ class TestGPTOSSChat:
        )

        msg = tool_choice_none.choices[0].message
-        assert len(msg.tool_calls) == 0
+        assert msg.tool_calls is None


 class TestGPTOSSSpeculativeChat:
@@ -386,8 +386,15 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch):
        reasoning_parser="qwen3",
    )

+    request = ResponsesRequest(input="hi", tools=[], stream=False)
+    response_parser = serving._make_response_parser(
+        request,
+        tokenizer,
+        serving._effective_chat_template_kwargs(request),
+    )
+
    # Build a SimpleContext with thinking tokens in the output.
-    context = SimpleContext()
+    context = SimpleContext(response_parser=response_parser)
    token_ids = [1, 10, 2, 20]  # <think> 10 </think> 20 -> reasoning token count = 1
    completion = CompletionOutput(
        index=0,
@@ -412,7 +419,6 @@ async def test_reasoning_tokens_counted_for_text_reasoning_model(monkeypatch):
    async def dummy_result_generator():
        yield None

-    request = ResponsesRequest(input="hi", tools=[], stream=False)
    sampling_params = SamplingParams(max_tokens=16)
    metadata = RequestResponseMetadata(request_id="req")

@@ -636,9 +642,9 @@ class TestHarmonyPreambleStreaming:
        assert "response.output_text.done" not in type_names


-def _make_simple_context_with_output(text, token_ids):
+def _make_simple_context_with_output(text, token_ids, response_parser=None):
    """Create a SimpleContext with a RequestOutput containing the given text."""
-    ctx = SimpleContext()
+    ctx = SimpleContext(response_parser=response_parser)
    completion = CompletionOutput(
        index=0,
        text=text,
@@ -719,6 +725,7 @@ def _mock_parser_with_reasoning(serving, delta_sequence: list[DeltaMessage]):
    mock_parser_instance.parse_delta = mock_parse_delta
    mock_parser_instance.is_reasoning_end = MagicMock(return_value=False)
    serving.parser = MagicMock(return_value=mock_parser_instance)
+    return mock_parser_instance


 class TestStreamingReasoningToContentTransition:
@@ -745,12 +752,12 @@ class TestStreamingReasoningToContentTransition:
            DeltaMessage(reasoning=" end", content="hello"),  # mixed delta
            DeltaMessage(content=" world"),
        ]
-        _mock_parser_with_reasoning(serving, delta_sequence)
+        response_parser = _mock_parser_with_reasoning(serving, delta_sequence)
        # Create contexts for each streaming chunk
        contexts = [
-            _make_simple_context_with_output("chunk1", [10]),
-            _make_simple_context_with_output("chunk2", [20]),
-            _make_simple_context_with_output("chunk3", [30]),
+            _make_simple_context_with_output("chunk1", [10], response_parser),
+            _make_simple_context_with_output("chunk2", [20], response_parser),
+            _make_simple_context_with_output("chunk3", [30], response_parser),
        ]

        async def result_generator():
@@ -767,7 +774,7 @@ class TestStreamingReasoningToContentTransition:
            request=request,
            sampling_params=sampling_params,
            result_generator=result_generator(),
-            context=SimpleContext(),
+            context=SimpleContext(response_parser=response_parser),
            model_name="test-model",
            tokenizer=MagicMock(),
            request_metadata=metadata,
@@ -813,11 +820,11 @@ class TestStreamingReasoningToContentTransition:
            DeltaMessage(reasoning="thinking"),
            DeltaMessage(content="answer"),
        ]
-        _mock_parser_with_reasoning(serving, delta_sequence)
+        response_parser = _mock_parser_with_reasoning(serving, delta_sequence)

        contexts = [
-            _make_simple_context_with_output("chunk1", [10]),
-            _make_simple_context_with_output("chunk2", [20]),
+            _make_simple_context_with_output("chunk1", [10], response_parser),
+            _make_simple_context_with_output("chunk2", [20], response_parser),
        ]

        async def result_generator():
@@ -834,7 +841,7 @@ class TestStreamingReasoningToContentTransition:
            request=request,
            sampling_params=sampling_params,
            result_generator=result_generator(),
-            context=SimpleContext(),
+            context=SimpleContext(response_parser=response_parser),
            model_name="test-model",
            tokenizer=MagicMock(),
            request_metadata=metadata,
@@ -875,11 +882,11 @@ class TestStreamingReasoningToContentTransition:
            DeltaMessage(reasoning="step 1"),
            DeltaMessage(reasoning=" step 2"),
        ]
-        _mock_parser_with_reasoning(serving, delta_sequence)
+        response_parser = _mock_parser_with_reasoning(serving, delta_sequence)

        contexts = [
-            _make_simple_context_with_output("chunk1", [10]),
-            _make_simple_context_with_output("chunk2", [20]),
+            _make_simple_context_with_output("chunk1", [10], response_parser),
+            _make_simple_context_with_output("chunk2", [20], response_parser),
        ]

        async def result_generator():
@@ -896,7 +903,7 @@ class TestStreamingReasoningToContentTransition:
            request=request,
            sampling_params=sampling_params,
            result_generator=result_generator(),
-            context=SimpleContext(),
+            context=SimpleContext(response_parser=response_parser),
            model_name="test-model",
            tokenizer=MagicMock(),
            request_metadata=metadata,
@@ -936,10 +943,10 @@ class TestAutoToolStreaming:
    @staticmethod
    async def _collect_events(delta_sequence: list[DeltaMessage]):
        serving = _make_serving_instance_with_reasoning()
-        _mock_parser_with_reasoning(serving, delta_sequence)
+        response_parser = _mock_parser_with_reasoning(serving, delta_sequence)

        contexts = [
-            _make_simple_context_with_output("chunk", [i])
+            _make_simple_context_with_output("chunk", [i], response_parser)
            for i in range(len(delta_sequence))
        ]

@@ -974,7 +981,7 @@ class TestAutoToolStreaming:
            request=request,
            sampling_params=sampling_params,
            result_generator=result_generator(),
-            context=SimpleContext(),
+            context=SimpleContext(response_parser=response_parser),
            model_name="test-model",
            tokenizer=MagicMock(),
            request_metadata=metadata,
@@ -2,8 +2,25 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import pytest
+from openai.types.chat.chat_completion import ChatCompletion as OpenAIChatCompletion
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk

-from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    ChatCompletionResponseChoice,
+    ChatCompletionResponseStreamChoice,
+    ChatCompletionStreamResponse,
+    ChatMessage,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+    DeltaFunctionCall,
+    DeltaMessage,
+    DeltaToolCall,
+    FunctionCall,
+    ToolCall,
+    UsageInfo,
+)
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
 from vllm.parser.abstract_parser import DelegatingParser

@@ -86,3 +103,101 @@ def test_responses_parser_allows_named_tool_choice_with_none_content():

    assert content is None
    assert tool_calls == []
+
+
+def _chat_response(message: ChatMessage) -> ChatCompletionResponse:
+    return ChatCompletionResponse(
+        model="test-model",
+        choices=[
+            ChatCompletionResponseChoice(
+                index=0,
+                message=message,
+                finish_reason="stop",
+            )
+        ],
+        usage=UsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=2),
+    )
+
+
+def test_chat_completion_response_omits_empty_tool_calls_payload():
+    response = _chat_response(ChatMessage(role="assistant", content="done"))
+
+    payload = response.model_dump()
+    payload_exclude_unset = response.model_dump(exclude_unset=True)
+
+    assert "tool_calls" not in payload["choices"][0]["message"]
+    assert "tool_calls" not in payload_exclude_unset["choices"][0]["message"]
+    parsed = OpenAIChatCompletion.model_validate(payload)
+    assert parsed.choices[0].message.tool_calls is None
+
+
+def test_chat_completion_response_keeps_non_empty_tool_calls_payload():
+    response = _chat_response(
+        ChatMessage(
+            role="assistant",
+            content="",
+            tool_calls=[
+                ToolCall(
+                    function=FunctionCall(
+                        name="get_weather",
+                        arguments='{"city": "Beijing"}',
+                    )
+                )
+            ],
+        )
+    )
+
+    message = response.model_dump()["choices"][0]["message"]
+
+    assert len(message["tool_calls"]) == 1
+    assert message["tool_calls"][0]["function"]["name"] == "get_weather"
+
+
+def _stream_response(delta: DeltaMessage) -> ChatCompletionStreamResponse:
+    return ChatCompletionStreamResponse(
+        id="chatcmpl-test",
+        object="chat.completion.chunk",
+        created=1,
+        model="test-model",
+        choices=[
+            ChatCompletionResponseStreamChoice(
+                index=0,
+                delta=delta,
+                finish_reason=None,
+            )
+        ],
+    )
+
+
+def test_chat_completion_stream_response_omits_empty_tool_calls_payload():
+    response = _stream_response(DeltaMessage(content="done"))
+
+    payload = response.model_dump(exclude_unset=True)
+    payload_json = response.model_dump_json(exclude_unset=True)
+
+    assert "tool_calls" not in payload["choices"][0]["delta"]
+    parsed = ChatCompletionChunk.model_validate_json(payload_json)
+    assert parsed.choices[0].delta.tool_calls is None
+
+
+def test_chat_completion_stream_response_keeps_non_empty_tool_calls_payload():
+    response = _stream_response(
+        DeltaMessage(
+            tool_calls=[
+                DeltaToolCall(
+                    index=0,
+                    id="call-test",
+                    type="function",
+                    function=DeltaFunctionCall(
+                        name="get_weather",
+                        arguments='{"city": "Beijing"}',
+                    ),
+                )
+            ]
+        )
+    )
+
+    delta = response.model_dump(exclude_unset=True)["choices"][0]["delta"]
+
+    assert len(delta["tool_calls"]) == 1
+    assert delta["tool_calls"][0]["function"]["name"] == "get_weather"
@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+model_name: "openai/gpt-oss-20b"
+metric_threshold: 0.568
+reasoning_effort: "low"
@@ -0,0 +1,2 @@
+# DGX Spark model configurations for GPQA evaluation
+gpt-oss-20b-sm120.yaml
@@ -2,4 +2,5 @@ model_name: "nvidia/Qwen3-30B-A3B-FP4"
 accuracy_threshold: 0.89
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 4096"
@@ -0,0 +1,6 @@
+model_name: "amd/Qwen3-30B-A3B-Thinking-2507-FP8"
+accuracy_threshold: 0.81
+num_questions: 1319
+num_fewshot: 5
+max_tokens: 1024
+server_args: "--max-model-len 4096 --gpu-memory-utilization 0.85"
@@ -0,0 +1,6 @@
+model_name: "amd/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8"
+accuracy_threshold: 0.81
+num_questions: 1319
+num_fewshot: 5
+max_tokens: 1024
+server_args: "--max-model-len 4096 --gpu-memory-utilization 0.85"
@@ -2,6 +2,7 @@ model_name: "nm-testing/Qwen3-Next-80B-A3B-Instruct-NVFP4"
 accuracy_threshold: 0.75
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: >-
  --enforce-eager
  --max-model-len 4096
@@ -3,6 +3,7 @@ accuracy_threshold: 0.89
 tolerance: 0.03
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1800
 server_args: >-
  --max-model-len 4096
  --tensor-parallel-size 2
@@ -0,0 +1,5 @@
+model_name: "google/gemma-4-E4B-it-qat-mobile-ct"
+accuracy_threshold: 0.50
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 4096"
@@ -0,0 +1,12 @@
+model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
+accuracy_threshold: 0.86
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --quantization humming
+  --kernel-config.enable_flashinfer_autotune=False
+env:
+  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'
@@ -0,0 +1,10 @@
+model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
+accuracy_threshold: 0.86
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --quantization humming
+  --kernel-config.enable_flashinfer_autotune=False
@@ -0,0 +1,2 @@
+gpt-oss-20b-humming-act-fp8.yaml
+Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml
@@ -0,0 +1,2 @@
+gpt-oss-20b-humming.yaml
+Qwen3-30B-A3B-MXFP4A16-humming.yaml
@@ -0,0 +1,11 @@
+model_name: "openai/gpt-oss-20b"
+accuracy_threshold: 0.30
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --moe-backend humming
+env:
+  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'
@@ -0,0 +1,9 @@
+model_name: "openai/gpt-oss-20b"
+accuracy_threshold: 0.30
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --moe-backend humming
@@ -3,3 +3,5 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
 Qwen1.5-MoE-W4A16-CT.yaml
 DeepSeek-V2-Lite-Instruct-FP8.yaml
 Qwen3-Next-FP8-EP2_MI355.yaml
+Qwen3-30B-A3B-Thinking-2507-FP8.yaml
+Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml
@@ -4,4 +4,5 @@ Llama-3-8B-Instruct-nonuniform-CT.yaml
 Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
 Qwen1.5-MoE-W4A16-CT.yaml
 DeepSeek-V2-Lite-Instruct-FP8.yaml
-Qwen3-30B-A3B-MXFP4A16.yaml
+Qwen3-30B-A3B-MXFP4A16.yaml
+gemma-4-E4B-it-qat-mobile-ct.yaml
@@ -2,4 +2,5 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass"
@@ -2,4 +2,5 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_trtllm"
@@ -2,4 +2,5 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
@@ -2,4 +2,5 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4"
 accuracy_threshold: 0.88
 num_questions: 1319
 num_fewshot: 5
+startup_max_wait_seconds: 1200
 server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"
@@ -48,6 +48,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
    results = evaluate_gsm8k(
        num_questions=eval_config["num_questions"],
        num_shots=eval_config["num_fewshot"],
+        max_tokens=eval_config.get("max_tokens", 256),
        host=host,
        port=port,
        request_timeout_seconds=request_timeout_seconds,
@@ -69,6 +70,15 @@ def test_gsm8k_correctness(config_filename):
            "Marlin kernels are not supported."
        )

+    if (
+        not current_platform.is_cuda()
+        and "gemma-4-E4B-it-qat-mobile-ct" in eval_config["model_name"]
+    ):
+        pytest.skip(
+            "Skipping gemma-4-E4B-it-qat-mobile-ct on non-CUDA platforms. "
+            "Its W2A16 (uint2b2) scheme has no kernel outside CUDA."
+        )
+
    # TODO(akaratza): Enable DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms
    if current_platform.is_rocm() and (
        "deepseek-ai/DeepSeek-V3.2" in eval_config["model_name"]
@@ -134,6 +134,7 @@ def _reference_index_topk(
    topk: int,
    init_blocks: int,
    local_blocks: int,
+    sm_scale: float = 1.0,
 ) -> torch.Tensor:
    total_q, num_idx_heads, _ = idx_q.shape
    out = torch.full(
@@ -149,7 +150,7 @@ def _reference_index_topk(
        num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
        pages = block_table[req_id, :num_blocks]
        k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)
-        score = torch.einsum("qhd,kd->hqk", q.float(), k.float())
+        score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale

        q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)
        k_pos = torch.arange(k.shape[0], device=idx_q.device)
@@ -244,6 +245,270 @@ def test_prefill_index_topk_correctness():
    _assert_topk_indices_equal_unordered(actual, expected)


+# MSA indexer (SM100): fmha_sm100 OnlyScore for the per-block scores, then the
+# Triton minimax_m3_index_topk for selection (no sparse_topk_select). Uses a
+# deterministic construction (idx_q == 1, distinct e4m3-exact per-block values)
+# so scores are strictly monotonic in the block id -> exact top-k agreement.
+def _fmha_indexer_topk(
+    idx_q: torch.Tensor,  # [total_q, H, 128] bf16/e4m3
+    index_cache: torch.Tensor,  # [num_pages, 128, 128] bf16/e4m3
+    block_table: torch.Tensor,
+    q_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    sm_scale: float,
+    topk: int,
+) -> torch.Tensor:
+    """Replicate MiniMaxM3IndexerMSAImpl's score path (single decode/prefill side)."""
+    from vllm.third_party.fmha_sm100.api import _fmha_sm100, _fmha_sm100_plan
+
+    num_idx_heads, head_dim = idx_q.shape[1], idx_q.shape[2]
+    nvp = [(s + 127) // 128 for s in seq_lens.tolist()]
+    kv_indices = torch.cat([block_table[r, : nvp[r]] for r in range(len(nvp))]).to(
+        torch.int32
+    )
+
+    qo = q_lens.cpu().to(torch.int32)
+    kv = seq_lens.cpu().to(torch.int32)
+    plan = _fmha_sm100_plan(
+        qo,
+        kv,
+        num_idx_heads,
+        num_kv_heads=1,
+        qo_offset=kv - qo,
+        page_size=128,
+        output_maxscore=True,
+        causal=True,
+        num_kv_splits=1,
+    )
+    k_pages = index_cache.view(index_cache.shape[0], 1, 128, head_dim)
+    _, max_score = _fmha_sm100(
+        idx_q,
+        k_pages,
+        k_pages,
+        plan,
+        kv_indices=kv_indices,
+        output_o=False,
+        output_maxscore=True,
+        sm_scale=sm_scale,
+    )
+
+    batch = q_lens.numel()
+    cu = torch.zeros(batch + 1, dtype=torch.int32, device=idx_q.device)
+    cu[1:] = q_lens.to(torch.int32).cumsum(0)
+    # max_score [H, k_tiles, total_q] -> transpose to [H, total_q, k_tiles].
+    return minimax_m3_index_topk(
+        max_score.transpose(1, 2),
+        cu,
+        prefix_lens.to(torch.int32),
+        int(q_lens.max()),
+        topk,
+        0,  # init_blocks
+        0,  # local_blocks
+    )
+
+
+# e4m3-exact, strictly-increasing per-block values: with idx_q == 1 (also exact)
+# the per-block scores are exact and distinct in BOTH bf16 and e4m3, so the fp8
+# score path selects the same top-k as the reference (no quantization ties).
+_E4M3_EXACT_VALUES = [
+    *range(1, 17),  # 1..16  (step 1)
+    *range(18, 33, 2),  # 18..32 (step 2)
+    *range(36, 65, 4),  # 36..64 (step 4)
+    *range(72, 129, 8),  # 72..128 (step 8)
+]
+
+
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="fmha_sm100 indexer requires SM100 (Blackwell).",
+)
+@pytest.mark.parametrize("index_dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize(
+    ("q_lens", "prefix_lens"),
+    [
+        ((4, 3), (2048, 2560)),  # prefill: every token sees >= 16 causal blocks
+        ((1, 1, 1), (2048, 3000, 4096)),  # decode: one query token per request
+    ],
+)
+def test_fmha_sm100_indexer_matches_reference(q_lens, prefix_lens, index_dtype):
+    torch.manual_seed(0)
+    num_idx_heads, head_dim = 4, HEAD_DIM
+    device = "cuda"
+
+    q_lens_t = torch.tensor(q_lens, device=device, dtype=torch.int32)
+    prefix_lens_t = torch.tensor(prefix_lens, device=device, dtype=torch.int32)
+    seq_lens = prefix_lens_t + q_lens_t
+    batch = len(q_lens)
+    max_blocks = (int(seq_lens.max()) + BLOCK_SIZE - 1) // BLOCK_SIZE
+    assert max_blocks <= len(_E4M3_EXACT_VALUES)
+    num_pages = batch * max_blocks
+    block_table = torch.randperm(num_pages, device=device, dtype=torch.int32).reshape(
+        batch, max_blocks
+    )
+
+    idx_q = torch.ones(
+        int(q_lens_t.sum()), num_idx_heads, head_dim, device=device, dtype=index_dtype
+    )
+    index_cache = torch.empty(
+        num_pages, BLOCK_SIZE, head_dim, device=device, dtype=index_dtype
+    )
+    for r in range(batch):
+        for b in range(max_blocks):
+            index_cache[block_table[r, b]] = float(_E4M3_EXACT_VALUES[b])
+
+    sm_scale = head_dim**-0.5
+    actual = _fmha_indexer_topk(
+        idx_q,
+        index_cache,
+        block_table,
+        q_lens_t,
+        seq_lens,
+        prefix_lens_t,
+        sm_scale,
+        TOPK,
+    )
+    expected = _reference_index_topk(
+        idx_q,
+        index_cache,
+        block_table,
+        q_lens_t,
+        seq_lens,
+        prefix_lens_t,
+        TOPK,
+        init_blocks=0,
+        local_blocks=0,
+        sm_scale=sm_scale,
+    )
+    _assert_topk_indices_equal_unordered(actual, expected)
+
+
+# Full impl-level parity: drive both MiniMaxM3IndexerMSAImpl (fmha_sm100 score +
+# Triton top-k) and MiniMaxM3IndexerTritonImpl through their real metadata
+# builders on the SAME CommonAttentionMetadata + index cache, and assert the
+# selected blocks agree. This exercises all the metadata the impl/kernels consume
+# (decode/prefill split, cu_seqlens_q rebasing, prefix_lens, kv_indices gather,
+# decode_pages split) -- a metadata bug on either side shifts the causal window
+# or the block->page mapping and breaks the comparison.
+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="fmha_sm100 indexer requires SM100 (Blackwell).",
+)
+@pytest.mark.parametrize("topk", [8, 16])
+def test_msa_indexer_impl_matches_triton(topk, monkeypatch):
+    import vllm.models.minimax_m3.common.indexer as indexer_mod
+    from tests.v1.attention.utils import (
+        BatchSpec,
+        create_common_attn_metadata,
+        create_vllm_config,
+    )
+    from vllm.config import set_current_vllm_config
+    from vllm.forward_context import set_forward_context
+    from vllm.models.minimax_m3.common.indexer import (
+        MiniMaxM3IndexerTritonImpl,
+        MiniMaxM3IndexerTritonMetadataBuilder,
+    )
+    from vllm.models.minimax_m3.nvidia.indexer_msa import (
+        MiniMaxM3IndexerMSAImpl,
+        MiniMaxM3IndexerMSAMetadataBuilder,
+    )
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    num_idx_heads, head_dim = 4, HEAD_DIM
+    # TP=1: avoid requiring an initialized distributed group in a unit test.
+    monkeypatch.setattr(indexer_mod, "get_tensor_model_parallel_world_size", lambda: 1)
+
+    vllm_config = create_vllm_config(
+        block_size=BLOCK_SIZE, max_model_len=8192, max_num_batched_tokens=8192
+    )
+    vllm_config.model_config.hf_config.sparse_attention_config = {
+        "sparse_num_index_heads": num_idx_heads
+    }
+
+    # Decode-first mixed batch: 2 decode reqs (q_len 1) then 2 prefill reqs. Long
+    # prefixes so every token sees > TOPK causal blocks (non-trivial selection).
+    batch = BatchSpec(seq_lens=[2305, 2561, 2624, 2720], query_lens=[1, 1, 64, 96])
+    common = create_common_attn_metadata(
+        batch, BLOCK_SIZE, device, arange_block_indices=True
+    )
+    num_tokens = batch.compute_num_tokens()
+
+    # Deterministic index cache: distinct, monotonic per-logical-block values so
+    # the top-k is unambiguous (both kernels pick the same blocks, no fp ties).
+    block_table = common.block_table_tensor
+    num_pages = int(block_table.max().item()) + 1
+    index_cache = torch.zeros(
+        num_pages, BLOCK_SIZE, head_dim, device=device, dtype=DTYPE
+    )
+    for r, seq_len in enumerate(batch.seq_lens):
+        for b in range((seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE):
+            index_cache[block_table[r, b]] = float(b + 1)
+    index_q = torch.ones(
+        num_tokens, num_idx_heads * head_dim, device=device, dtype=DTYPE
+    )
+
+    spec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=head_dim, dtype=DTYPE
+    )
+    impl_kwargs = dict(
+        num_kv_heads=num_idx_heads,
+        scale=head_dim**-0.5,
+        topk_blocks=topk,
+        sparse_block_size=BLOCK_SIZE,
+        num_index_heads=num_idx_heads,
+        index_head_dim=head_dim,
+        init_blocks=0,
+        local_blocks=0,
+    )
+
+    with set_current_vllm_config(vllm_config):
+        msa_impl = MiniMaxM3IndexerMSAImpl(prefix="idx_msa", **impl_kwargs)
+        triton_impl = MiniMaxM3IndexerTritonImpl(prefix="idx_triton", **impl_kwargs)
+        msa_builder = MiniMaxM3IndexerMSAMetadataBuilder(
+            spec, [msa_impl.index_cache.prefix], vllm_config, device
+        )
+        triton_builder = MiniMaxM3IndexerTritonMetadataBuilder(
+            spec, [triton_impl.index_cache.prefix], vllm_config, device
+        )
+
+    # Both impls score against the same index keys.
+    msa_impl.index_cache.kv_cache = index_cache
+    triton_impl.index_cache.kv_cache = index_cache
+
+    # Exercise the shared persistent top-k buffer for BOTH impls: each must write
+    # decode ([:, :nd]) and prefill ([:, nd:]) into its buffer and return views.
+    # Separate buffers so the two forwards don't clobber each other.
+    nd = sum(q for q in batch.query_lens if q <= 1)
+    msa_impl.topk_indices_buffer = torch.full(
+        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
+    )
+    triton_impl.topk_indices_buffer = torch.full(
+        (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device
+    )
+
+    attn_metadata = {
+        msa_impl.index_cache.prefix: msa_builder.build(0, common),
+        triton_impl.index_cache.prefix: triton_builder.build(0, common),
+    }
+    with set_forward_context(attn_metadata, vllm_config):
+        msa_decode, msa_prefill = msa_impl(index_q)
+        tri_decode, tri_prefill = triton_impl(index_q)
+
+    assert msa_decode is not None and tri_decode is not None
+    assert msa_prefill is not None and tri_prefill is not None
+    _assert_topk_indices_equal_unordered(msa_decode, tri_decode)
+    _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill)
+    # decode/prefill outputs are views into each impl's persistent buffer.
+    for impl, dec, pre in (
+        (msa_impl, msa_decode, msa_prefill),
+        (triton_impl, tri_decode, tri_prefill),
+    ):
+        buf = impl.topk_indices_buffer
+        assert dec.data_ptr() == buf[:, :nd, :].data_ptr()
+        assert pre.data_ptr() == buf[:, nd:, :].data_ptr()
+
+
@pytest.mark.parametrize(
    ("decode_query_len", "max_decode_query_len"),
    [
@@ -317,6 +582,65 @@ def test_decode_index_topk_correctness(
    _assert_topk_indices_equal_unordered(actual, expected)


+@pytest.mark.skipif(
+    not current_platform.is_device_capability_family(100),
+    reason="fp8 e4m3 indexer cache is the SM100 (MSA) path.",
+)
+@pytest.mark.parametrize("num_idx_heads", [1, 4])
+def test_decode_index_topk_fp8(num_idx_heads: int):
+    """The fp8 (e4m3) indexer cache feeds the Triton decode kernel on the MSA
+    path. The kernel must score in fp32 (no scaling) so its top-k matches a
+    reference computed from the dequantized fp8 values."""
+    torch.manual_seed(0)
+    topk, init_blocks, local_blocks, head_dim = 8, 0, 1, 128
+    decode_query_len = 1
+    active_seq_lens = torch.tensor((129, 1025, 4097), device="cuda", dtype=torch.int32)
+    q_lens = torch.full_like(active_seq_lens, decode_query_len)
+    prefix_lens = active_seq_lens - decode_query_len
+    batch = active_seq_lens.numel()
+    max_seq_len = int(active_seq_lens.max())
+    max_blocks = (max_seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE
+    num_pages = batch * max_blocks
+    block_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
+        batch, max_blocks
+    )
+    idx_q = torch.randn(
+        batch * decode_query_len, num_idx_heads, head_dim, device="cuda"
+    ).to(torch.float8_e4m3fn)
+    index_kv_cache = torch.randn(num_pages, BLOCK_SIZE, head_dim, device="cuda").to(
+        torch.float8_e4m3fn
+    )
+
+    actual = minimax_m3_index_decode(
+        idx_q,
+        index_kv_cache,
+        block_table,
+        active_seq_lens,
+        max_seq_len=max_seq_len,
+        topk=topk,
+        init_blocks=init_blocks,
+        local_blocks=local_blocks,
+        num_kv_heads=num_idx_heads,
+        sm_scale=head_dim**-0.5,
+        decode_query_len=decode_query_len,
+    )
+    # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK
+    # in fp32, so it must match an fp32 matmul of the same e4m3 values).
+    expected = _reference_index_topk(
+        idx_q.float(),
+        index_kv_cache.float(),
+        block_table,
+        q_lens,
+        active_seq_lens,
+        prefix_lens,
+        topk,
+        init_blocks,
+        local_blocks,
+        head_dim**-0.5,
+    )
+    _assert_topk_indices_equal_unordered(actual, expected)
+
+
 # Sparse attention kernels.
 def _reference_sparse_attn(
    q: torch.Tensor,
@@ -325,3 +325,24 @@ class TestWeightLoadingWithPaddedHiddenSize:
                shard_id="w2",
                expert_id=0,
            )
+
+
+class TestPerTensorScaleCoercion:
+    """Regression test for shape-(1,) per-tensor scales (issue #43297).
+
+    llm-compressor NVFP4 emits per-tensor weight and input scales as
+    shape-(1,) tensors. `_to_scalar` collapses them to a 0-D scalar so the
+    scalar-slot assignments in the weight loader neither broadcast nor raise.
+    """
+
+    def test_collapses_to_scalar(self):
+        # shape-(1,) and 0-D both reduce to a 0-D scalar.
+        for loaded_weight in (torch.tensor([0.5]), torch.tensor(0.5)):
+            scalar = RoutedExperts._to_scalar(loaded_weight)
+            assert scalar.shape == ()
+            assert scalar.item() == pytest.approx(0.5)
+
+    def test_rejects_non_scalar(self):
+        # numel > 1 must fail loudly instead of silently picking an element.
+        with pytest.raises(RuntimeError):
+            RoutedExperts._to_scalar(torch.tensor([0.1, 0.2]))
@@ -75,7 +75,7 @@ def get_ref_results(
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
-@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm", "b12x"])
+@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass", "cudnn", "trtllm", "b12x"])
@pytest.mark.parametrize("autotune", [False, True])
@torch.inference_mode()
 def test_flashinfer_nvfp4_gemm(
@@ -88,6 +88,8 @@ def test_flashinfer_nvfp4_gemm(
 ) -> None:
    if "trtllm" in backend and dtype == torch.float16:
        pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
+    if backend == "cute-dsl" and not current_platform.is_device_capability_family(100):
+        pytest.skip("FlashInfer cutedsl backend is only supported on SM10x")
    if backend == "b12x" and not current_platform.has_device_capability(120):
        pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
    if backend == "b12x" and not has_flashinfer_b12x_gemm():
@@ -278,3 +278,99 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
    torch.testing.assert_close(
        index_cache.view(-1, HEAD_DIM), expected_index_cache, rtol=0, atol=0
    )
+
+
+# ── Test 3: fp8 (e4m3) index outputs ─────────────────────────────────────────
+# The fp8 score path stores index_q and the index-K cache as e4m3 while q/k/v +
+# q_out stay bf16. Asserts: (1) q/k/v/q_out are bit-identical to the bf16 run
+# (the index dtype must not perturb the main branch), and (2) the e4m3 index
+# outputs dequantize close to the bf16 reference.
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9),
+    reason="e4m3 conversion requires CUDA SM89+.",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 64, 513])
+@pytest.mark.parametrize("block_size", [16, 64])
+def test_sparse_full_fp8_index(num_tokens, block_size):
+    torch.manual_seed(1)
+    device, dtype, eps = "cuda", torch.bfloat16, 1e-6
+    base, max_pos = 5_000_000.0, 4096
+    num_heads, num_kv_heads, num_idx_heads = 16, 4, 4
+
+    q_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    k_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    iq_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    ik_w = torch.randn(HEAD_DIM, dtype=dtype, device=device) * 0.1
+    cos_sin = make_cos_sin_cache(max_pos, ROTARY_DIM, base, dtype, device)
+    positions = torch.randint(
+        0, max_pos, (num_tokens,), dtype=torch.int64, device=device
+    )
+
+    qsz, kvsz = num_heads * HEAD_DIM, num_kv_heads * HEAD_DIM
+    iqsz, iksz = num_idx_heads * HEAD_DIM, HEAD_DIM
+    qkv0 = torch.randn(
+        num_tokens, qsz + 2 * kvsz + iqsz + iksz, dtype=dtype, device=device
+    )
+
+    num_blocks = (num_tokens + block_size - 1) // block_size + 1
+    slot_mapping = torch.randperm(
+        num_blocks * block_size, dtype=torch.int64, device=device
+    )[:num_tokens]
+    index_slot_mapping = torch.roll(slot_mapping, shifts=1)
+
+    def run(index_dtype):
+        qkv = qkv0.clone()
+        kv_cache = torch.zeros(
+            num_blocks,
+            2,
+            block_size,
+            num_kv_heads,
+            HEAD_DIM,
+            dtype=dtype,
+            device=device,
+        )
+        index_cache = torch.zeros(
+            num_blocks, block_size, HEAD_DIM, dtype=index_dtype, device=device
+        )
+        q_out = torch.empty(num_tokens, qsz, dtype=dtype, device=device)
+        index_q = torch.empty(num_tokens, iqsz, dtype=index_dtype, device=device)
+        ops.fused_minimax_m3_qknorm_rope_kv_insert(
+            qkv,
+            q_w,
+            k_w,
+            cos_sin,
+            positions,
+            num_heads,
+            num_kv_heads,
+            ROTARY_DIM,
+            eps,
+            iq_w,
+            ik_w,
+            num_idx_heads,
+            slot_mapping,
+            index_slot_mapping,
+            kv_cache,
+            index_cache,
+            block_size,
+            q_out,
+            index_q,
+        )
+        return qkv, kv_cache, index_cache, q_out, index_q
+
+    qkv_bf, kvc_bf, idxc_bf, qo_bf, iq_bf = run(torch.bfloat16)
+    qkv_fp, kvc_fp, idxc_fp, qo_fp, iq_fp = run(torch.float8_e4m3fn)
+
+    assert iq_fp.dtype == torch.float8_e4m3fn
+    assert idxc_fp.dtype == torch.float8_e4m3fn
+
+    # (1) The main branch (q/k/v in qkv, q_out, kv cache) must be bit-identical:
+    # the index output dtype must not perturb anything else.
+    torch.testing.assert_close(qo_fp, qo_bf, rtol=0, atol=0)
+    torch.testing.assert_close(qkv_fp, qkv_bf, rtol=0, atol=0)
+    torch.testing.assert_close(kvc_fp, kvc_bf, rtol=0, atol=0)
+
+    # (2) Dequantized e4m3 index outputs match the bf16 reference within fp8 ulp.
+    torch.testing.assert_close(iq_fp.float(), iq_bf.float(), rtol=0.13, atol=0.05)
+    torch.testing.assert_close(idxc_fp.float(), idxc_bf.float(), rtol=0.13, atol=0.05)
@@ -14,6 +14,70 @@ TOP_K_VALUES = [2048, 3000]
 BATCH_SIZE = [1, 2, 2048]
 NEXT_N = [1, 8]
 DATA_GENERATION = ["random", "10LSBits"]
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+
+def _has_device_capability(major: int) -> bool:
+    return current_platform.is_cuda() and current_platform.has_device_capability(major)
+
+
+COOPERATIVE_TOPK_BACKEND = pytest.param(
+    "cooperative_topk",
+    marks=pytest.mark.skipif(
+        not _has_device_capability(90),
+        reason="cooperative_topk requires SM90+",
+    ),
+)
+WORKSPACE_TOPK_BACKENDS = ["persistent_topk", COOPERATIVE_TOPK_BACKEND]
+TOPK_BACKENDS = ["top_k_per_row_decode", *WORKSPACE_TOPK_BACKENDS]
+
+
+def _run_topk_backend(
+    backend: str,
+    logits: torch.Tensor,
+    lengths: torch.Tensor,
+    indices: torch.Tensor,
+    top_k: int,
+    max_seq_len: int,
+    next_n: int = 1,
+) -> None:
+    if backend == "top_k_per_row_decode":
+        torch.ops._C.top_k_per_row_decode(
+            logits,
+            next_n,
+            lengths,
+            indices,
+            indices.shape[0],
+            logits.stride(0),
+            logits.stride(1),
+            top_k,
+        )
+    elif backend == "persistent_topk":
+        workspace = torch.empty(
+            RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda"
+        )
+        torch.ops._C.persistent_topk(
+            logits, lengths, indices, workspace, top_k, max_seq_len
+        )
+    elif backend == "cooperative_topk":
+        if indices.shape[0] > 32:
+            pytest.skip(
+                "cooperative_topk supports <=32 rows; "
+                "persistent_topk covers larger batches"
+            )
+        if logits.stride(0) % 4 != 0:
+            pytest.skip(
+                "cooperative_topk requires row stride divisible by 4; "
+                "persistent_topk covers unaligned strides"
+            )
+        workspace = torch.empty(
+            RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda"
+        )
+        torch.ops._C.cooperative_topk(
+            logits, lengths, indices, workspace, top_k, max_seq_len
+        )
+    else:
+        raise ValueError(f"Unknown top-k backend: {backend}")


 def create_random_logits(
@@ -322,16 +386,19 @@ def test_top_k_per_row_decode_large_vocab_size(clean_logits: bool) -> None:
@pytest.mark.parametrize("clean_logits", [True, False])
@pytest.mark.parametrize("top_k", [2048])
@pytest.mark.parametrize("next_n", [1, 4])
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
@torch.inference_mode()
-def test_deepseek_persistent_topk(
+def test_deepseek_workspace_topk(
    seq_len_range: tuple[int, int],
    test_id: str,
    clean_logits: bool,
    top_k: int,
    next_n: int,
+    backend: str,
 ) -> None:
    """
-    Test persistent_topk with varying sequence lengths and speculative decoding.
+    Test workspace top-k backends with varying sequence lengths and speculative
+    decoding.
    Supports speculative decoding with next_n > 1.
    """
    set_random_seed(42 if test_id == "short_sequences" else 43)
@@ -347,6 +414,7 @@ def test_deepseek_persistent_topk(
        dtype=torch.int32,
        device="cuda",
    )
+    seq_lens = (seq_lens + 3) & ~3  # align to 4 for TMA

    # Compute row boundaries for speculative decoding
    row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda")
@@ -366,14 +434,11 @@ def test_deepseek_persistent_topk(
        offsets = torch.arange(next_n, device=logits.device, dtype=torch.int32)
        lengths = (seq_lens.unsqueeze(1) - next_n + 1 + offsets).flatten()

-    workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
    max_seq_len = int(seq_lens.max().item())
-    torch.ops._C.persistent_topk(
-        logits, lengths, indices, workspace, top_k, max_seq_len
-    )
+    _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len, next_n)

    validate_topk_against_reference(
-        logits, indices, row_starts, row_ends, top_k, f"persistent_topk ({test_id})"
+        logits, indices, row_starts, row_ends, top_k, f"{backend} ({test_id})"
    )


@@ -383,9 +448,10 @@ def run_large_context_topk_test(
    top_k: int,
    data_type: str = "random",
    seed: int = 42,
+    backend: str = "cooperative_topk",
 ) -> None:
    """
-    Helper to run persistent_topk kernel test with given parameters.
+    Helper to run a top-k backend test with given parameters.

    Args:
        batch_size: Number of rows/sequences
@@ -393,6 +459,7 @@ def run_large_context_topk_test(
        top_k: Number of top elements to select
        data_type: Type of test data to generate
        seed: Random seed for reproducibility
+        backend: Top-k backend to test
    """
    torch.set_default_device("cuda:0")
    set_random_seed(seed)
@@ -449,11 +516,8 @@ def run_large_context_topk_test(
    # Create output tensor
    indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")

-    workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
    max_seq_len = max(seq_lens)
-    torch.ops._C.persistent_topk(
-        logits, lengths, indices, workspace, top_k, max_seq_len
-    )
+    _run_topk_backend(backend, logits, lengths, indices, top_k, max_seq_len)

    torch.accelerator.synchronize()

@@ -605,8 +669,9 @@ def run_large_context_topk_test(
        ),
    ],
 )
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
@torch.inference_mode()
-def test_persistent_topk_correctness(test_config: dict) -> None:
+def test_workspace_topk_correctness(test_config: dict, backend: str) -> None:
    """
    Comprehensive correctness tests covering:
    - Sequence length edge cases (trivial, boundary, varied)
@@ -620,6 +685,7 @@ def test_persistent_topk_correctness(test_config: dict) -> None:
        seq_lens=test_config["seq_lens"],
        top_k=test_config["top_k"],
        data_type=test_config.get("data_type", "random"),
+        backend=backend,
    )


@@ -668,8 +734,9 @@ def test_persistent_topk_correctness(test_config: dict) -> None:
        ),
    ],
 )
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
@torch.inference_mode()
-def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
+def test_workspace_topk_algorithm_paths(test_config: dict, backend: str) -> None:
    """
    Test different algorithm execution paths (capped at 163840 for DeepSeek V3.2):
    - Batch size scalability (1, 4, 32, 256)
@@ -680,12 +747,14 @@ def test_persistent_topk_algorithm_paths(test_config: dict) -> None:
        batch_size=test_config["batch_size"],
        seq_lens=[test_config["seq_len"]] * test_config["batch_size"],
        top_k=test_config["top_k"],
+        backend=backend,
    )


@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
@torch.inference_mode()
-def test_persistent_topk_stress() -> None:
+def test_workspace_topk_stress(backend: str) -> None:
    """
    Stress test with random configurations to catch edge cases.
    Capped at 163840 (DeepSeek V3.2 max context) for realistic testing.
@@ -700,16 +769,73 @@ def test_persistent_topk_stress() -> None:
        batch_size = torch.randint(1, 32, (1,)).item()

        # Random sequence lengths capped at DeepSeek V3.2 max context
-        seq_lens = torch.randint(100, 163840, (batch_size,)).tolist()
+        seq_lens_tensor = torch.randint(100, 163840, (batch_size,))
+        if backend == "cooperative_topk":
+            seq_lens = ((seq_lens_tensor + 3) & ~3).tolist()
+        else:
+            seq_lens = seq_lens_tensor.tolist()

        run_large_context_topk_test(
            batch_size=batch_size,
            seq_lens=seq_lens,
            top_k=top_k,
            seed=seed,
+            backend=backend,
        )


+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@pytest.mark.parametrize("backend", TOPK_BACKENDS)
+@pytest.mark.parametrize("top_k", [512, 1024, 2048])
+@torch.inference_mode()
+def test_deepseek_topk_backends_no_error_and_reference(
+    backend: str,
+    top_k: int,
+) -> None:
+    """Exercise every production top-k backend on the same inputs."""
+    run_large_context_topk_test(
+        batch_size=4,
+        seq_lens=[2049, 4097, 8191, 12000],
+        top_k=top_k,
+        data_type="random",
+        seed=123,
+        backend=backend,
+    )
+
+
+@pytest.mark.skipif(not _has_device_capability(90), reason="This test requires SM90+")
+@torch.inference_mode()
+def test_cooperative_topk_512_tie_workspace_is_per_row() -> None:
+    """Regression test for TopK=512 tie workspace row overlap."""
+    torch.set_default_device("cuda:0")
+
+    top_k = 512
+    num_rows = 2
+    stride = 65536
+    lengths = torch.tensor([40960, 65536], dtype=torch.int32, device="cuda")
+    logits = torch.full(
+        (num_rows, stride), float("-inf"), dtype=torch.float32, device="cuda"
+    )
+
+    # Row 0 must never select these low indices: many better row-0 ties exist.
+    logits[0, :2048] = -10.0
+    logits[0, 2048 : lengths[0]] = 1.0
+    # Row 1 has higher exact tie scores. With the old row * TopK tie_ws stride,
+    # these row-1 ties could overwrite row 0's TopK=512 refinement workspace.
+    logits[1, : lengths[1]] = 2.0
+
+    indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
+    workspace = torch.empty(RADIX_TOPK_WORKSPACE_SIZE, dtype=torch.uint8, device="cuda")
+    torch.ops._C.cooperative_topk(logits, lengths, indices, workspace, top_k, stride)
+    torch.accelerator.synchronize()
+
+    row0 = indices[0].cpu()
+    assert torch.all(row0 >= 2048), (
+        "cooperative_topk TopK=512 selected row-0 low-score indices, likely "
+        "from overlapping tie_ws rows"
+    )
+
+
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
@pytest.mark.parametrize(
    "test_config",
@@ -774,10 +900,11 @@ def test_persistent_topk_stress() -> None:
    ],
 )
@pytest.mark.parametrize("top_k", [512, 2048])
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
@torch.inference_mode()
-def test_persistent_topk(test_config: dict, top_k: int) -> None:
+def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
    """
-    Tests specific to the persistent_topk kernel:
+    Tests specific to workspace top-k backends:
    - Mixed medium/large rows in the same batch (dynamic per-row dispatch)
    - Boundary around LARGE_THRESHOLD (32K)
    - Trivial + medium + large rows in a single batch
@@ -787,15 +914,17 @@ def test_persistent_topk(test_config: dict, top_k: int) -> None:
        seq_lens=test_config["seq_lens"],
        top_k=top_k,
        data_type=test_config.get("data_type", "random"),
+        backend=backend,
    )


@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
@pytest.mark.parametrize("top_k", [512, 2048])
+@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)
@torch.inference_mode()
-def test_persistent_topk_padded_stride(top_k: int) -> None:
+def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None:
    """
-    Test persistent_topk with padded logits (large stride, small seq_len)
+    Test workspace top-k backends with padded logits (large stride, small seq_len)
    to simulate the e2e CUDAGraph scenario where fp8_paged_mqa_logits
    returns [B, max_model_len] with max_model_len=163840.
    """
@@ -818,11 +947,7 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:

    lengths = torch.tensor(actual_seq_lens, dtype=torch.int32, device="cuda")
    indices = torch.empty((batch_size, top_k), dtype=torch.int32, device="cuda")
-    workspace = torch.empty(1024 * 1024, dtype=torch.uint8, device="cuda")
-
-    torch.ops._C.persistent_topk(
-        logits, lengths, indices, workspace, top_k, max(actual_seq_lens)
-    )
+    _run_topk_backend(backend, logits, lengths, indices, top_k, max(actual_seq_lens))
    torch.accelerator.synchronize()

    # Validate against torch.topk
@@ -840,6 +965,6 @@ def test_persistent_topk_padded_stride(top_k: int) -> None:
            expected_vals = logits[i, expected].cpu().sort(descending=True)[0]
            actual_vals = logits[i, actual].cpu().sort(descending=True)[0]
            assert torch.allclose(expected_vals, actual_vals, rtol=1e-4, atol=1e-4), (
-                f"Row {i}: persistent_topk with padded stride doesn't match. "
+                f"Row {i}: {backend} with padded stride doesn't match. "
                f"seq_len={sl}, stride={padded_stride}"
            )
@@ -1,60 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import sys
-from hashlib import sha256
-from pathlib import Path
-from types import SimpleNamespace
-
-from vllm.model_executor.warmup import kernel_warmup
-
-
-def test_resolve_flashinfer_autotune_file_default_layout(
-    monkeypatch, tmp_path: Path
-) -> None:
-    fake_jit = SimpleNamespace(
-        env=SimpleNamespace(
-            FLASHINFER_WORKSPACE_DIR=Path("/flashinfer-cache/0.6.11.post2/103a")
-        )
-    )
-    fake_flashinfer = SimpleNamespace(jit=fake_jit)
-    monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
-    monkeypatch.setitem(sys.modules, "flashinfer.jit", fake_jit)
-    monkeypatch.setattr(
-        kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
-    )
-    monkeypatch.setattr(kernel_warmup.envs, "VLLM_CACHE_ROOT", str(tmp_path))
-    monkeypatch.setattr(kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None)
-
-    runner = SimpleNamespace(vllm_config=SimpleNamespace())
-    cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
-
-    path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
-
-    assert path == (
-        tmp_path
-        / "flashinfer_autotune_cache"
-        / "0.6.11.post2"
-        / "103a"
-        / cache_hash
-        / "autotune_configs.json"
-    )
-    assert path.parent.is_dir()
-
-
-def test_resolve_flashinfer_autotune_file_uses_override_dir(
-    monkeypatch, tmp_path: Path
-) -> None:
-    monkeypatch.setattr(
-        kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)
-    )
-    monkeypatch.setattr(
-        kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
-    )
-
-    runner = SimpleNamespace(vllm_config=SimpleNamespace())
-    cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
-
-    path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
-
-    assert path == tmp_path / cache_hash / "autotune_configs.json"
@@ -12,7 +12,7 @@ from .embed_utils import run_client_embeddings

 MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
 MAX_MODEL_LEN = 4000
-ATOL = 0.002
+ATOL = 2.3e-3


 def _arr(arr):
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.assets.audio import AudioAsset
+from vllm.model_executor.models.moss_audio import MOSS_AUDIO_PLACEHOLDER
+from vllm.platforms import current_platform
+
+from ...registry import HF_EXAMPLE_MODELS
+from ...utils import check_logprobs_close
+
+CORE_MODEL = pytest.param(
+    "OpenMOSS-Team/MOSS-Audio-4B-Instruct",
+    marks=pytest.mark.core_model,
+    id="4b-instruct",
+)
+
+EXTENDED_MODELS = [
+    "OpenMOSS-Team/MOSS-Audio-4B-Thinking",
+    "OpenMOSS-Team/MOSS-Audio-8B-Instruct",
+    "OpenMOSS-Team/MOSS-Audio-8B-Thinking",
+]
+
+ACCURACY_MODELS = [CORE_MODEL, *EXTENDED_MODELS]
+
+PARALLEL_SMOKE_CASES = [
+    pytest.param({"tensor_parallel_size": 2}, id="tp2"),
+    pytest.param({"pipeline_parallel_size": 2}, id="pp2"),
+    pytest.param(
+        {"tensor_parallel_size": 2, "pipeline_parallel_size": 2},
+        id="tp2_pp2",
+    ),
+]
+
+HF_ACCURACY_SKIP_REASON = (
+    "HF AutoModelForCausalLM cannot load remote MOSS-Audio configs; "
+    "vLLM generation coverage is provided by the smoke tests below."
+)
+
+
+@pytest.mark.core_model
+def test_moss_audio_generation_smoke(vllm_runner) -> None:
+    model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+    model_info.check_transformers_version(on_fail="skip")
+
+    prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
+    audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
+
+    with vllm_runner(
+        model,
+        dtype="half",
+        enforce_eager=True,
+        max_model_len=1024,
+        limit_mm_per_prompt={"audio": 1},
+        trust_remote_code=True,
+    ) as vllm_model:
+        outputs = vllm_model.generate_greedy(
+            prompts,
+            max_tokens=4,
+            audios=audios,
+        )
+
+    assert len(outputs) == 1
+    assert len(outputs[0][1]) > 0
+
+
+@pytest.mark.skip(reason=HF_ACCURACY_SKIP_REASON)
+@pytest.mark.parametrize("model", ACCURACY_MODELS)
+@pytest.mark.parametrize("dtype", ["half"])
+@pytest.mark.parametrize("max_tokens", [8])
+@pytest.mark.parametrize("num_logprobs", [5])
+def test_moss_audio_hf_vllm_accuracy(
+    hf_runner,
+    vllm_runner,
+    model: str,
+    dtype: str,
+    max_tokens: int,
+    num_logprobs: int,
+) -> None:
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+    model_info.check_transformers_version(on_fail="skip")
+
+    prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."]
+    audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
+
+    with vllm_runner(
+        model,
+        dtype=dtype,
+        enforce_eager=True,
+        max_model_len=1024,
+        limit_mm_per_prompt={"audio": 1},
+        trust_remote_code=True,
+    ) as vllm_model:
+        vllm_outputs = vllm_model.generate_greedy_logprobs(
+            prompts,
+            max_tokens,
+            num_logprobs=num_logprobs,
+            audios=audios,
+        )
+
+    with hf_runner(model, dtype=dtype, trust_remote_code=True) as hf_model:
+        hf_outputs = hf_model.generate_greedy_logprobs_limit(
+            prompts,
+            max_tokens,
+            num_logprobs=num_logprobs,
+            audios=audios,
+        )
+
+    check_logprobs_close(
+        outputs_0_lst=hf_outputs,
+        outputs_1_lst=vllm_outputs,
+        name_0="hf",
+        name_1="vllm",
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.parametrize("parallel_kwargs", PARALLEL_SMOKE_CASES)
+def test_moss_audio_parallel_smoke(vllm_runner, parallel_kwargs) -> None:
+    model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
+    required_gpus = parallel_kwargs.get(
+        "tensor_parallel_size", 1
+    ) * parallel_kwargs.get("pipeline_parallel_size", 1)
+    if current_platform.device_count() < required_gpus:
+        # TP/PP integration smoke runs on local or multi-GPU CI only.
+        pytest.skip(f"Requires at least {required_gpus} GPUs")
+
+    model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
+    model_info.check_available_online(on_fail="skip")
+    model_info.check_transformers_version(on_fail="skip")
+
+    prompts = [f"{MOSS_AUDIO_PLACEHOLDER}\nBriefly describe this audio."]
+    audios = [[AudioAsset("mary_had_lamb").audio_and_sample_rate[0]]]
+
+    with vllm_runner(
+        model,
+        dtype="half",
+        enforce_eager=True,
+        max_model_len=1024,
+        limit_mm_per_prompt={"audio": 1},
+        trust_remote_code=True,
+        **parallel_kwargs,
+    ) as vllm_model:
+        outputs = vllm_model.generate_greedy(
+            prompts,
+            max_tokens=4,
+            audios=audios,
+        )
+
+    assert len(outputs) == 1
+    assert len(outputs[0][1]) > 0
@@ -432,6 +432,12 @@ def test_processing_correctness(
        )
    if model_id == "CohereLabs/cohere-transcribe-03-2026":
        pytest.skip("Fix later")
+    if model_id.startswith("OpenMOSS-Team/MOSS-Audio-"):
+        pytest.skip(
+            "MOSS-Audio uses a custom processor that dynamically expands "
+            "audio placeholders from processed audio lengths. Its vLLM "
+            "processor paths are covered by test_moss_audio.py."
+        )

    _test_processing_correctness(
        model_id,
@@ -0,0 +1,694 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from transformers import Qwen3Config
+
+from vllm.model_executor.models.interfaces import SupportsLoRA, supports_lora
+from vllm.model_executor.models.moss_audio import (
+    MOSS_AUDIO_BOS_TOKEN,
+    MOSS_AUDIO_BOS_TOKEN_ID,
+    MOSS_AUDIO_EOS_TOKEN,
+    MOSS_AUDIO_EOS_TOKEN_ID,
+    MOSS_AUDIO_PLACEHOLDER,
+    MOSS_AUDIO_TOKEN,
+    MOSS_AUDIO_TOKEN_ID,
+    GatedMLP,
+    MossAudioConfig,
+    MossAudioDummyInputsBuilder,
+    MossAudioEncoder,
+    MossAudioEncoderConfig,
+    MossAudioModel,
+    MossAudioMultiModalProcessor,
+    MossAudioProcessingInfo,
+    MossAudioProcessor,
+    MossQwen3ForCausalLM,
+    MossQwen3Model,
+)
+from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
+from vllm.multimodal.inputs import batched_tensors_equal
+from vllm.sequence import IntermediateTensors
+
+
+class _Tokenizer:
+    def encode(self, text, add_special_tokens=False):
+        del add_special_tokens
+        return [ord(char) for char in text]
+
+    def decode(self, token_ids, **kwargs):
+        del kwargs
+        return "".join(chr(token_id) for token_id in token_ids)
+
+    def batch_decode(self, batch_token_ids, **kwargs):
+        return [self.decode(token_ids, **kwargs) for token_ids in batch_token_ids]
+
+
+class _MMConfig:
+    enable_mm_embeds = False
+    mm_processor_cache_gb = 1
+
+    def merge_mm_processor_kwargs(self, kwargs):
+        return dict(kwargs)
+
+    def get_limit_per_prompt(self, modality):
+        del modality
+        return 3
+
+
+class _ModelConfig:
+    def __init__(self):
+        self.model = "OpenMOSS-Team/MOSS-Audio-4B-Instruct"
+        self.revision = None
+        self.max_model_len = 4096
+        self.encoder_config = {}
+        self.dtype = torch.float32
+        self.hf_config = MossAudioConfig(language_config=Qwen3Config())
+        self.multimodal_config = _MMConfig()
+
+    def get_multimodal_config(self):
+        return self.multimodal_config
+
+    def get_inputs_embeds_size(self):
+        return None
+
+
+class _ProcessingContext:
+    def __init__(self):
+        self.model_config = _ModelConfig()
+        self.tokenizer = _Tokenizer()
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def get_hf_config(self):
+        return self.model_config.hf_config
+
+    def get_mm_config(self):
+        return self.model_config.get_multimodal_config()
+
+    def get_merged_mm_kwargs(self, kwargs):
+        return self.get_mm_config().merge_mm_processor_kwargs(kwargs)
+
+    def call_hf_processor(self, hf_processor, data, kwargs):
+        merged_kwargs = self.get_merged_mm_kwargs(kwargs)
+        merged_kwargs.setdefault("return_tensors", "pt")
+        return hf_processor(**data, **merged_kwargs)
+
+
+class _TestMossAudioProcessingInfo(MossAudioProcessingInfo):
+    def _get_processor_config_defaults(self):
+        return {}
+
+
+def _vllm_config(tensor_parallel_size=1, pipeline_parallel_size=1, hf_config=None):
+    if hf_config is None:
+        hf_config = MossAudioConfig(language_config=Qwen3Config())
+    return SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_config=hf_config,
+            multimodal_config=None,
+        ),
+        quant_config=None,
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=tensor_parallel_size,
+            pipeline_parallel_size=pipeline_parallel_size,
+        ),
+    )
+
+
+class _FakeAudioEncoder:
+    dtype = torch.float32
+
+    def __init__(self, deepstack_layers=0):
+        self.deepstack_layers = deepstack_layers
+        self.output_deepstack_hidden_states = None
+        self.input_shape = None
+        self.feature_lens = None
+
+    def __call__(self, audio_data, *, feature_lens, output_deepstack_hidden_states):
+        self.input_shape = tuple(audio_data.shape)
+        self.feature_lens = feature_lens.detach().cpu().clone()
+        self.output_deepstack_hidden_states = output_deepstack_hidden_states
+        lengths = MossAudioEncoder._compute_downsampled_length(feature_lens)
+        hidden_states = torch.ones(1, int(lengths.sum().item()), 8)
+        if not output_deepstack_hidden_states:
+            return hidden_states, None
+        return hidden_states, [
+            hidden_states * scale for scale in range(2, 2 + self.deepstack_layers)
+        ]
+
+
+def _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=1, tp_rank=0):
+    import vllm.model_executor.layers.linear as linear_layers
+    import vllm.model_executor.models.moss_audio as moss_audio_module
+    import vllm.model_executor.parameter as parameter_module
+
+    for module in (moss_audio_module, linear_layers, parameter_module):
+        monkeypatch.setattr(
+            module, "get_tensor_model_parallel_world_size", lambda: tp_size
+        )
+    monkeypatch.setattr(
+        linear_layers, "get_tensor_model_parallel_rank", lambda: tp_rank
+    )
+    monkeypatch.setattr(
+        parameter_module, "get_tensor_model_parallel_rank", lambda: tp_rank
+    )
+    monkeypatch.setattr(
+        linear_layers, "tensor_model_parallel_all_reduce", lambda tensor: tensor
+    )
+
+
+def _build_moss_audio_processor(cache=None):
+    ctx = _ProcessingContext()
+    info = _TestMossAudioProcessingInfo(ctx)
+    return (
+        MossAudioMultiModalProcessor(
+            info,
+            MossAudioDummyInputsBuilder(info),
+            cache=cache,
+        ),
+        ctx,
+    )
+
+
+def _assert_mm_inputs_equal(left, right):
+    assert left["prompt_token_ids"] == right["prompt_token_ids"]
+    assert left["mm_hashes"] == right["mm_hashes"]
+
+    left_placeholder = left["mm_placeholders"]["audio"][0]
+    right_placeholder = right["mm_placeholders"]["audio"][0]
+    assert left_placeholder.offset == right_placeholder.offset
+    assert left_placeholder.length == right_placeholder.length
+    assert left_placeholder.is_embed.tolist() == right_placeholder.is_embed.tolist()
+
+    assert batched_tensors_equal(
+        left["mm_kwargs"].get_data(),
+        right["mm_kwargs"].get_data(),
+    )
+
+
+@pytest.mark.parametrize(
+    ("prompt", "prefix"),
+    [
+        (
+            f"before {MOSS_AUDIO_PLACEHOLDER} after",
+            [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
+        ),
+        (
+            f"before {MOSS_AUDIO_BOS_TOKEN}{MOSS_AUDIO_TOKEN}"
+            f"{MOSS_AUDIO_TOKEN}{MOSS_AUDIO_EOS_TOKEN} after",
+            [*[ord(char) for char in "before "], MOSS_AUDIO_BOS_TOKEN_ID],
+        ),
+        ("Describe this audio.", [MOSS_AUDIO_BOS_TOKEN_ID]),
+    ],
+)
+def test_moss_audio_processor_expands_audio_placeholders(prompt, prefix):
+    raw_mel_len = 17
+    processed = MossAudioProcessor(_Tokenizer())(
+        text=prompt, audio=[torch.zeros(160 * raw_mel_len)]
+    )
+    input_ids = processed["input_ids"][0].tolist()
+
+    assert input_ids[: len(prefix)] == prefix
+    assert input_ids.count(MOSS_AUDIO_BOS_TOKEN_ID) == 1
+    assert input_ids.count(MOSS_AUDIO_EOS_TOKEN_ID) == 1
+    assert input_ids.count(MOSS_AUDIO_TOKEN_ID) == (
+        MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
+    )
+    assert processed["audio_data"].shape == (1, 128, raw_mel_len)
+    assert processed["audio_data_seqlens"].tolist() == [raw_mel_len]
+
+
+def test_moss_audio_processor_preserves_placeholder_without_audio():
+    processed = MossAudioProcessor(_Tokenizer())(
+        text=f"before {MOSS_AUDIO_PLACEHOLDER} after"
+    )
+
+    assert processed["input_ids"][0].tolist() == [
+        *[ord(char) for char in "before "],
+        MOSS_AUDIO_BOS_TOKEN_ID,
+        MOSS_AUDIO_TOKEN_ID,
+        MOSS_AUDIO_EOS_TOKEN_ID,
+        *[ord(char) for char in " after"],
+    ]
+    assert "audio_data" not in processed
+    assert "audio_data_seqlens" not in processed
+
+
+def test_moss_audio_multimodal_processor_handles_token_and_cache_paths():
+    raw_mel_len = 17
+    audio = np.zeros(160 * raw_mel_len, dtype=np.float32)
+    prompt = f"{MOSS_AUDIO_PLACEHOLDER}\nTranscribe this audio."
+
+    baseline_processor, ctx = _build_moss_audio_processor()
+    mm_items = baseline_processor.info.parse_mm_data({"audio": [audio]})
+    token_prompt = ctx.get_tokenizer().encode(prompt, add_special_tokens=False)
+
+    baseline_text = baseline_processor(
+        prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+    baseline_token = baseline_processor(
+        token_prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+
+    cache = MultiModalProcessorOnlyCache(ctx.model_config)
+    cached_processor, _ = _build_moss_audio_processor(cache=cache)
+    cached_text_miss = cached_processor(
+        prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+    cached_text_hit = cached_processor(
+        prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+    cached_token_hit = cached_processor(
+        token_prompt,
+        mm_items=mm_items,
+        hf_processor_mm_kwargs={},
+    )
+
+    expected_audio_tokens = MossAudioEncoder.compute_num_audio_tokens(raw_mel_len)
+    prompt_token_ids = baseline_text["prompt_token_ids"]
+    assert prompt_token_ids.count(MOSS_AUDIO_TOKEN_ID) == expected_audio_tokens
+    assert baseline_text["mm_placeholders"]["audio"][0].length == (
+        expected_audio_tokens + 2
+    )
+
+    _assert_mm_inputs_equal(baseline_text, baseline_token)
+    _assert_mm_inputs_equal(baseline_text, cached_text_miss)
+    _assert_mm_inputs_equal(baseline_text, cached_text_hit)
+    _assert_mm_inputs_equal(baseline_text, cached_token_hit)
+
+
+def test_moss_audio_supports_language_model_lora_only():
+    assert supports_lora(MossAudioModel)
+
+    model = object.__new__(MossAudioModel)
+    assert isinstance(model, SupportsLoRA)
+
+    mapping = model.get_mm_mapping()
+    assert mapping.language_model == ["language_model."]
+    assert mapping.tower_model == []
+    assert mapping.connector == []
+
+
+def test_moss_audio_error_paths():
+    model = object.__new__(MossAudioModel)
+    with pytest.raises(ValueError, match="DeepStack audio token count mismatch"):
+        model._cache_deepstack_input_embeds(
+            inputs_embeds=torch.zeros(4, 8),
+            deepstack_embeddings=((torch.ones(1, 8),),),
+            is_multimodal=torch.tensor([False, True, True, False]),
+        )
+
+    with pytest.raises(ValueError, match="too short"):
+        MossAudioProcessor(_Tokenizer())(
+            text=MOSS_AUDIO_PLACEHOLDER, audio=[torch.empty(0)]
+        )
+    with pytest.raises(ValueError, match="too short"):
+        model._parse_and_validate_audio_input(
+            audio_data=torch.zeros(1, 128, 1),
+            audio_data_seqlens=torch.tensor([0], dtype=torch.long),
+        )
+
+
+def test_moss_audio_validates_tp_config():
+    vllm_config = _vllm_config(tensor_parallel_size=2)
+    vllm_config.model_config.hf_config.adapter_hidden_size = 7
+
+    with pytest.raises(ValueError, match="adapter_hidden_size"):
+        MossAudioModel(vllm_config=vllm_config)
+
+    vllm_config = _vllm_config(tensor_parallel_size=2)
+    vllm_config.model_config.hf_config.audio_config.d_model = 6
+    vllm_config.model_config.hf_config.audio_config.encoder_attention_heads = 3
+    with pytest.raises(ValueError, match="encoder_attention_heads"):
+        MossAudioModel(vllm_config=vllm_config)
+
+
+def test_moss_audio_rejects_audio_data_list_seqlen_count_mismatch():
+    model = object.__new__(MossAudioModel)
+
+    with pytest.raises(ValueError, match="audio_data batch size"):
+        model._parse_and_validate_audio_input(
+            audio_data=[torch.zeros(128, 8), torch.zeros(128, 11)],
+            audio_data_seqlens=torch.tensor([8], dtype=torch.long),
+        )
+
+
+@pytest.mark.parametrize("deepstack_scales", [(), (7, 11)])
+def test_moss_audio_embed_multimodal_packs_by_audio(deepstack_scales):
+    model = object.__new__(MossAudioModel)
+    model.audio_encoder = _FakeAudioEncoder(len(deepstack_scales))
+    model.audio_adapter = lambda hidden_states: hidden_states * 5
+    model.deepstack_audio_merger_list = [
+        lambda hidden_states, scale=scale: hidden_states * scale
+        for scale in deepstack_scales
+    ]
+    model.deepstack_input_embeds = None
+
+    embeddings = model.embed_multimodal(
+        audio_data=torch.zeros(2, 128, 9),
+        audio_data_seqlens=torch.tensor([8, 9], dtype=torch.long),
+    )
+
+    assert model.audio_encoder.output_deepstack_hidden_states is bool(deepstack_scales)
+    assert [embeds.shape for embeds in embeddings] == [
+        torch.Size([1, 8 * (1 + len(deepstack_scales))]),
+        torch.Size([2, 8 * (1 + len(deepstack_scales))]),
+    ]
+    if not deepstack_scales:
+        assert model.deepstack_input_embeds is None
+        return
+
+    main_embeddings, deepstack_embeddings = model._split_multimodal_embeddings(
+        embeddings, hidden_size=8
+    )
+    assert [embeds.shape for embeds in main_embeddings] == [
+        torch.Size([1, 8]),
+        torch.Size([2, 8]),
+    ]
+    assert [[e.shape for e in layer] for layer in deepstack_embeddings] == [
+        [torch.Size([1, 8]), torch.Size([2, 8])] for _ in deepstack_scales
+    ]
+    assert torch.equal(main_embeddings[0], torch.full((1, 8), 5.0))
+    for idx, scale in enumerate(deepstack_scales):
+        assert torch.equal(
+            deepstack_embeddings[idx][0],
+            torch.full((1, 8), float((idx + 2) * scale)),
+        )
+
+
+def test_moss_audio_embed_input_ids_caches_packed_deepstack():
+    class _FakeLanguageModel:
+        def embed_input_ids(self, input_ids):
+            return torch.zeros(input_ids.shape[0], 8)
+
+    model = object.__new__(MossAudioModel)
+    model.language_model = _FakeLanguageModel()
+    model.deepstack_audio_merger_list = [object(), object()]
+    model.deepstack_input_embeds = None
+    multimodal_embeddings = (
+        torch.cat([torch.full((1, 8), x) for x in (5.0, 14.0, 33.0)], dim=-1),
+        torch.cat([torch.full((2, 8), x) for x in (7.0, 22.0, 44.0)], dim=-1),
+    )
+    is_multimodal = torch.tensor([False, True, True, True, False])
+
+    inputs_embeds = model.embed_input_ids(
+        input_ids=torch.arange(5),
+        multimodal_embeddings=multimodal_embeddings,
+        is_multimodal=is_multimodal,
+    )
+
+    assert torch.equal(inputs_embeds[1], torch.full((8,), 5.0))
+    assert torch.equal(inputs_embeds[2], torch.full((8,), 7.0))
+    assert torch.equal(inputs_embeds[3], torch.full((8,), 7.0))
+    assert model.deepstack_input_embeds is not None
+    tensors = model.deepstack_input_embeds.tensors
+    assert set(tensors) == {"deepstack_input_embeds_0", "deepstack_input_embeds_1"}
+    for tensor in tensors.values():
+        assert tensor[is_multimodal].abs().sum() > 0
+        assert torch.equal(tensor[~is_multimodal], torch.zeros(2, 8))
+
+
+def _patch_pp_group(monkeypatch, *, first=True, last=True):
+    import vllm.model_executor.models.moss_audio as moss_audio_module
+
+    monkeypatch.setattr(
+        moss_audio_module,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=first, is_last_rank=last),
+    )
+
+
+def test_moss_audio_pp_forward_routes_deepstack(monkeypatch):
+    for first in (True, False):
+        calls: list[dict[str, object]] = []
+
+        def fake_lm(*args, _calls=calls, **kwargs):
+            del args
+            _calls.append(kwargs)
+            return torch.ones(1, 1)
+
+        _patch_pp_group(monkeypatch, first=first)
+        model = object.__new__(MossAudioModel)
+        torch.nn.Module.__init__(model)
+        model.language_model = fake_lm
+        cached = IntermediateTensors({"deepstack_input_embeds_0": torch.ones(3, 8)})
+        inter = IntermediateTensors(
+            {
+                "hidden_states": torch.ones(3, 8),
+                "residual": torch.zeros(3, 8),
+                "deepstack_input_embeds_0": torch.full((3, 8), 5.0),
+            }
+        )
+        inputs_embeds = torch.full((3, 8), 9.0)
+        model.deepstack_input_embeds = cached
+
+        model.forward(
+            input_ids=None,
+            positions=torch.arange(3),
+            intermediate_tensors=None if first else inter,
+            inputs_embeds=inputs_embeds if first else None,
+        )
+
+        kwargs = calls[0]
+        assert kwargs["inputs_embeds"] is (inputs_embeds if first else None)
+        assert kwargs["deepstack_input_embeds"] is (cached if first else inter)
+        assert model.deepstack_input_embeds is None
+
+    calls = []
+
+    def fake_lm_non_first_rank(*args, **kwargs):
+        del args
+        calls.append(kwargs)
+        return torch.ones(1, 1)
+
+    _patch_pp_group(monkeypatch, first=False)
+    model = object.__new__(MossAudioModel)
+    torch.nn.Module.__init__(model)
+    model.language_model = fake_lm_non_first_rank
+    model.deepstack_input_embeds = IntermediateTensors({})
+    inter = IntermediateTensors(
+        {
+            "hidden_states": torch.ones(3, 8),
+            "residual": torch.zeros(3, 8),
+        }
+    )
+
+    model.forward(
+        input_ids=None,
+        positions=torch.arange(3),
+        intermediate_tensors=inter,
+        inputs_embeds=torch.ones(3, 8),
+    )
+    assert calls[0]["inputs_embeds"] is None
+    assert calls[0]["deepstack_input_embeds"] is inter
+
+
+def test_moss_qwen3_deepstack_keys_for_pp(monkeypatch):
+    class AddOne(torch.nn.Module):
+        def forward(self, positions, hidden_states, residual):
+            del positions, residual
+            return hidden_states + 1, torch.zeros_like(hidden_states)
+
+    def make_model(num_layers, deepstack_layers=None):
+        model = object.__new__(MossQwen3Model)
+        torch.nn.Module.__init__(model)
+        model.start_layer, model.end_layer = 0, num_layers
+        model.layers = torch.nn.ModuleList([AddOne() for _ in range(num_layers)])
+        model.norm = lambda hidden_states, residual: (hidden_states, residual)
+        model._maybe_add_hidden_state = lambda aux, *args: aux
+        model.deepstack_inject_layer_indices = (
+            range(0) if deepstack_layers is None else deepstack_layers
+        )
+        return model
+
+    _patch_pp_group(monkeypatch, first=True, last=True)
+    output = make_model(3).forward(
+        input_ids=None,
+        positions=torch.arange(2),
+        inputs_embeds=torch.zeros(2, 4),
+        deepstack_input_embeds=IntermediateTensors(
+            {
+                "deepstack_input_embeds_2": torch.full((2, 4), 5.0),
+            }
+        ),
+    )
+    assert torch.equal(output, torch.full((2, 4), 8.0))
+
+    _patch_pp_group(monkeypatch, first=True, last=False)
+    deepstack = IntermediateTensors(
+        {
+            "deepstack_input_embeds_0": torch.full((2, 4), 7.0),
+            "deepstack_input_embeds_3": torch.full((2, 4), 11.0),
+        }
+    )
+    output = make_model(2, range(4)).forward(
+        input_ids=None,
+        positions=torch.arange(2),
+        inputs_embeds=torch.zeros(2, 4),
+        deepstack_input_embeds=deepstack,
+    )
+    assert isinstance(output, IntermediateTensors)
+    assert set(output.tensors) == {
+        "hidden_states",
+        "residual",
+        "deepstack_input_embeds_2",
+        "deepstack_input_embeds_3",
+    }
+    assert torch.equal(output["hidden_states"], torch.full((2, 4), 9.0))
+    assert torch.equal(output["deepstack_input_embeds_2"], torch.zeros(2, 4))
+    assert output["deepstack_input_embeds_3"] is deepstack["deepstack_input_embeds_3"]
+
+    inner_model = make_model(0, range(2))
+    inner_model.make_empty_intermediate_tensors = lambda batch, dtype, device: (
+        IntermediateTensors(
+            {
+                "hidden_states": torch.zeros(batch, 4, dtype=dtype, device=device),
+                "residual": torch.zeros(batch, 4, dtype=dtype, device=device),
+            }
+        )
+    )
+    language_model = object.__new__(MossQwen3ForCausalLM)
+    torch.nn.Module.__init__(language_model)
+    language_model.model = inner_model
+    language_model.config = SimpleNamespace(hidden_size=4)
+    language_model.deepstack_inject_layer_indices = range(2)
+
+    tensors = MossQwen3ForCausalLM.make_empty_intermediate_tensors(
+        language_model,
+        batch_size=3,
+        dtype=torch.float16,
+        device=torch.device("cpu"),
+    )
+
+    assert set(tensors.tensors) == {
+        "hidden_states",
+        "residual",
+        "deepstack_input_embeds_0",
+        "deepstack_input_embeds_1",
+    }
+    assert tensors["deepstack_input_embeds_0"].shape == (3, 4)
+    assert tensors["deepstack_input_embeds_0"].dtype == torch.float16
+
+    _patch_pp_group(monkeypatch, first=True, last=False)
+    forward_tensors = inner_model.forward(
+        input_ids=None,
+        positions=torch.arange(3),
+        inputs_embeds=torch.ones(3, 4, dtype=torch.float16),
+        deepstack_input_embeds=None,
+    )
+    assert isinstance(forward_tensors, IntermediateTensors)
+    assert set(forward_tensors.tensors) == set(tensors.tensors)
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_moss_audio_gated_mlp_tp_shapes_and_loading(monkeypatch, tp_size):
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.device import DeviceConfig
+
+    _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=tp_size)
+    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+        mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
+
+    params = dict(mlp.named_parameters())
+    assert params["gate_up_proj.weight"].shape == torch.Size([16 // tp_size, 4])
+    assert params["down_proj.weight"].shape == torch.Size([6, 8 // tp_size])
+
+    gate_weight = torch.arange(32, dtype=torch.float32).reshape(8, 4)
+    up_weight = torch.arange(100, 132, dtype=torch.float32).reshape(8, 4)
+    down_weight = torch.arange(48, dtype=torch.float32).reshape(6, 8)
+    loaded = mlp.load_weights(
+        [
+            ("gate_proj.weight", gate_weight),
+            ("up_proj.weight", up_weight),
+            ("down_proj.weight", down_weight),
+        ]
+    )
+
+    assert loaded == {"gate_up_proj.weight", "down_proj.weight"}
+    shard = 8 // tp_size
+    assert torch.equal(params["gate_up_proj.weight"][:shard], gate_weight[:shard])
+    assert torch.equal(params["gate_up_proj.weight"][shard:], up_weight[:shard])
+    assert torch.equal(params["down_proj.weight"], down_weight[:, : 8 // tp_size])
+
+    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+        packed_mlp = GatedMLP(input_size=4, hidden_size=8, output_size=6)
+    packed_params = dict(packed_mlp.named_parameters())
+    loaded = packed_mlp.load_weights(
+        [("gate_up_proj.weight", torch.cat([gate_weight, up_weight], dim=0))]
+    )
+    assert loaded == {"gate_up_proj.weight"}
+    assert torch.equal(
+        packed_params["gate_up_proj.weight"][:shard],
+        gate_weight[:shard],
+    )
+    assert torch.equal(
+        packed_params["gate_up_proj.weight"][shard:],
+        up_weight[:shard],
+    )
+
+
+def test_moss_audio_encoder_loads_realistic_attention_weight_names(monkeypatch):
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.config.device import DeviceConfig
+
+    _patch_tensor_parallel_for_linear_layers(monkeypatch, tp_size=2)
+    config = MossAudioEncoderConfig(
+        d_model=8,
+        output_dim=8,
+        num_mel_bins=8,
+        encoder_layers=1,
+        encoder_attention_heads=2,
+        encoder_ffn_dim=16,
+        downsample_hidden_size=2,
+        deepstack_encoder_layer_indexes=[],
+    )
+    with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+        encoder = MossAudioEncoder(config)
+
+    attention = encoder.layers[0].self_attn
+    assert all(hasattr(attention, name) for name in ("q_proj", "k_proj", "v_proj"))
+    assert hasattr(attention, "out_proj")
+    assert not hasattr(attention, "qkv")
+    assert attention.k_proj.bias is None
+
+    weight_names = [
+        "layers.0.self_attn.q_proj.weight",
+        "layers.0.self_attn.q_proj.bias",
+        "layers.0.self_attn.k_proj.weight",
+        "layers.0.self_attn.v_proj.weight",
+        "layers.0.self_attn.v_proj.bias",
+        "layers.0.self_attn.out_proj.weight",
+        "layers.0.self_attn.out_proj.bias",
+        "conv1.weight",
+        "conv1.bias",
+    ]
+    params = dict(encoder.named_parameters(remove_duplicate=False))
+    assert "layers.0.self_attn.k_proj.bias" not in params
+    weights = {
+        name: torch.full_like(params[name], fill_value=float(i + 1))
+        for i, name in enumerate(weight_names)
+    }
+
+    loaded = AutoWeightsLoader(encoder).load_weights(weights.items())
+
+    assert "load_weights" not in MossAudioEncoder.__dict__
+    assert loaded == set(weight_names)
+    assert not any(".qkv." in name for name in loaded)
+    assert torch.equal(
+        params["layers.0.self_attn.q_proj.weight"],
+        weights["layers.0.self_attn.q_proj.weight"],
+    )
@@ -90,6 +90,7 @@ def test_models(example_prompts, model_name) -> None:
 EAGER = [True, False]

 SM_100_NVFP4_BACKENDS = [
+    "flashinfer_cutedsl",
    "flashinfer_cudnn",
    "flashinfer_trtllm",
    "flashinfer_cutlass",
@@ -102,12 +103,18 @@ SM_100_NVFP4_BACKENDS = [
    "backend",
    [
        "emulation",
+        "flashinfer_cutedsl",
        "flashinfer_cudnn",
        "flashinfer_trtllm",  # the small seq_len ensures trtllm_8x4_layout backend is used
        "flashinfer_cutlass",
    ],
 )
 def test_nvfp4(vllm_runner, model, eager, backend):
+    if backend == "flashinfer_cutedsl" and not (
+        current_platform.is_device_capability_family(100)
+    ):
+        pytest.skip("The flashinfer_cutedsl backend is only supported on SM10x")
+
    if (
        not current_platform.has_device_capability(100)
        and backend in SM_100_NVFP4_BACKENDS
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Varun Sundar Rabindranathandkhluu	ee0da84ab9	[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`)	2026-06-28 00:04:08 -07:00
Kevin H. Luu GitHub Claude	217c64a976	[CI] Raise gsm8k startup timeout for MoE Refactor Qwen3 NVFP4 configs (#46882 ) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-06-26 16:33:47 -07:00
Kevin H. Luu GitHub Claude	cfe8a4d063	[CI] Raise gsm8k startup timeout for Qwen3 NVFP4 trtllm configs (#46881 ) Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-06-26 16:15:00 -07:00
Robert Shawandkhluu	6d37570a1c	Fix P/D with DP Supervisor (#46628 ) Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> (cherry picked from commit `c5e3c40877`)	2026-06-25 13:15:20 -07:00
Gabriel Wuandkhluu	f85a9f112a	[Bugfix] FLASHINFER_MLA_SPARSE_SM120 compatibility with GLM-5 NVFP4 (#46506 ) (cherry picked from commit `8fa36fbbeb`)	2026-06-25 13:15:20 -07:00
Micah Williamsonandkhluu	836b5acb1b	[ROCm] Begin Deprecation Window for CUDA_VISIBLE_DEVICES on ROCm (#46636 ) Signed-off-by: Micah Williamson <micah.williamson@amd.com> (cherry picked from commit `d6696e2385`)	2026-06-25 13:15:02 -07:00
Lucas Wilkinsonandkhluu	b36db10f27	[KV Offload] Gate packed HMA KV cache on cross-layer config (#46252 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> (cherry picked from commit `e7df232288`)	2026-06-25 13:15:02 -07:00
Wentao Yeandkhluu	b70c13ea47	[Bug] Fix `IndentationError: expected an indented block after 'with' statement` (#46627 ) Signed-off-by: yewentao256 <zhyanwentao@126.com> (cherry picked from commit `1cd3e0e945`)	2026-06-25 13:15:02 -07:00
Cyrus Leungandkhluu	6829a6d55f	[Bugfix] Re-enable FP8 MoE on NVIDIA Thor (#46339 ) Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> (cherry picked from commit `24d5186138`)	2026-06-25 13:15:02 -07:00
meihancandkhluu	6ed56e04ff	[Bugfix] Fix illegal memory access from a forward during a partial wake_up (#44483 ) Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Nick Hill <nickhill123@gmail.com> (cherry picked from commit `93ec645878`)	2026-06-25 13:15:02 -07:00
Mike Gandkhluu	21089b4762	[NVFP4 MoE/Deepseek V4] Marlin: wire SwiGLU clamp + allow it for clamped models on non-Blackwell (#45836 ) Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com> (cherry picked from commit `0775b882ba`) Signed-off-by: khluu <khluu000@gmail.com>	2026-06-25 13:15:02 -07:00
Mohammad Miadh Angkadandkhluu	7b3d595eb1	[CI/Build] Fix topk histogram build on SM75 (#46550 ) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> (cherry picked from commit `191826ec61`)	2026-06-24 01:01:53 -07:00
hurukawaandkhluu	e66b766bc4	feat: support to OpenMOSS-Team (#44124 ) Signed-off-by: nagisa-kun <1434936049@qq.com> Signed-off-by: nagisa19 <1434936049@qq.com> Signed-off-by: nagisa <1434936049@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> (cherry picked from commit `489abadfb8`)	2026-06-24 01:01:53 -07:00
Jee Jee Liandkhluu	51000a95a1	[Kernel] GLM5 Router GEMM (#46385 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> (cherry picked from commit `9d6fdc2901`)	2026-06-24 01:01:53 -07:00
Roberto L. Castroandkhluu	2e0c5f52f7	[Perf][DSv4/DSv3.2] Add cluster-cooperative topK kernel for low-latency scenarios (#43008 ) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> (cherry picked from commit `855cd4d787`)	2026-06-24 01:01:53 -07:00
Yongye Zhuandkhluu	0085058ca8	[Kernel] Add FlashInferCutedslMxfp8LinearKernel (cute-dsl mm_mxfp8) (#46393 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> (cherry picked from commit `11b56b2ff2`)	2026-06-24 01:01:53 -07:00
Gabriel Wuandkhluu	fdd8e4efd3	[Bugfix] Allow flashinfer_cutlass as a clamped NVFP4 MoE backend (#46492 ) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> (cherry picked from commit `0d4d164488`) Signed-off-by: khluu <khluu000@gmail.com>	2026-06-24 01:01:48 -07:00
Yongye Zhuandkhluu	c6561c2fb6	[Minimax-M3] BF16/FP8 Indexer using MSA (#45892 ) Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> (cherry picked from commit `6691f087a6`)	2026-06-24 00:53:06 -07:00
jack GitHub QwertyJack Chauncey	6c427dd401	[BugFix] Omit empty tool_calls from OpenAI chat responses (#44105 ) Signed-off-by: QwertyJack <7554089+QwertyJack@users.noreply.github.com> Signed-off-by: Chauncey <chaunceyjiang@gmail.com> Co-authored-by: QwertyJack <7554089+QwertyJack@users.noreply.github.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>	2026-06-23 13:43:53 +08:00
Chauncey GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	3ce5823762	[Refactor] Responses API parser state into conversation context (#46030 ) Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-23 13:42:58 +08:00
Woosuk KwonandGitHub	04c2a8deac	[DeepEP V2] Fill invalid recv_topk_idx with -1 (#46432 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 21:45:49 -07:00
Tan Pin Siang GitHub vllmellm Hongxia Yang Jun Kang Chow Chun Fang TianDi101 functionstackx	7e47fb72b5	[ROCm][P/D] Fix MoRIIO WRITE mode for mixed KV layouts (#46290 ) Signed-off-by: Tan Pin Siang <tanpinsiang@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Jun Kang Chow <junkangchow@gmail.com> Co-authored-by: Chun Fang <chun.fang@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com>	2026-06-23 12:12:51 +08:00
Bugen Zhao GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	a8481be7a9	[Rust Frontend][Perf] Use dedicated runtime for HTTP/request-processing/ZMQ (#46051 ) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-06-23 04:03:20 +00:00
Kunshang JiandGitHub	9d3317172c	[XPU][CI]fix xpu kv cache layout test (#46429 ) Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>	2026-06-23 03:43:29 +00:00
Dao007forever GitHub Claude Opus 4.8	430a95ae3a	[v1][kvcache] Honor prefix-cache retention interval for Mamba/linear attention (#45845 ) Signed-off-by: Dao Le <daole@inferact.ai> Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>	2026-06-22 19:51:11 -07:00
Mike GandGitHub	56e5797511	[Quant] Enable modelopt_mixed on Turing (SM75) (#45375 ) Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com>	2026-06-22 19:30:49 -07:00
Rui Yin GitHub abinggo	8db12169a4	fix: stream Qwen3 tool call string arguments (#46351 ) Signed-off-by: Rui Yin <2260891073@qq.com> Co-authored-by: abinggo <107740309+abinggo@users.noreply.github.com>	2026-06-23 10:26:37 +08:00
MichaelCaoo GitHub Claude	33f50773cb	[Doc] Fix typos, grammar, and broken commands across docs (#46398 ) Signed-off-by: MichaelCaoo <a992033227@163.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-06-23 02:01:22 +00:00
Micah WilliamsonandGitHub	fa36f86d77	[CI] Torch 2.11 flaky test_spec_decode_logprobs and gritlm tests (#45772 ) Signed-off-by: Micah Williamson <micah.williamson@amd.com>	2026-06-23 01:26:54 +00:00
Michael Goin GitHub Claude Opus 4.8	8207ce0850	[Bugfix] Fix humming lm_head crash and FusedMoE weight_shape coercion (#46420 ) Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>	2026-06-22 18:19:29 -07:00
Woosuk Kwon GitHub Roy Wang gnovack Claudemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	e48592066e	[DeepEP V2] Bound num_max_tokens_per_rank in do_expand=False (#46404 ) Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Roy Wang <jasonailu87@gmail.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-22 18:14:53 -07:00
stefankoncarevic GitHub Andreas Karatzas	91ba720b75	[ROCm][CI] Only require q_scale==1.0 for fp8 query in RocmAttention (#46148 ) Signed-off-by: stefankoncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-06-22 18:25:43 -05:00
fxmarty-amdandGitHub	6ead164e52	[CI] Add TP=4 requirement to `test_mixed_precision_model_accuracies` (#46161 ) Signed-off-by: Felix Marty <Felix.Marty@amd.com>	2026-06-22 18:19:43 -05:00
Bowen Bao GitHub Andreas Karatzas	c97e8f99d6	[ROCm][Quantization][4/N] refactor quark_moe fp8 w/ oracle (#43721 ) Signed-off-by: Bowen Bao <bowenbao@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-06-22 15:58:03 -07:00
Guipeng Zhang GitHub Codex Michael Goin	183b5f27ea	[Bugfix][V1][TurboQuant] Reserve workspace before CUDA graph capture (#44053 ) Signed-off-by: Guipeng Zhang <zhangguipeng23z@ict.ac.cn> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-06-22 15:47:48 -07:00
ZewenShen-Cohere GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	ca5b24695b	Fix static actorder handling for compressed-tensors WNA16 MoE (#41161 ) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-22 15:46:46 -07:00
Charlie FuandGitHub	6f6bd3b8fe	[ROCm][CI] Increase the max wait time for server startup (#46417 ) Signed-off-by: charlifu <charlifu@amd.com>	2026-06-22 17:46:31 -05:00
Andreas KaratzasandGitHub	70ef4d3009	[ROCm][CI] Purging away redundant test group definitions (#46418 ) Signed-off-by: Andreas Karatzas <akaratza@amd.com>	2026-06-22 15:42:47 -07:00
Tyler Michael Smith GitHub Claude	e2fe837572	[CI] Fix CPU-Multi-Modal Model Tests timeout by adding a 4th shard (#46388 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-06-22 22:08:00 +00:00
Aarushi JainandGitHub	fbf9ff7cf4	[CI][ROCm] Restrict MLA cross-layer KV cache test to supported backends on ROCm (#46401 ) Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com>	2026-06-22 17:05:26 -05:00
Michael Goin GitHub Claude Opus 4.8	6cc2c9ba3a	[CI] Add DGX Spark GPQA smoke test (#39541 ) Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-22 14:52:38 -07:00
Varshith GitHub Claude Michael Goin	c0b2d8f471	[Bugfix] FusedMoE: coerce shape-(1,) per-tensor scales to 0-D scalar … (#43362 ) Signed-off-by: Varshith <kvarshithgowda@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-06-22 13:26:53 -07:00
Mohammad Miadh AngkadandGitHub	d1a38c2762	[Kernel][Performance] Add FlashInfer cutedsl NVFP4 GEMM backend (#42235 ) Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>	2026-06-22 16:17:18 -04:00
stefankoncarevic GitHub Andreas Karatzas	2b4a7491ec	[ROCm][CI] Query total device memory via amdsmi to avoid HIP init (#46141 ) Signed-off-by: stefankoncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-06-22 15:12:24 -05:00
SaddssandGitHub	82ede09a5a	[Bugfix][KVConnector] Fix SimpleCPUOffloadConnector GPU->CPU store race (#46278 )	2026-06-22 13:08:47 -07:00
Nick HillandGitHub	fbf520cf3a	[MRV2] Generalize use of `WhisperModelState` (#46096 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-06-22 12:40:02 -07:00
Gabriel Wu GitHub Claude Opus 4.8 Yongye Zhu	44d95069e9	Enable DeepSeek V4 and GLM-5.1 on SM120 (#43477 ) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>	2026-06-22 11:54:14 -07:00
Dao007forever GitHub Claude Opus 4.8	3ce15fd574	[v1][kvconnector] DecodeBenchConnector: fill list/tuple (Mamba/KDA) KV caches (#45080 ) Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-22 11:54:00 -07:00
Jinzhen Lin GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	e4b3da3feb	[Quantization][CI] add humming lm-eval test (#43752 ) Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-22 11:23:55 -07:00
Nick Hill GitHub Claude Opus 4.8 Roger Wang	3e6529cc0e	[Bugfix][Spec Decode] Fix EAGLE drafter multimodal encoder cache misses (#46315 ) Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2026-06-22 18:14:02 +00:00
Ilya Markov GitHub Markov Ilya	ac614587f5	[EPLB] Enable nixl eplb communicator for elastic ep (#45013 ) Signed-off-by: Markov Ilya <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com>	2026-06-22 10:54:08 -07:00