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

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

.buildkite/hardware_tests/cpu.yaml

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

.buildkite/image_build/image_build_arm64.sh

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

.buildkite/intel_jobs/expert_parallelism_intel.yaml

@@ -1,5 +1,5 @@
 group: Expert Parallelism
-depends_on: 
+depends_on:
   - image-build-xpu
 steps:
 - label: EPLB Algorithm

.buildkite/intel_jobs/models_multimodal_intel.yaml

@@ -1,5 +1,5 @@
 group: Models - Multimodal
-depends_on: 
+depends_on:
   - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"

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

@@ -34,7 +34,7 @@ case "${test_suite}" in
     pytest -v -s v1/worker --ignore=v1/worker/test_gpu_model_runner.py --ignore=v1/worker/test_worker_memory_snapshot.py
     pytest -v -s v1/structured_output
     pytest -v -s v1/test_serial_utils.py
-    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py
+    pytest -v -s v1/spec_decode --ignore=v1/spec_decode/test_max_len.py --ignore=v1/spec_decode/test_speculators_eagle3.py --ignore=v1/spec_decode/test_acceptance_length.py --ignore=v1/spec_decode/test_speculators_correctness.py
     pytest -v -s v1/kv_connector/unit --ignore=v1/kv_connector/unit/test_multi_connector.py --ignore=v1/kv_connector/unit/test_example_connector.py --ignore=v1/kv_connector/unit/test_lmcache_integration.py --ignore=v1/kv_connector/unit/test_hf3fs_client.py --ignore=v1/kv_connector/unit/test_hf3fs_connector.py --ignore=v1/kv_connector/unit/test_hf3fs_metadata_server.py --ignore=v1/kv_connector/unit/test_offloading_connector.py
     ;;
   server)

.buildkite/test-amd.yaml

@@ -161,24 +161,6 @@ steps:
   commands:
   - "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
 
-- label: PyTorch Fullgraph # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  torch_nightly: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/compilation/
-  - vllm/model_executor/
-  - vllm/v1/attention/
-  - vllm/config/compilation.py
-  - csrc/
-  - tests/compile
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
-
 - label: PyTorch Fullgraph Smoke Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -263,37 +245,6 @@ steps:
   - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s v1/shutdown
   - pytest -v -s v1/worker/test_worker_memory_snapshot.py
 
-- label: Elastic EP Scaling Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/
-  - vllm/engine/
-  - vllm/executor/
-  - vllm/compilation/
-  - tests/distributed/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s distributed/test_elastic_ep.py
-
-- label: EPLB Execution # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_4
-  num_gpus: 4
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/distributed/eplb
-  - tests/distributed/test_eplb_execute.py
-  - tests/distributed/test_eplb_spec_decode.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s distributed/test_eplb_execute.py
-  - pytest -v -s distributed/test_eplb_spec_decode.py
-
 - label: Pipeline + Context Parallelism (4 GPUs) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -315,78 +266,8 @@ steps:
   - pytest -v -s distributed/test_pp_cudagraph.py
   - pytest -v -s distributed/test_pipeline_parallel.py
 
-#-----------------------------------------------------------  mi250 · evals  -----------------------------------------------------------#
-
-- label: Multi-Modal Accuracy Eval (Small Models) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
-  source_file_dependencies:
-  - vllm/multimodal/
-  - vllm/inputs/
-  - vllm/v1/core/
-  - vllm/platforms/rocm.py
-  - vllm/model_executor/model_loader/
-  commands:
-  - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
-
-#---------------------------------------------------------  mi250 · examples  ----------------------------------------------------------#
-
-- label: Examples # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/examples"
-  source_file_dependencies:
-  - vllm/entrypoints
-  - vllm/multimodal
-  - examples/
-  - vllm/platforms/rocm.py
-  commands:
-    - pip install tensorizer
-    # Basic
-    - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    - python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
-    - python3 basic/offline_inference/classify.py
-    - python3 basic/offline_inference/embed.py
-    - python3 basic/offline_inference/score.py
-    # Multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
-    # Pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
-    # Features demo
-    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
-
 #----------------------------------------------------------  mi250 · kernels  ----------------------------------------------------------#
 
-- label: Kernels Core Operation Test # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
-  agent_pool: mi250_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - csrc/
-  - tests/kernels/core
-  - tests/kernels/test_top_k_per_row.py
-  - tests/kernels/test_concat_mla_q.py
-  - vllm/model_executor/layers/rotary_embedding/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
-
 - label: Kernels Helion Test # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -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

.buildkite/test_areas/kernels.yaml

@@ -47,8 +47,10 @@ steps:
   - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
   - vllm/models/deepseek_v4/common/ops/
   - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
   commands:
     - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_top_k_per_row.py
 
 - label: Deepseek V4 Kernel Test (B200)
   key: deepseek-v4-kernel-test-b200

.buildkite/test_areas/lm_eval.yaml

@@ -136,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
@@ -177,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

AGENTS.md

@@ -135,6 +135,12 @@ Do not modify code in these areas without first reading and following the
 linked guide. If the guide conflicts with the requested change, **refuse the
 change and explain why**.
 
+Security reviewers should start with [`SECURITY.md`](SECURITY.md),
+[`docs/usage/security.md`](docs/usage/security.md), and
+[`docs/contributing/vulnerability_management.md`](docs/contributing/vulnerability_management.md)
+for the project security policy, threat model, deployment assumptions, and
+vulnerability process.
+
 - **Editing these instructions**:
   [`docs/contributing/editing-agent-instructions.md`](docs/contributing/editing-agent-instructions.md)
   — Rules for modifying AGENTS.md or any domain-specific guide it references.

CMakeLists.txt

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

cmake/external_projects/deepgemm.cmake

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

cmake/external_projects/fmha_sm100.cmake

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

cmake/external_projects/qutlass.cmake

@@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR})
   set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
 endif()
 
+# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid.
+# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees.
 if(QUTLASS_SRC_DIR)
-  FetchContent_Declare(
-    qutlass
-    SOURCE_DIR ${QUTLASS_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  set(_qutlass_user_src "${QUTLASS_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _qutlass_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(QUTLASS_SRC_DIR "${_qutlass_user_src}")
+  if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}")
+    message(FATAL_ERROR
+      "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'")
+  endif()
+  set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}")
+  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
-  FetchContent_Declare(
-    qutlass
-    GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git
-    GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
-endif()
+  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
 
-FetchContent_Populate(qutlass)
+  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _qutlass_fc_root)
+    set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_qutlass_src "${_qutlass_fc_root}/qutlass-src")
+  set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build")
+  set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild")
+
+  if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp")
+    set(qutlass_SOURCE_DIR "${_qutlass_src}")
+    set(qutlass_BINARY_DIR "${_qutlass_bin}")
+  else()
+    FetchContent_Populate(
+      qutlass
+      SUBBUILD_DIR "${_qutlass_sub}"
+      SOURCE_DIR "${_qutlass_src}"
+      BINARY_DIR "${_qutlass_bin}"
+      GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}"
+      GIT_TAG "${_QUTLASS_UPSTREAM_TAG}"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+endif()
 
 if(NOT qutlass_SOURCE_DIR)
   message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")

csrc/cpu/cpu_wna16.cpp

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

csrc/cpu/micro_gemm/cpu_micro_gemm_rvv.hpp

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

csrc/cpu/utils.hpp

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

csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu

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

csrc/libtorch_stable/cache_kernels.cu

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

csrc/libtorch_stable/cooperative_topk.cu

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

csrc/libtorch_stable/cooperative_topk.cuh

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

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -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());

csrc/libtorch_stable/moe/dsv3_router_gemm_bf16_out.cu

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

csrc/libtorch_stable/moe/dsv3_router_gemm_entry.cu

@@ -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);
     }
   }
 }

csrc/libtorch_stable/moe/dsv3_router_gemm_float_out.cu

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

csrc/libtorch_stable/ops.h

@@ -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,

csrc/libtorch_stable/persistent_topk.cuh

@@ -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_

csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu

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

csrc/libtorch_stable/topk_histogram_4096.cuh

@@ -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_

csrc/libtorch_stable/torch_bindings.cpp

@@ -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",

docker/Dockerfile.xpu

@@ -132,7 +132,7 @@ CMD ["/bin/bash"]
 FROM vllm-base AS ucx-nixl-build
 
 ARG UCX_VERSION=v1.21.0-rc2
-ARG NIXL_VERSION=0.10.1
+ARG NIXL_VERSION=v1.2.0
 
 # Build-time only: compiler, autotools, and verbs dev headers
 RUN apt-get update -y && apt-get install -y --no-install-recommends \
@@ -149,25 +149,25 @@ RUN apt-get update -y && apt-get install -y --no-install-recommends \
 # patchelf (installed via uv) is used by the NIXL wheel build to rewrite
 # RPATH entries, making the wheel portable across stages.
 RUN --mount=type=cache,target=/root/.cache/uv \
-    git clone https://github.com/openucx/ucx /tmp/ucx_source && \
-    cd /tmp/ucx_source && git checkout "${UCX_VERSION}" && \
+    git clone --depth 1 --branch "${UCX_VERSION}" https://github.com/openucx/ucx /tmp/ucx_source && \
+    cd /tmp/ucx_source && \
     bash autogen.sh && \
     ./configure --prefix=/tmp/ucx_install --with-ze=yes --enable-examples --enable-mt && \
-    make CFLAGS="-Wno-error=incompatible-pointer-types" -j8 && make install && \
-    git clone https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
-    cd /tmp/nixl_source && git checkout "${NIXL_VERSION}" && \
+    make CFLAGS="-Wno-error=incompatible-pointer-types" -j"$(nproc)" && make install && \
+    git clone --depth 1 --branch "${NIXL_VERSION}" https://github.com/ai-dynamo/nixl /tmp/nixl_source && \
+    cd /tmp/nixl_source && \
     uv pip install --upgrade meson pybind11 patchelf && \
     uv pip install -r requirements.txt && \
     PKG_CONFIG_PATH=/tmp/ucx_install/lib/pkgconfig \
     LD_LIBRARY_PATH=/tmp/ucx_install/lib \
     python -m pip wheel --no-deps . -w /tmp/nixl_wheels/ && \
     find /tmp/ucx_install -type f \( -name '*.a' -o -name '*.la' \) -delete && \
-    rm -rf /tmp/ucx_install/include /tmp/ucx_install/share /tmp/ucx_install/etc /tmp/ucx_install/lib/cmake /tmp/ucx_install/bin && \
-    rm -rf /tmp/ucx_source /tmp/nixl_source
+    rm -rf /tmp/ucx_install/{include,share,etc,bin} /tmp/ucx_install/lib/cmake \
+           /tmp/ucx_source /tmp/nixl_source
 
 FROM vllm-base AS vllm-openai
 
-ARG NIXL_VERSION=0.10.1
+ARG NIXL_VERSION=v1.2.0
 
 # Copy compiled UCX runtime libraries and the pre-built NIXL wheel.
 # No compiler or autotools are installed in this stage.
@@ -192,7 +192,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     librdmacm1t64 \
     && rm -rf /var/lib/apt/lists/* \
     && uv pip install --no-deps /tmp/nixl_wheels/nixl*.whl \
-    && uv pip install nixl==${NIXL_VERSION} \
+    && uv pip install nixl==${NIXL_VERSION} && uv pip uninstall nixl-cu13 \
     && rm -rf /tmp/nixl_wheels
 
 RUN --mount=type=cache,target=/root/.cache/uv \

docs/benchmarking/cli.md

@@ -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
 ```

docs/configuration/optimization.md

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

docs/contributing/README.md

@@ -304,9 +304,15 @@ review process:
   resources. The reviewer will add `ready` label to the PR when the PR is
   ready to merge or a full CI run is needed.
 
-### Escalating Stalled Contributions
+### Pull Request Limits and Escalation
 
-If you have an important contribution that has not yet received maintainer attention, please email us at:
+vLLM uses GitHub's [pull request limit](https://github.blog/open-source/maintainers/how-pull-request-limits-are-cutting-down-the-noise/)
+for contributors without write access. The current cap is 6 open PRs. If this
+blocks well-intentioned critical work, contact a committer to request bypass
+list access.
+
+If you need an expedited review for an important contribution, please email us
+at:
 
 <pr-review-request@vllm.ai>
 

docs/design/attention_backends.md

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

docs/design/cuda_graphs.md

@@ -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"""
 ```

docs/design/metrics.md

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

docs/design/prefix_caching.md

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

docs/features/quantization/gptqmodel.md

@@ -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 \

docs/features/quantization/llm_compressor/int8_w8a8.md

@@ -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:

docs/features/speculative_decoding/README.md

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

docs/features/tool_calling.md

@@ -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:
 

docs/models/pooling_models/README.md

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

docs/models/pooling_models/scoring.md

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

docs/models/supported_models.md

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

requirements/cuda.txt

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

rust/src/chat/src/multimodal.rs

@@ -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}"))
         })

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

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

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

@@ -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"),

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

@@ -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:?}")]

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

@@ -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;

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

@@ -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)
+}

rust/src/server/Cargo.toml

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

rust/src/server/src/lib.rs

@@ -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(),
             );

rust/src/server/src/middleware/mod.rs

@@ -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;

rust/src/server/src/middleware/offload.rs

@@ -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"));
+    }
+}

rust/src/server/src/routes.rs

@@ -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,

rust/src/server/src/runtime.rs

@@ -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)
+}

rust/src/server/src/state.rs

@@ -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,

setup.py

@@ -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",
     ]
 }
 

tests/benchmarks/test_throughput_cli.py

@@ -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,
+                ],
+            }
+        ]
+    ]

tests/distributed/test_elastic_ep.py

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

tests/distributed/test_eplb_execute.py

@@ -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,
+    )

tests/entrypoints/openai/chat_completion/test_chat.py

@@ -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"]
 
 

tests/entrypoints/openai/chat_completion/test_completion_with_function_calling.py

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

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -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:

tests/entrypoints/openai/responses/test_serving_responses.py

@@ -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,

tests/entrypoints/openai/test_tool_choice_content_none.py

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

tests/evals/gpt_oss/configs/gpt-oss-20b-sm120.yaml

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

tests/evals/gpt_oss/configs/models-spark.txt

@@ -0,0 +1,2 @@
+# DGX Spark model configurations for GPQA evaluation
+gpt-oss-20b-sm120.yaml

tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4.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"

tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-FP8.yaml

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

tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml

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

tests/evals/gsm8k/configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml

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

tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml

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

tests/evals/gsm8k/configs/gemma-4-E4B-it-qat-mobile-ct.yaml

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

tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml

@@ -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"}'

tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming.yaml

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

tests/evals/gsm8k/configs/humming/config-act-fp8.txt

@@ -0,0 +1,2 @@
+gpt-oss-20b-humming-act-fp8.yaml
+Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml

tests/evals/gsm8k/configs/humming/config.txt

@@ -0,0 +1,2 @@
+gpt-oss-20b-humming.yaml
+Qwen3-30B-A3B-MXFP4A16-humming.yaml

tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming-act-fp8.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"}'

tests/evals/gsm8k/configs/humming/gpt-oss-20b-humming.yaml

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

tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt

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

tests/evals/gsm8k/configs/models-small.txt

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

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.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"

tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml

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

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.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 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm"

tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml

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

tests/evals/gsm8k/test_gsm8k_correctness.py

@@ -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"]

tests/kernels/attention/test_cutlass_mla_decode.py

@@ -212,3 +212,69 @@ def test_cutlass_mla_decode(
     print(
         f"{t:.3f} ms, {FLOPS / 10**9 / t:.0f} TFLOPS,", f"{bytes / 10**6 / t:.0f} GB/s"
     )
+
+
+@pytest.mark.skipif(
+    not current_platform.has_device_capability(100),
+    reason=CUTLASS_MLA_UNSUPPORTED_REASON,
+)
+@torch.inference_mode()
+def test_cutlass_mla_decode_cross_layer_view():
+    """The kernel must read the cache's page-dim stride instead of assuming
+    pages are packed back-to-back. A per-layer view into a cross-layer
+    (block-major) cache has stride(0) inflated by num_layers; outputs must
+    match a contiguous cache holding the same data exactly."""
+    device = torch.device("cuda:0")
+    torch.set_default_dtype(torch.bfloat16)
+    torch.set_default_device(device)
+    torch.manual_seed(42)
+
+    b, mean_sk, d, dv, block_size = 4, 512, 576, 512, 64
+    num_layers, layer_idx = 3, 1
+    scale = math.sqrt(d) ** (-1)
+
+    num_pages = b * (mean_sk // block_size)
+    cache_seqlens = torch.full((b,), mean_sk, dtype=torch.int32)
+    block_table = torch.arange(num_pages, dtype=torch.int32).view(
+        b, mean_sk // block_size
+    )
+
+    kv_contig = torch.randn(num_pages, block_size, d)
+    # Neighbor layers hold random data so packed-pages addressing reads
+    # garbage rather than zeros.
+    kv_cross_layer = torch.randn(num_pages, num_layers, block_size, d)
+    kv_view = kv_cross_layer[:, layer_idx]
+    kv_view.copy_(kv_contig)
+    assert kv_view.stride(0) == num_layers * block_size * d
+
+    q_nope = torch.randn(b, 128, dv)
+    q_pe = torch.randn(b, 128, d - dv)
+    sm_count = num_compute_units(device.index)
+    workspace_size = ops.sm100_cutlass_mla_get_workspace_size(
+        mean_sk, b, sm_count, num_kv_splits=1
+    )
+    workspace = torch.empty(workspace_size, dtype=torch.uint8)
+
+    def run(cache):
+        out = torch.empty(b, 128, dv)
+        lse = torch.empty(b, 128, dtype=torch.float32)
+        ops.sm100_cutlass_mla_decode(
+            out,
+            lse,
+            q_nope,
+            q_pe,
+            cache,
+            cache_seqlens,
+            block_table,
+            workspace,
+            scale,
+            1,
+        )
+        return out, lse
+
+    out_contig, lse_contig = run(kv_contig)
+    out_view, lse_view = run(kv_view)
+
+    # Same data and same compute order; only addressing differs.
+    assert torch.equal(out_contig, out_view)
+    assert torch.equal(lse_contig, lse_view)

tests/kernels/attention/test_minimax_m3.py

@@ -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,

tests/kernels/attention/test_mla_cross_layer_kernel_equivalence.py

@@ -0,0 +1,566 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bit-exact kernel equivalence for MLA decode/write kernels on the
+cross-layer (block-major) KV cache layout.
+
+The cross-layer layout carves each layer's per-block page out of a single
+unified slot, so the per-layer view has an inflated ``stride(0)`` (the full
+unified slot) and a non-zero storage offset. These tests confirm the MLA
+kernels behind the backends that opt in to the layout (FlashMLA dense,
+FlashInfer MLA dense, FlashMLA fp8 sparse, plus the ``concat_and_cache_mla``
+write) honor that strided view bit-identically to a contiguous per-layer
+cache, and that writes do not bleed into neighbouring layers' segments.
+"""
+
+import pytest
+import torch
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="MLA cache kernels require CUDA"
+)
+
+
+def test_concat_and_cache_mla_into_unified_slot_view():
+    """concat_and_cache_mla must write correctly into a per-layer view whose
+    block stride is the full unified slot (block-major), with zero bleed into
+    the other layers' segments of the same slot."""
+    from vllm import _custom_ops as ops
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    kv_lora_rank = 512
+    pe = 64
+    entry = kv_lora_rank + pe
+    page = 64
+    num_blocks = 32
+    ntok = 200
+
+    kv_c = torch.randn(ntok, kv_lora_rank, device=dev, dtype=torch.bfloat16)
+    k_pe = torch.randn(ntok, pe, device=dev, dtype=torch.bfloat16)
+    slot = torch.randperm(num_blocks * page, device=dev, dtype=torch.int64)[:ntok]
+    scale = torch.tensor(1.0, device=dev)
+
+    def write(cache):
+        ops.concat_and_cache_mla(kv_c, k_pe, cache, slot, "auto", scale)
+
+    # Contiguous per-layer reference: (num_blocks, page, entry).
+    ref = torch.zeros(num_blocks, page, entry, device=dev, dtype=torch.bfloat16)
+    write(ref)
+
+    # Unified slot holding three layer pages per block. Carve the middle
+    # layer's view (non-zero offset, block stride == full unified slot).
+    layer_page_elems = page * entry
+    n_layers = 3
+    unified_slot_elems = n_layers * layer_page_elems
+    big = torch.zeros(num_blocks, unified_slot_elems, device=dev, dtype=torch.bfloat16)
+    flat = big.view(-1)
+    offset = layer_page_elems  # middle layer
+    view = torch.as_strided(
+        flat,
+        size=(num_blocks, page, entry),
+        stride=(unified_slot_elems, entry, 1),
+        storage_offset=offset,
+    )
+    assert not view.is_contiguous()
+    assert view.stride(0) == unified_slot_elems
+    write(view)
+
+    # Bit-exact equivalence and zero bleed into the neighbour segments.
+    max_diff = (ref.float() - view.float()).abs().max().item()
+    assert max_diff == 0.0, f"max|Δ| = {max_diff}"
+
+    neighbour_lo = torch.as_strided(
+        flat, (num_blocks, layer_page_elems), (unified_slot_elems, 1), 0
+    )
+    neighbour_hi = torch.as_strided(
+        flat,
+        (num_blocks, layer_page_elems),
+        (unified_slot_elems, 1),
+        2 * layer_page_elems,
+    )
+    assert neighbour_lo.abs().max().item() == 0.0
+    assert neighbour_hi.abs().max().item() == 0.0
+
+
+def test_flashmla_dense_decode_unified_slot_view():
+    """FlashMLA dense decode (FLASHMLA backend, e.g. Kimi-K2-style dense MLA
+    on Hopper) must read a unified-slot block-major view bit-identically to a
+    contiguous per-layer cache."""
+    import vllm.v1.attention.ops.flashmla as fm
+
+    ok, reason = fm.is_flashmla_dense_supported()
+    if not ok:
+        pytest.skip(reason)
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    dt = torch.bfloat16
+    head_dim = 576
+    hdv = 512
+    h_q = 128
+    page = 64
+    num_blocks = 64
+    bs = 4
+    n_layers = 3
+    layer = 1
+
+    q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=dt) * 0.1
+    kv_data = torch.randn(num_blocks, page, 1, head_dim, device=dev, dtype=dt) * 0.1
+
+    # (A) contiguous per-layer reference.
+    cache_contiguous = kv_data.clone().contiguous()
+
+    # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
+    unified = (
+        torch.randn(num_blocks, n_layers, page, 1, head_dim, device=dev, dtype=dt) * 0.1
+    )
+    unified[:, layer].copy_(kv_data)
+    cache_view = unified[:, layer]
+    assert not cache_view.is_contiguous()
+    assert cache_view.stride(0) == n_layers * page * 1 * head_dim
+
+    max_blk = num_blocks // bs
+    block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
+        bs, max_blk
+    )
+    cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
+
+    def run(kc):
+        meta, num_splits = fm.get_mla_metadata()
+        out, _ = fm.flash_mla_with_kvcache(
+            q=q,
+            k_cache=kc,
+            block_table=block_table,
+            cache_seqlens=cache_seqlens,
+            head_dim_v=hdv,
+            tile_scheduler_metadata=meta,
+            num_splits=num_splits,
+            softmax_scale=head_dim**-0.5,
+            causal=True,
+        )
+        return out.clone().float()
+
+    out_ref = run(cache_contiguous)
+    out_view = run(cache_view)
+    assert torch.isfinite(out_ref).all()
+    assert out_ref.abs().max().item() > 0.0
+    assert (out_ref - out_view).abs().max().item() == 0.0
+
+
+def test_flashinfer_mla_dense_decode_unified_slot_view():
+    """FlashInfer MLA dense decode must read a unified-slot block-major view
+    (inflated stride(0), non-zero storage offset) bit-identically to a
+    contiguous per-layer cache."""
+    try:
+        from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+    except ImportError:
+        pytest.skip("flashinfer is not available")
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    dt = torch.bfloat16
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+    qk_nope_head_dim = 128
+    head_dim = kv_lora_rank + qk_rope_head_dim  # 576
+    num_qo_heads = 128
+    page = 64
+    num_blocks = 64
+    bs = 4
+    n_layers = 3  # >1 so the per-layer view's block stride is inflated.
+    layer = 1
+
+    q = torch.randn(bs, 1, num_qo_heads, head_dim, device=dev, dtype=dt)
+    kv_data = torch.randn(num_blocks, 1, page, head_dim, device=dev, dtype=dt)
+
+    # (A) contiguous per-layer reference.
+    kv_contiguous = kv_data.clone().contiguous()
+
+    # (B) unified slot: block b of every layer packed together; view one layer
+    # -> stride(0) is n_layers x larger and storage offset is non-zero.
+    unified = torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev, dtype=dt)
+    unified[:, layer].copy_(kv_data)
+    kv_view = unified[:, layer]
+    assert not kv_view.is_contiguous()
+    assert kv_view.stride(0) == n_layers * 1 * page * head_dim
+
+    max_blk = num_blocks // bs
+    block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
+        bs, max_blk
+    )
+    seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
+    ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
+    scale = head_dim**-0.5
+
+    def run(kv):
+        return trtllm_batch_decode_with_kv_cache_mla(
+            query=q,
+            kv_cache=kv,
+            workspace_buffer=ws,
+            qk_nope_head_dim=qk_nope_head_dim,
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            block_tables=block_tables,
+            seq_lens=seq_lens,
+            max_seq_len=int(seq_lens.max().item()),
+            bmm1_scale=scale,
+            bmm2_scale=1.0,
+        ).clone()
+
+    out_ref = run(kv_contiguous).float()
+    out_view = run(kv_view).float()
+    assert torch.isfinite(out_ref).all()
+    assert (out_ref - out_view).abs().max().item() == 0.0
+
+
+def test_flashmla_fp8_sparse_decode_unified_slot_view():
+    """FlashMLA fp8 sparse decode (DeepSeek V3.2/V4 DSA path) must read a
+    unified-slot block-major view bit-identically to a contiguous fp8_ds_mla
+    cache, with finite nonzero output."""
+    import vllm.v1.attention.ops.flashmla as fm
+
+    ok, reason = fm.is_flashmla_sparse_supported()
+    if not ok:
+        pytest.skip(reason)
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    entry = 656  # fp8_ds_mla bytes per token
+    page = 64
+    num_blocks = 32
+    h_q = 128
+    head_dim = 576
+    hdv = 512
+    batch = 2
+    topk = 128
+    n_layers = 3
+    layer = 1
+
+    q = torch.randn(batch, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
+
+    # Structurally valid fp8 ds_mla payload: 512B fp8 + 16B f32 scales + 128B
+    # bf16 rope (random bytes corrupt the scale region and yield NaNs).
+    nope = (torch.randn(num_blocks, page, 1, 512, device=dev) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    scales = torch.ones(num_blocks, page, 1, 4, device=dev, dtype=torch.float32)
+    rope = (torch.randn(num_blocks, page, 1, 64, device=dev) * 0.1).to(torch.bfloat16)
+    payload = torch.cat(
+        [
+            nope.view(torch.uint8).view(num_blocks, page, 1, 512),
+            scales.view(torch.uint8).view(num_blocks, page, 1, 16),
+            rope.view(torch.uint8).view(num_blocks, page, 1, 128),
+        ],
+        dim=-1,
+    ).contiguous()
+    assert payload.shape[-1] == entry and payload.dtype == torch.uint8
+
+    # (A) contiguous reference.
+    cache_contiguous = payload.clone().contiguous()
+
+    # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
+    unified = torch.randint(
+        0, 256, (num_blocks, n_layers, page, 1, entry), device=dev, dtype=torch.uint8
+    )
+    unified[:, layer].copy_(payload)
+    cache_view = unified[:, layer]
+    assert not cache_view.is_contiguous()
+    assert cache_view.stride(0) == n_layers * page * 1 * entry
+
+    # Sparse indices: each batch uses its own disjoint blocks.
+    blocks_per_batch = num_blocks // batch
+    idx = torch.full((batch, 1, topk), -1, device=dev, dtype=torch.int32)
+    for b in range(batch):
+        slots: list[int] = []
+        for blk in range(b * blocks_per_batch, (b + 1) * blocks_per_batch):
+            slots.extend(blk * page + off for off in range(page))
+        slots_t = torch.tensor(slots[:topk], device=dev, dtype=torch.int32)
+        idx[b, 0, : slots_t.numel()] = slots_t
+
+    def run(kc):
+        meta, num_splits = fm.get_mla_metadata()
+        out, _ = fm.flash_mla_with_kvcache(
+            q=q,
+            k_cache=kc,
+            block_table=None,
+            cache_seqlens=None,
+            head_dim_v=hdv,
+            tile_scheduler_metadata=meta,
+            is_fp8_kvcache=True,
+            indices=idx,
+            softmax_scale=head_dim**-0.5,
+        )
+        return out.clone().float()
+
+    out_ref = run(cache_contiguous)
+    out_view = run(cache_view)
+    assert torch.isfinite(out_ref).all()
+    assert out_ref.abs().max().item() > 0.0
+    assert (out_ref - out_view).abs().max().item() == 0.0
+
+
+def test_indexer_k_quant_and_cache_into_unified_slot_view():
+    """indexer_k_quant_and_cache (DeepSeek V3.2/V4 DSA indexer K write) must
+    write correctly into a per-layer view whose block stride is the full
+    unified slot, with zero bleed into the other layers' segments."""
+    from vllm import _custom_ops as ops
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    head_dim = 128
+    quant_block_size = 128
+    block_size = 64
+    num_blocks = 16
+    ntok = 100
+    # Indexer cache layout per token: head_dim fp8 bytes followed by
+    # head_dim * 4 / quant_block_size scale bytes.
+    cache_stride = head_dim + head_dim * 4 // quant_block_size
+
+    k = torch.randn(ntok, head_dim, device=dev, dtype=torch.bfloat16)
+    slot = torch.randperm(num_blocks * block_size, device=dev, dtype=torch.int64)[:ntok]
+
+    def write(cache):
+        ops.indexer_k_quant_and_cache(k, cache, slot, quant_block_size, "ue8m0")
+
+    # Contiguous per-layer reference.
+    ref = torch.zeros(
+        num_blocks, block_size, cache_stride, device=dev, dtype=torch.uint8
+    )
+    write(ref)
+
+    # Unified slot holding three layer pages per block; carve the middle one.
+    n_layers = 3
+    layer = 1
+    unified = torch.zeros(
+        num_blocks, n_layers, block_size, cache_stride, device=dev, dtype=torch.uint8
+    )
+    view = unified[:, layer]
+    assert not view.is_contiguous()
+    assert view.stride(0) == n_layers * block_size * cache_stride
+    write(view)
+
+    assert torch.equal(ref, view.contiguous())
+    # Zero bleed into the neighbour layers' segments.
+    assert unified[:, 0].abs().max().item() == 0
+    assert unified[:, 2].abs().max().item() == 0
+
+
+def test_flashattn_mla_dense_decode_unified_slot_view():
+    """FA3 decode (FLASH_ATTN_MLA backend) must read a unified-slot
+    block-major view bit-identically to a contiguous per-layer cache."""
+    try:
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+    except ImportError:
+        pytest.skip("vllm_flash_attn is not available")
+    from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
+
+    if not flash_attn_supports_mla():
+        pytest.skip("FA3 MLA requires a Hopper device")
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    dt = torch.bfloat16
+    kv_lora_rank = 512
+    rope_dim = 64
+    entry = kv_lora_rank + rope_dim  # 576
+    h_q = 16
+    page = 64
+    num_blocks = 64
+    bs = 4
+    n_layers = 3
+    layer = 1
+
+    q_pe = torch.randn(bs, h_q, rope_dim, device=dev, dtype=dt) * 0.1
+    q_nope = torch.randn(bs, h_q, kv_lora_rank, device=dev, dtype=dt) * 0.1
+    kv_data = torch.randn(num_blocks, page, entry, device=dev, dtype=dt) * 0.1
+
+    # (A) contiguous per-layer reference.
+    cache_contiguous = kv_data.clone().contiguous()
+
+    # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
+    unified = torch.randn(num_blocks, n_layers, page, entry, device=dev, dtype=dt) * 0.1
+    unified[:, layer].copy_(kv_data)
+    cache_view = unified[:, layer]
+    assert not cache_view.is_contiguous()
+    assert cache_view.stride(0) == n_layers * page * entry
+
+    max_blk = num_blocks // bs
+    block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
+        bs, max_blk
+    )
+    seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
+    cu_seqlens_q = torch.arange(bs + 1, device=dev, dtype=torch.int32)
+
+    def run(cache):
+        kv_c_cache = cache[..., :kv_lora_rank]
+        k_pe_cache = cache[..., kv_lora_rank:]
+        out = flash_attn_varlen_func(
+            q=q_pe,
+            k=k_pe_cache.unsqueeze(-2),  # Add head dim of 1
+            v=kv_c_cache.unsqueeze(-2),  # Add head dim of 1
+            q_v=q_nope,
+            max_seqlen_q=1,
+            cu_seqlens_q=cu_seqlens_q,
+            max_seqlen_k=int(seq_lens.max().item()),
+            seqused_k=seq_lens,
+            block_table=block_table,
+            softmax_scale=entry**-0.5,
+            causal=True,
+            fa_version=3,
+        )
+        return out.clone().float()
+
+    out_ref = run(cache_contiguous)
+    out_view = run(cache_view)
+    assert torch.isfinite(out_ref).all()
+    assert out_ref.abs().max().item() > 0.0
+    assert (out_ref - out_view).abs().max().item() == 0.0
+
+
+def test_flashmla_dense_fp8_decode_unified_slot_view():
+    """FlashMLA dense fp8 decode (FLASHMLA backend with quantized KV cache)
+    must read a unified-slot block-major view bit-identically to a contiguous
+    per-layer fp8 cache."""
+    import vllm.v1.attention.ops.flashmla as fm
+
+    ok, reason = fm.is_flashmla_dense_supported()
+    if not ok:
+        pytest.skip(reason)
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    head_dim = 576
+    hdv = 512
+    h_q = 128
+    page = 64
+    num_blocks = 64
+    bs = 4
+    n_layers = 3
+    layer = 1
+
+    q = torch.randn(bs, 1, h_q, head_dim, device=dev, dtype=torch.bfloat16) * 0.1
+    kv_data = (torch.randn(num_blocks, page, head_dim, device=dev) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+
+    # (A) contiguous per-layer reference.
+    cache_contiguous = kv_data.clone().contiguous()
+
+    # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
+    unified = (torch.randn(num_blocks, n_layers, page, head_dim, device=dev) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    unified[:, layer].copy_(kv_data)
+    cache_view = unified[:, layer]
+    assert not cache_view.is_contiguous()
+    assert cache_view.stride(0) == n_layers * page * head_dim
+
+    max_blk = num_blocks // bs
+    block_table = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
+        bs, max_blk
+    )
+    cache_seqlens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
+    descale = torch.ones(1, device=dev, dtype=torch.float32)
+
+    def run(kc):
+        tile_md, num_splits = fm.get_mla_metadata_dense_fp8(cache_seqlens, h_q, 1)
+        out, _ = fm.flash_mla_with_kvcache_fp8(
+            q=q,
+            k_cache=kc.unsqueeze(-2),  # Add head dim of 1
+            block_table=block_table,
+            cache_seqlens=cache_seqlens,
+            head_dim_v=hdv,
+            tile_scheduler_metadata=tile_md,
+            num_splits=num_splits,
+            softmax_scale=head_dim**-0.5,
+            causal=True,
+            descale_q=descale,
+            descale_k=descale,
+        )
+        return out.clone().float()
+
+    out_ref = run(cache_contiguous)
+    out_view = run(cache_view)
+    assert torch.isfinite(out_ref).all()
+    assert out_ref.abs().max().item() > 0.0
+    assert (out_ref - out_view).abs().max().item() == 0.0
+
+
+def test_flashinfer_mla_dense_fp8_decode_unified_slot_view():
+    """FlashInfer MLA dense decode with an fp8 KV cache must read a
+    unified-slot block-major view bit-identically to a contiguous per-layer
+    cache."""
+    try:
+        from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
+    except ImportError:
+        pytest.skip("flashinfer is not available")
+    from vllm.platforms import current_platform
+
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("FlashInfer trtllm-gen MLA requires sm100")
+
+    torch.manual_seed(0)
+    dev = "cuda"
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+    qk_nope_head_dim = 128
+    head_dim = kv_lora_rank + qk_rope_head_dim  # 576
+    num_qo_heads = 128
+    page = 64
+    num_blocks = 64
+    bs = 4
+    n_layers = 3
+    layer = 1
+
+    # With a quantized KV cache the decode query is quantized to fp8 as well
+    # (trtllm-gen has no bf16-query x fp8-cache decode kernel).
+    q = (torch.randn(bs, 1, num_qo_heads, head_dim, device=dev) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    kv_data = (torch.randn(num_blocks, 1, page, head_dim, device=dev) * 0.1).to(
+        torch.float8_e4m3fn
+    )
+
+    # (A) contiguous per-layer reference.
+    kv_contiguous = kv_data.clone().contiguous()
+
+    # (B) unified slot: view one layer -> inflated stride(0), non-zero offset.
+    unified = (
+        torch.randn(num_blocks, n_layers, 1, page, head_dim, device=dev) * 0.1
+    ).to(torch.float8_e4m3fn)
+    unified[:, layer].copy_(kv_data)
+    kv_view = unified[:, layer]
+    assert not kv_view.is_contiguous()
+    assert kv_view.stride(0) == n_layers * 1 * page * head_dim
+
+    max_blk = num_blocks // bs
+    block_tables = torch.arange(num_blocks, device=dev, dtype=torch.int32).view(
+        bs, max_blk
+    )
+    seq_lens = torch.full((bs,), max_blk * page, device=dev, dtype=torch.int32)
+    ws = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device=dev)
+    scale = head_dim**-0.5
+
+    def run(kv):
+        return trtllm_batch_decode_with_kv_cache_mla(
+            query=q,
+            kv_cache=kv,
+            workspace_buffer=ws,
+            qk_nope_head_dim=qk_nope_head_dim,
+            kv_lora_rank=kv_lora_rank,
+            qk_rope_head_dim=qk_rope_head_dim,
+            block_tables=block_tables,
+            seq_lens=seq_lens,
+            max_seq_len=int(seq_lens.max().item()),
+            bmm1_scale=scale,
+            bmm2_scale=1.0,
+        ).clone()
+
+    out_ref = run(kv_contiguous).float()
+    out_view = run(kv_view).float()
+    assert torch.isfinite(out_ref).all()
+    assert (out_ref - out_view).abs().max().item() == 0.0

tests/kernels/attention/test_triton_decode_attention.py

@@ -231,3 +231,95 @@ def test_decode_attention_fp8(B, L, H_Q, H_KV, D_QK, D_V, CACHE_SIZE, PAGE_SIZE)
 
     # FP8 tolerances match test_mla_backends.py test_backend_correctness.
     torch.testing.assert_close(o_ref, o_fp8, atol=5e-1, rtol=1e-2)
+
+
+@pytest.mark.parametrize(
+    "H_Q,H_KV,D_QK,D_V,is_mla",
+    [
+        (16, 1, 576, 512, True),  # MLA path (grouped kernel, v = trans(k))
+        (32, 8, 128, 128, False),  # GQA path (grouped kernel)
+        (32, 32, 128, 128, False),  # MHA path (normal kernel)
+    ],
+)
+@pytest.mark.parametrize("PAGE_SIZE", [16])
+def test_decode_attention_cross_layer_view(H_Q, H_KV, D_QK, D_V, is_mla, PAGE_SIZE):
+    """The kernel must honor the cache's page-dim stride, not assume pages are
+    packed back-to-back. A per-layer view into a cross-layer (block-major)
+    cache has stride(0) inflated by num_layers; outputs must match a
+    contiguous cache holding the same data exactly."""
+    B = 3
+    seq_len = 1027
+    CACHE_SIZE = 16384
+    NUM_LAYERS = 3
+    LAYER_IDX = 1
+    dtype = torch.bfloat16
+    sm_scale = 1.0 / (D_QK**0.5)
+    num_kv_splits = 8
+    num_pages = CACHE_SIZE // PAGE_SIZE
+
+    num_pages_per_batch = cdiv(seq_len, PAGE_SIZE)
+    req_to_page = torch.randint(
+        0, num_pages, (B, num_pages_per_batch), device=DEVICE_TYPE
+    )
+
+    q = torch.randn(B, H_Q, D_QK, dtype=dtype, device=DEVICE_TYPE)
+    b_seq_len = torch.full((B,), seq_len, device=DEVICE_TYPE)
+
+    # Reference: contiguous paged cache.
+    k_ref = torch.randn(
+        num_pages, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
+    )
+    if is_mla:
+        v_ref = k_ref[..., :D_V]
+    else:
+        v_ref = torch.randn(
+            num_pages, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
+        )
+
+    # Cross-layer cache: all layers' pages for a block are adjacent. The
+    # per-layer view has the same shape as the contiguous cache but
+    # stride(0) is NUM_LAYERS x larger. Neighbor layers hold random data so
+    # any packed-pages addressing reads garbage rather than zeros.
+    k_xl = torch.randn(
+        num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_QK, dtype=dtype, device=DEVICE_TYPE
+    )
+    k_view = k_xl[:, LAYER_IDX]
+    k_view.copy_(k_ref)
+    assert k_view.stride(0) == NUM_LAYERS * PAGE_SIZE * H_KV * D_QK
+    if is_mla:
+        v_view = k_view[..., :D_V]
+    else:
+        v_xl = torch.randn(
+            num_pages, NUM_LAYERS, PAGE_SIZE, H_KV, D_V, dtype=dtype, device=DEVICE_TYPE
+        )
+        v_view = v_xl[:, LAYER_IDX]
+        v_view.copy_(v_ref)
+
+    def run(k_buffer, v_buffer):
+        o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=DEVICE_TYPE)
+        lse = torch.zeros(B, H_Q, dtype=dtype, device=DEVICE_TYPE)
+        attn_logits = torch.empty(
+            (B, H_Q, num_kv_splits, D_V + 1), dtype=torch.float32, device=DEVICE_TYPE
+        )
+        decode_attention_fwd(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            lse,
+            req_to_page,
+            b_seq_len,
+            attn_logits,
+            num_kv_splits,
+            sm_scale,
+            PAGE_SIZE,
+            is_mla=is_mla,
+        )
+        return o, lse
+
+    o_ref, lse_ref = run(k_ref, v_ref)
+    o_xl, lse_xl = run(k_view, v_view)
+
+    # Same data and same compute order; only addressing differs.
+    assert torch.equal(o_ref, o_xl)
+    assert torch.equal(lse_ref, lse_xl)

tests/kernels/moe/test_moe_weight_loading_padded.py

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

tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py

@@ -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():

tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

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

tests/kernels/test_top_k_per_row.py

@@ -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}"
             )

tests/lora/test_lora_manager.py

@@ -124,6 +124,7 @@ def test_replace_submodules(default_vllm_config, dist_init, dummy_model):
             max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
         ),
         torch.device(DEVICES[0]),
+        default_vllm_config,
     )
     model = manager.model
     assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)
@@ -152,6 +153,7 @@ def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy
             max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
         ),
         torch.device(DEVICES[0]),
+        default_vllm_config,
     )
 
     assert isinstance(
@@ -172,6 +174,7 @@ def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):
             max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
         ),
         torch.device(DEVICES[0]),
+        default_vllm_config,
     )
 
     assert isinstance(
@@ -219,6 +222,7 @@ def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_
             max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
         ),
         torch.device(DEVICES[0]),
+        default_vllm_config,
     )
 
     canonical = manager.model.get_submodule("moe.gate")
@@ -263,6 +267,7 @@ def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model):
             max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
         ),
         torch.device(DEVICES[0]),
+        default_vllm_config,
     )
 
     # lm_head's special handling still ran: logits_processor got wrapped
@@ -293,6 +298,7 @@ def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_
             max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
         ),
         torch.device(DEVICES[0]),
+        default_vllm_config,
     )
 
     # Should not crash and should keep unsupported matched modules unchanged.
@@ -325,6 +331,7 @@ def test_target_modules_fail_closed_on_unsupported_matched_modules(
                 target_modules=["dense1"],
             ),
             torch.device(DEVICES[0]),
+            default_vllm_config,
         )
 
 
@@ -374,6 +381,7 @@ def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device)
             max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
         ),
         device=device,
+        vllm_config=default_vllm_config,
     )
     assert all(x is None for x in manager.lora_index_to_id)
     assert manager.add_adapter(model_lora1)
@@ -442,6 +450,7 @@ def test_lora_lru_cache_model_manager(
             max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
         ),
         device=device,
+        vllm_config=default_vllm_config,
     )
     assert all(x is None for x in manager.lora_index_to_id)
     assert manager.add_adapter(model_lora1)
@@ -535,6 +544,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev
             max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
         ),
         device=device,
+        vllm_config=default_vllm_config,
     )
     assert all(x is None for x in manager.lora_index_to_id)
 
@@ -642,9 +652,7 @@ def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, dev
 
 
 @pytest.mark.parametrize("device", DEVICES)
-def test_lru_cache_worker_adapter_manager(
-    default_vllm_config, dist_init, dummy_model, device, tmp_path
-):
+def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):
     lora_config = LoRAConfig(
         max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
     )
@@ -670,7 +678,7 @@ def test_lru_cache_worker_adapter_manager(
     worker_adapter_manager.max_num_seqs = 4
     worker_adapter_manager.max_num_batched_tokens = 2
 
-    worker_adapter_manager.create_lora_manager(dummy_model)
+    worker_adapter_manager.create_lora_manager(dummy_model, vllm_config)
 
     mapping = LoRAMapping([], [])
     worker_adapter_manager.set_active_adapters(
@@ -758,9 +766,7 @@ def test_lru_cache_worker_adapter_manager(
 
 
 @pytest.mark.parametrize("device", DEVICES)
-def test_worker_adapter_manager(
-    default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path
-):
+def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):
     # Should remove every LoRA not specified in the request.
     lora_config = LoRAConfig(
         max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
@@ -774,7 +780,7 @@ def test_worker_adapter_manager(
 
     worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
     worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
-    worker_adapter_manager.create_lora_manager(dummy_model_gate_up)
+    worker_adapter_manager.create_lora_manager(dummy_model_gate_up, vllm_config)
 
     dummy_lora_files = f"{tmp_path}/lora_adapter"
     os.makedirs(dummy_lora_files, exist_ok=True)
@@ -894,6 +900,7 @@ def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, devic
             max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
         ),
         device=device,
+        vllm_config=default_vllm_config,
     )
     model = manager.model
 
@@ -944,6 +951,7 @@ def _test_target_modules(
     device: str,
     expected_lora: list[tuple[str, type]],
     expected_no_lora: list[tuple[str, type]],
+    vllm_config,
 ):
     """Create a LoRAModelManager and assert which modules have LoRA applied."""
     LoRAModelManager(
@@ -959,6 +967,7 @@ def _test_target_modules(
             target_modules=target_modules,
         ),
         device=device,
+        vllm_config=vllm_config,
     )
     for module_path, lora_cls in expected_lora:
         assert isinstance(model.get_submodule(module_path), lora_cls)
@@ -981,6 +990,7 @@ def test_target_modules_config(default_vllm_config, dist_init, dummy_model, devi
             ("dense2", RowParallelLinearWithLoRA),
             ("layer1.dense2", RowParallelLinearWithLoRA),
         ],
+        vllm_config=default_vllm_config,
     )
 
 
@@ -998,6 +1008,7 @@ def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, de
             ("layer1.dense2", RowParallelLinearWithLoRA),
         ],
         expected_no_lora=[],
+        vllm_config=default_vllm_config,
     )
 
 
@@ -1017,6 +1028,7 @@ def test_target_modules_none_uses_all(
             ("layer1.dense2", RowParallelLinearWithLoRA),
         ],
         expected_no_lora=[],
+        vllm_config=default_vllm_config,
     )
 
 
@@ -1036,4 +1048,5 @@ def test_target_modules_match_packed_runtime_modules(
             ("layer1.dense1", ColumnParallelLinearWithLoRA),
             ("layer1.dense2", RowParallelLinearWithLoRA),
         ],
+        vllm_config=default_vllm_config,
     )

tests/model_executor/test_flashinfer_autotune_cache.py

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

tests/models/language/pooling/test_gritlm.py

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

tests/models/multimodal/generation/test_moss_audio.py

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