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

Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com> (cherry picked from commit 798185d438)
[CI] Raise gsm8k startup timeout for MoE Refactor Qwen3 NVFP4 configs (#46882 )
2026-06-28 00:04:08 -07:00 · 2026-06-26 16:33:47 -07:00 · 2026-06-26 16:15:00 -07:00 · 2026-06-25 13:15:20 -07:00 · 2026-06-25 13:15:20 -07:00 · 2026-06-25 13:15:02 -07:00
479 changed files with 22702 additions and 7477 deletions
@@ -136,7 +136,7 @@ steps:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 3
+  parallelism: 4

 - label: "Arm CPU Test"
  depends_on: []
@@ -21,6 +21,10 @@ steps:
    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -38,6 +42,10 @@ steps:
    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -55,6 +63,10 @@ steps:
    timeout_in_minutes: 30
    optional: true
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -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 \
@@ -5,6 +5,10 @@ steps:
 - label: XPU Sleep Mode
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -5,6 +5,10 @@ steps:
 - label: Engine (1 GPU)
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -1,11 +1,15 @@
 group: Expert Parallelism
-depends_on: 
+depends_on:
  - image-build-xpu
 steps:
 - label: EPLB Algorithm
  key: eplb-algorithm
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -5,6 +5,10 @@ steps:
 - label: vLLM IR Tests
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -5,6 +5,10 @@ steps:
 - label: LoRA Runtime + Utils
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
@@ -34,6 +38,10 @@ steps:
 - label: LoRA Fused/MoE Kernels
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -54,6 +62,10 @@ steps:
 - label: LoRA Punica Kernels
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -74,6 +86,10 @@ steps:
 - label: LoRA Punica FP8/XPU Ops
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -94,6 +110,10 @@ steps:
 - label: LoRA Models
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
@@ -117,6 +137,10 @@ steps:
 - label: LoRA Multimodal
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -5,6 +5,10 @@ steps:
 - label: V1 Core + KV + Metrics
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -31,6 +35,10 @@ steps:
 - label: V1 Sample + Logits
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -71,6 +79,10 @@ steps:
 - label: XPU CPU Offload
  timeout_in_minutes: 60
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -95,6 +107,10 @@ steps:
  key: regression
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -126,6 +142,10 @@ steps:
  timeout_in_minutes: 30
  num_devices: 2
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -157,6 +177,10 @@ steps:
  key: async-engine-inputs-utils-worker
  timeout_in_minutes: 30
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
@@ -5,6 +5,10 @@ steps:
 - label: Model Runner V2 Core Tests (Intel)
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 2+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -30,6 +34,10 @@ steps:
 - label: Model Runner V2 Examples (Intel)
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
@@ -1,11 +1,15 @@
 group: Models - Multimodal
-depends_on: 
+depends_on:
  - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
  key: multi-modal-models-standard-1-qwen2
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -27,6 +31,10 @@ steps:
  key: multi-modal-models-standard-2-qwen3-gemma
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -47,6 +55,10 @@ steps:
  key: multi-modal-models-standard-3-llava-qwen2-vl
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 24+
  no_plugin: true
  working_dir: "."
  env:
@@ -68,6 +80,10 @@ steps:
  key: multi-modal-models-standard-4-other-whisper
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -88,6 +104,10 @@ steps:
  key: multi-modal-processor
  timeout_in_minutes: 45
  device: intel_gpu
+  agent_tags:
+    label: production
+    gpu: 1+
+    mem: 16+
  no_plugin: true
  working_dir: "."
  env:
@@ -19,6 +19,10 @@ steps:
      - image-build-xpu
    timeout_in_minutes: 30
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 2+
+      mem: 24+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -49,6 +53,10 @@ steps:
      - image-build-xpu
    timeout_in_minutes: 30
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -74,6 +82,10 @@ steps:
      - image-build-xpu
    timeout_in_minutes: 30
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -93,6 +105,10 @@ steps:
      - image-build-xpu
    timeout_in_minutes: 30
    device: intel_gpu
+    agent_tags:
+      label: production
+      gpu: 1+
+      mem: 16+
    no_plugin: true
    env:
      REGISTRY: "public.ecr.aws/q9t5s3a7"
@@ -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)
@@ -4,6 +4,11 @@

 set -euo pipefail

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

 uv pip install --system -r "${REQUIREMENTS_FILE}"
@@ -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]
@@ -647,7 +439,7 @@ steps:
  - pytest -v -s v1/cudagraph/test_cudagraph_mode.py

 - label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 35
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
  optional: true
@@ -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]
@@ -1594,9 +1405,10 @@ steps:
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#

 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 55
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1627,10 +1439,11 @@ steps:
  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py

 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  parallelism: 4
+  optional: true
+  parallelism: 5
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -1740,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
@@ -1768,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
@@ -2075,19 +1956,6 @@ steps:
  - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
  - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test

- label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/
-  - tests/v1/e2e/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
-
 - label: e2e Scheduling (1 GPU) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2133,9 +2001,10 @@ steps:
  - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

 - label: Spec Decode Eagle # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/v1/spec_decode/
@@ -2346,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]
@@ -3053,7 +2937,7 @@ steps:
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#

 - label: Kernels (B200-MI355) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/"
@@ -3077,11 +2961,10 @@ steps:
  - pytest -v -s tests/kernels/attention/test_attention_selector.py

 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  parallelism: 2
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/attention/
@@ -3095,10 +2978,10 @@ steps:
  - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
-  parallelism: 4
+  parallelism: 5
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -3422,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
@@ -74,6 +74,16 @@ steps:
    - tests/v1/e2e/general/
  commands:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+  mirror:
+    amd:
+      device: mi250_1
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/
+      - tests/v1/e2e/general/
+      - vllm/platforms/rocm.py

 - label: V1 e2e (2 GPUs)
  key: v1-e2e-2-gpus
@@ -47,8 +47,10 @@ steps:
  - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
  - vllm/models/deepseek_v4/common/ops/
  - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
+  - tests/kernels/test_top_k_per_row.py # it runs on Blackwell too - some kernels have arch-specific optimizations
  commands:
    - pytest -v -s kernels/test_fused_deepseek_v4_*.py
+    - pytest -v -s kernels/test_top_k_per_row.py

 - label: Deepseek V4 Kernel Test (B200)
  key: deepseek-v4-kernel-test-b200
@@ -74,6 +76,20 @@ steps:
  commands:
    - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
  parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py

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

 - label: Kernels Mamba Test
  key: kernels-mamba-test
@@ -101,6 +101,16 @@ steps:
  num_devices: 8
  commands:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
+  mirror:
+    amd:
+      device: mi300_8
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt

 - label: MoE Refactor Integration Test (H100 - TEMPORARY)
  key: moe-refactor-integration-test-h100-temporary
@@ -126,6 +136,49 @@ steps:
  commands:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt

+- label: LM Eval Humming (A100 - TEMPORARY)
+  key: lm-eval-humming-a100
+  timeout_in_minutes: 30
+  device: a100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+
+- label: LM Eval Humming (H100 - TEMPORARY)
+  key: lm-eval-humming-h100
+  timeout_in_minutes: 30
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
+
+- label: LM Eval Humming (B200 - TEMPORARY)
+  key: lm-eval-humming-b200
+  timeout_in_minutes: 30
+  device: b200-k8s
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt

 - label: LM Eval TurboQuant KV Cache
  key: lm-eval-turboquant-kv-cache
@@ -167,6 +220,20 @@ steps:
    - uv pip install --system 'gpt-oss[eval]==0.0.5'
    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt

+- label: GPQA Eval (GPT-OSS) (DGX Spark)
+  key: gpqa-eval-gpt-oss-spark
+  timeout_in_minutes: 120
+  device: dgx-spark
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - tests/evals/gpt_oss/
+  commands:
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
+
 - label: MRCR Eval Small Models
  device: h200_35gb
  timeout_in_minutes: 30
@@ -105,6 +105,12 @@ steps:
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 60
+      depends_on:
+      - image-build-amd

 - label: V1 Others (CPU)
  key: v1-others-cpu
@@ -68,7 +68,6 @@ steps:
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
  mirror:
    amd:
-      soft_fail: true
      device: mi325_1
      depends_on:
      - image-build-amd
@@ -107,6 +107,12 @@ steps:
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd

 - label: PyTorch Fullgraph Smoke Test
  key: pytorch-fullgraph-smoke-test
@@ -12,6 +12,20 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/worker/gpu/spec_decode/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/sample/
+      - vllm/model_executor/layers/
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py

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

 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
-/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
-/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
+/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/model_loader @22quinn
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/ir @ProExpertProg
@@ -199,7 +199,9 @@ cython_debug/
 .vscode/

 # Claude
-.claude/
+.claude/*
+!.claude/skills/
+!.claude/skills/**

 # Codex
 .codex/
@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
 - Keep comments and docstrings minimal and concise.
 - Assume the reader is familiar with vLLM.

-### Diagnosing CI failures
-
-Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
-
-```bash
-# All failed-job logs for a PR's latest build (current branch's PR if omitted):
-.buildkite/scripts/ci-fetch-log.sh --pr <PR>
-# Any Buildkite build or job URL also works:
-.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
-```
-
 ### Commit messages

 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
@@ -146,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.
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
 endif()

 #
-# _C extension
+# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
 #

-set(VLLM_EXT_SRC
-  "csrc/quantization/activation_kernels.cu"
-  "csrc/torch_bindings.cpp")
-
-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
-
-  # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
-  set(CUTLASS_REVISION "v4.4.2")
-
-  # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
-  if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
-    set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
-  endif()
-
-  if(VLLM_CUTLASS_SRC_DIR)
-    if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
-      get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
-    endif()
-    message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
-    FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
-  else()
-    FetchContent_Declare(
-        cutlass
-        GIT_REPOSITORY https://github.com/nvidia/cutlass.git
-        # Please keep this in sync with CUTLASS_REVISION line above.
-        GIT_TAG ${CUTLASS_REVISION}
-        GIT_PROGRESS TRUE
-
-        # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
-        # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
-        # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
-        GIT_SHALLOW TRUE
-    )
-  endif()
-  FetchContent_MakeAvailable(cutlass)
-
-  set_gencode_flags_for_srcs(
-    SRCS "${VLLM_EXT_SRC}"
-    CUDA_ARCHS "${CUDA_ARCHS}")
-
-# if CUDA endif
-endif()
-
-if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
-  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
-  list(APPEND VLLM_EXT_SRC
+if(VLLM_GPU_LANG STREQUAL "HIP")
+  set(VLLM_EXT_SRC
+    "csrc/torch_bindings.cpp"
    "csrc/custom_quickreduce.cu"
    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu"
-  )
-# if ROCM endif
-endif()
+    "csrc/libtorch_stable/cuda_utils_kernels.cu")

-message(STATUS "Enabling C extension.")
-define_extension_target(
-  _C
-  DESTINATION vllm
-  LANGUAGE ${VLLM_GPU_LANG}
-  SOURCES ${VLLM_EXT_SRC}
-  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-  ARCHITECTURES ${VLLM_GPU_ARCHES}
-  INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
-  INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
-  USE_SABI 3
-  WITH_SOABI)
+  message(STATUS "Enabling C extension.")
+  define_extension_target(
+    _C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${VLLM_EXT_SRC}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
+    INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
+    USE_SABI 3
+    WITH_SOABI)

-# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
-# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
-# driver API. This causes problems when linking with earlier versions of CUDA.
-# Setting this variable sidesteps the issue by calling the driver directly.
-target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  # If CUTLASS is compiled on NVCC >= 12.5, it by default uses
+  # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
+  # driver API. This causes problems when linking with earlier versions of CUDA.
+  # Setting this variable sidesteps the issue by calling the driver directly.
+  target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+endif() # _C HIP endif

 if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  #
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  set(VLLM_STABLE_EXT_SRC
    "csrc/libtorch_stable/torch_bindings.cpp"
    "csrc/libtorch_stable/activation_kernels.cu"
+    "csrc/libtorch_stable/quantization/activation_kernels.cu"
    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
    "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
    "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -428,7 +382,57 @@ 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")
+
+    # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
+    set(CUTLASS_REVISION "v4.4.2")
+
+    # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
+    if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
+      set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
+    endif()
+
+    if(VLLM_CUTLASS_SRC_DIR)
+      if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
+        get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
+      endif()
+      message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
+      FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
+    else()
+      FetchContent_Declare(
+          cutlass
+          GIT_REPOSITORY https://github.com/nvidia/cutlass.git
+          # Please keep this in sync with CUTLASS_REVISION line above.
+          GIT_TAG ${CUTLASS_REVISION}
+          GIT_PROGRESS TRUE
+
+          # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
+          # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
+          # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
+          GIT_SHALLOW TRUE
+      )
+    endif()
+    FetchContent_MakeAvailable(cutlass)
+
    list(APPEND VLLM_STABLE_EXT_SRC
      "csrc/libtorch_stable/cuda_view.cu"
      "csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -512,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.
@@ -857,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")
@@ -929,7 +941,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${FP4_SM120_SRCS}"
      CUDA_ARCHS "${FP4_SM120_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
    message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +973,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${FP4_SM100_SRCS}"
      CUDA_ARCHS "${FP4_SM100_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
    message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
@@ -1065,6 +1075,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    target_compile_definitions(_C_stable_libtorch PRIVATE
      TORCH_TARGET_VERSION=0x020B000000000000ULL)
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
+    if(COOPERATIVE_TOPK_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_COOPERATIVE_TOPK=1)
+    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -8,43 +8,73 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR})
  set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
 endif()

+# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download).
+# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW
+# disallows one-argument Populate(dep) after Declare; MakeAvailable would run
+# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load).
 if(DEEPGEMM_SRC_DIR)
-  FetchContent_Declare(
-    deepgemm
-    SOURCE_DIR ${DEEPGEMM_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  # cmake_path(ABSOLUTE_PATH <var> ...) reads the path from <var>; NORMALIZE is a
+  # flag (no trailing path argument). Resolve relative paths against vLLM root.
+  set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _deepgemm_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}")
+  if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}")
+    message(FATAL_ERROR
+      "DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'")
+  endif()
+  set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}")
+  message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
 else()
-  # This ref should be kept in sync with tools/install_deepgemm.sh
-  FetchContent_Declare(
-    deepgemm
-    GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
-    GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
-    GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  # Keep in sync with tools/install_deepgemm.sh
+  set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
+  set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
+
+  set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _deepgemm_fc_root)
+    set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src")
+  set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build")
+  set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild")
+
+  if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp")
+    set(deepgemm_SOURCE_DIR "${_deepgemm_src}")
+    set(deepgemm_BINARY_DIR "${_deepgemm_bin}")
+  else()
+    FetchContent_Populate(
+      deepgemm
+      SUBBUILD_DIR "${_deepgemm_sub}"
+      SOURCE_DIR "${_deepgemm_src}"
+      BINARY_DIR "${_deepgemm_bin}"
+      GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
+      GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
+      GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+  message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
 endif()

-# Use FetchContent_Populate (not MakeAvailable) to avoid processing
-# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls.
-FetchContent_GetProperties(deepgemm)
-if(NOT deepgemm_POPULATED)
-  FetchContent_Populate(deepgemm)
-endif()
-message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
-
-# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100
+# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream),
+# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x
+# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants.
 set(DEEPGEMM_SUPPORT_ARCHS)
 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
  list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
 endif()
-if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
-  list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
-elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
-  list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
+if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
+  else()
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f")
+  else()
+    list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a")
+  endif()
 endif()

 cuda_archs_loose_intersection(DEEPGEMM_ARCHS
@@ -17,7 +17,7 @@ else()
  FetchContent_Declare(
    fmha_sm100
    GIT_REPOSITORY https://github.com/vllm-project/MSA.git
-    GIT_TAG 544eee5e09ae2dfa774d5b06739013f9b7402c57
+    GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5
    GIT_PROGRESS TRUE
    CONFIGURE_COMMAND ""
    BUILD_COMMAND ""
@@ -36,13 +36,38 @@ set(FMHA_SM100_PY_ROOT "${fmha_sm100_SOURCE_DIR}/python/fmha_sm100")

 install(FILES
  "${FMHA_SM100_PY_ROOT}/__init__.py"
+  "${FMHA_SM100_PY_ROOT}/api.py"
+  "${FMHA_SM100_PY_ROOT}/bench_utils.py"
+  "${FMHA_SM100_PY_ROOT}/jit.py"
  "${FMHA_SM100_PY_ROOT}/sparse.py"
+  "${FMHA_SM100_PY_ROOT}/sparse_fmha_adapter.py"
  DESTINATION vllm/third_party/fmha_sm100
  COMPONENT fmha_sm100)

+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/csrc/"
+  DESTINATION vllm/third_party/fmha_sm100/csrc
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
+
 install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cute/"
  DESTINATION vllm/third_party/fmha_sm100/cute
  COMPONENT fmha_sm100
  PATTERN "__pycache__" EXCLUDE
  PATTERN "*.pyc" EXCLUDE
  PATTERN ".git*" EXCLUDE)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/include/"
+  DESTINATION vllm/third_party/fmha_sm100/cutlass/include
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
+
+install(DIRECTORY "${FMHA_SM100_PY_ROOT}/cutlass/tools/util/include/"
+  DESTINATION vllm/third_party/fmha_sm100/cutlass/tools/util/include
+  COMPONENT fmha_sm100
+  PATTERN "__pycache__" EXCLUDE
+  PATTERN "*.pyc" EXCLUDE
+  PATTERN ".git*" EXCLUDE)
@@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR})
  set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
 endif()

+# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid.
+# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees.
 if(QUTLASS_SRC_DIR)
-  FetchContent_Declare(
-    qutlass
-    SOURCE_DIR ${QUTLASS_SRC_DIR}
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
+  set(_qutlass_user_src "${QUTLASS_SRC_DIR}")
+  cmake_path(ABSOLUTE_PATH _qutlass_user_src
+    BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
+    NORMALIZE)
+  set(QUTLASS_SRC_DIR "${_qutlass_user_src}")
+  if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}")
+    message(FATAL_ERROR
+      "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'")
+  endif()
+  set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}")
+  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
-  FetchContent_Declare(
-    qutlass
-    GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git
-    GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65
-    GIT_PROGRESS TRUE
-    CONFIGURE_COMMAND ""
-    BUILD_COMMAND ""
-  )
-endif()
+  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
+  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")

-FetchContent_Populate(qutlass)
+  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
+  if(NOT _qutlass_fc_root)
+    set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps")
+  endif()
+  set(_qutlass_src "${_qutlass_fc_root}/qutlass-src")
+  set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build")
+  set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild")
+
+  if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp")
+    set(qutlass_SOURCE_DIR "${_qutlass_src}")
+    set(qutlass_BINARY_DIR "${_qutlass_bin}")
+  else()
+    FetchContent_Populate(
+      qutlass
+      SUBBUILD_DIR "${_qutlass_sub}"
+      SOURCE_DIR "${_qutlass_src}"
+      BINARY_DIR "${_qutlass_bin}"
+      GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}"
+      GIT_TAG "${_QUTLASS_UPSTREAM_TAG}"
+      GIT_PROGRESS TRUE
+    )
+  endif()
+endif()

 if(NOT qutlass_SOURCE_DIR)
  message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")
@@ -60,6 +82,7 @@ endif()

 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
  set(QUTLASS_SOURCES
+    csrc/qutlass_registration.cpp
    ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
    ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
    ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +101,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)

  if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
    list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
+    if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
+       EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
+      list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
+    else()
+      get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
+      if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
+        list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
+      endif()
+    endif()
  elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
-    list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
+    list(APPEND QUTLASS_INCLUDES
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
    message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
  else()
    message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +125,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    CUDA_ARCHS "${QUTLASS_ARCHS}"
  )

-  target_sources(_C PRIVATE ${QUTLASS_SOURCES})
-  target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
-  target_compile_definitions(_C PRIVATE
+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
+  define_extension_target(
+    _qutlass_C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${QUTLASS_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_definitions(_qutlass_C PRIVATE
    QUTLASS_DISABLE_PYBIND=1
    TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-  )
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

  set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +156,5 @@ else()
      "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
      "CUDA_ARCHS='${CUDA_ARCHS}'.")
  endif()
+  add_custom_target(_qutlass_C)
 endif()
@@ -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;
+      }
    }
  });
 }
@@ -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
@@ -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);
  }
@@ -136,8 +136,12 @@ typename T::Fmha::Arguments args_from_options(
  StrideQ stride_Q_pe = cute::make_tuple(
      static_cast<int64_t>(q_pe.stride(1)), _1{}, static_cast<int64_t>(q_pe.stride(0)));

+  // Read the token and page strides from the cache tensor instead of assuming
+  // packed pages, so strided views (e.g. per-layer views into a cross-layer
+  // block-major cache) are addressed correctly.
  StrideK stride_C = cute::make_tuple(
-      static_cast<int64_t>(0 + D_latent + D_rope), _1{}, static_cast<int64_t>(page_size * (D_latent + D_rope)));
+      static_cast<int64_t>(kv_c_and_k_pe_cache.stride(1)), _1{},
+      static_cast<int64_t>(kv_c_and_k_pe_cache.stride(0)));
  StrideLSE stride_PT = cute::make_stride(_1{}, page_count_per_seq);
  StrideLSE stride_LSE = cute::make_tuple(_1{}, 0 + H);
  StrideO stride_O = cute::make_tuple(static_cast<int64_t>(0 + D_latent), _1{}, static_cast<int64_t>(0 + H * D_latent));
@@ -268,9 +272,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
  using TileShapeD = typename MlaSm100Type::TileShapeD;
  arguments.problem_shape =
      cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
-  // Assumes device 0 when getting sm_count.
-  arguments.hw_info.sm_count =
-      sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
+  if (sm_count <= 0) {
+    int current_device = 0;
+    cudaGetDevice(&current_device);
+    arguments.hw_info.sm_count =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
+  } else {
+    arguments.hw_info.sm_count = sm_count;
+  }
  arguments.split_kv = static_cast<int>(num_kv_splits);
  MlaSm100Type::Fmha::set_split_kv(arguments);

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

 template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>(
    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+// Template instantiations for GLM-5 (DEFAULT_NUM_EXPERTS, hidden_dim=6144)
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
+
+template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 6144>(
+    float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t);
@@ -9,7 +9,7 @@
 #include <torch/headeronly/core/ScalarType.h>

 #include "../../cuda_compat.h"
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/dispatch_utils.h"
 #include "libtorch_stable/torch_utils.h"

@@ -2,9 +2,25 @@

 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/util/Exception.h>

 #include <optional>
 #include <string>
+#include <vector>
+
+#include <torch/csrc/stable/ops.h>
+
+inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
+  // Ensure tensor is on CUDA
+  STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
+
+  // Get the raw data pointer
+  void* data_ptr = tensor.mutable_data_ptr();
+
+  /// Create a new tensor from the raw data pointer
+  return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
+                                  tensor.device(), tensor.scalar_type());
+}

 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                               torch::stable::Tensor& output_q,
@@ -327,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,
@@ -371,6 +395,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
                        torch::stable::Tensor& input, double limit,
                        double alpha = 1.0, double beta = 0.0);
+
+void silu_and_mul_quant(torch::stable::Tensor& out,
+                        torch::stable::Tensor& input,
+                        torch::stable::Tensor& scale);
+
+void persistent_masked_m_silu_mul_quant(
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
+    bool use_ue8m0);
+
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,
@@ -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_
@@ -1,16 +1,12 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "libtorch_stable/torch_utils.h"

 #include <cmath>
-#include "core/math.hpp"
-#include "../cuda_compat.h"
-#include "dispatch_utils.h"

+#include "libtorch_stable/core/math.hpp"
+#include "cuda_compat.h"
+#include "libtorch_stable/dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"

-#include <c10/util/Float8_e4m3fn.h>
-
 #ifndef USE_ROCM
  #include <cuda_bf16.h>
  #include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
  #endif
 #endif

-#include "core/registration.h"
 namespace vllm {

 template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
 }  // namespace vllm

 // Launch activation, gating, and quantize kernel.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                               \
-  int d = input.size(-1) / 2;                                               \
-  int64_t num_tokens = input.numel() / input.size(-1);                      \
-  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);     \
-  dim3 block(std::min(d, 512));                                             \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
-  VLLM_DISPATCH_FLOATING_TYPES(                                             \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                      \
-        VLLM_DISPATCH_FP8_TYPES(                                            \
-            out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] {  \
-              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,    \
-                                             fp8_t>                         \
-                  <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),       \
-                                               input.data_ptr<scalar_t>(),  \
-                                               scale.data_ptr<float>(), d); \
-            });                                                             \
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                             \
+  int d = input.size(-1) / 2;                                             \
+  int64_t num_tokens = input.numel() / input.size(-1);                    \
+  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);   \
+  dim3 block(std::min(d, 512));                                           \
+  const torch::stable::accelerator::DeviceGuard device_guard(             \
+      input.get_device_index());                                          \
+  const cudaStream_t stream =                                             \
+      get_current_cuda_stream(input.get_device_index());                  \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                    \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                    \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                   \
+            out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
+              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,  \
+                                             fp8_t>                       \
+                  <<<grid, block, 0, stream>>>(                           \
+                      out.mutable_data_ptr<fp8_t>(),                      \
+                      input.const_data_ptr<scalar_t>(),                   \
+                      scale.const_data_ptr<float>(), d);                  \
+            });                                                           \
      });

-void silu_and_mul_quant(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        torch::Tensor& scale) {
-  TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
-              out.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.dtype() == torch::kFloat16 ||
-              input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(input.size(-1) % 2 == 0);
+void silu_and_mul_quant(torch::stable::Tensor& out,    // [..., d]
+                        torch::stable::Tensor& input,  // [..., 2 * d]
+                        torch::stable::Tensor& scale) {
+  STD_TORCH_CHECK(
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::Half ||
+          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Input must be FP16 or BF16");
+  STD_TORCH_CHECK(input.size(-1) % 2 == 0);
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
 }

 void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,              // (E, T, 2*H)
-    const at::Tensor& tokens_per_expert,  // (E)
-    at::Tensor& y_q,                      // (E, T, H) [OUT]
-    at::Tensor& y_s,                      // (E, T, H//group_size) [OUT]
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
    bool cast_scale_ue8m0) {
 #ifndef USE_ROCM

@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
  // fixed GROUP_SIZE of 128.
  static constexpr int GROUP_SIZE = 128;

-  TORCH_CHECK(input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
-              y_q.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
+  STD_TORCH_CHECK(input.scalar_type() ==
+                  torch::headeronly::ScalarType::BFloat16);
+  STD_TORCH_CHECK(
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);

  bool const is_packed_ue8m0 =
-      (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
+      (y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
+       cast_scale_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
+                  is_packed_ue8m0);

  using Idx_t = int64_t;

@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(

  int const NUM_GROUPS = H / GROUP_SIZE;

-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());

  // TODO: Get this from cuda_arch ?
  static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
    static constexpr int max_shared_mem_bytes =                                \
        GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2;                               \
    dim3 grid(sms), block(THREAD_COUNT);                                       \
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));          \
-    VLLM_DISPATCH_FP8_TYPES(                                                   \
+    const torch::stable::accelerator::DeviceGuard device_guard(                \
+        input.get_device_index());                                             \
+    VLLM_STABLE_DISPATCH_FP8_TYPES(                                            \
        y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] {        \
          vllm::silu_mul_fp8_quant_deep_gemm_kernel<                           \
              BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
              Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES>                           \
              <<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>(  \
-                  reinterpret_cast<__nv_bfloat16*>(input.data_ptr()),          \
-                  (fp8_t*)y_q.data_ptr(),                                      \
-                  reinterpret_cast<scale_t*>(y_s.data_ptr()),                  \
-                  reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
-                  T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,       \
+                  reinterpret_cast<const __nv_bfloat16*>(                      \
+                      input.const_data_ptr()),                                 \
+                  y_q.mutable_data_ptr<fp8_t>(),                               \
+                  reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()),          \
+                  reinterpret_cast<const int32_t*>(                            \
+                      tokens_per_expert.const_data_ptr()),                     \
+                  E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,    \
                  stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T,          \
                  STRIDE_YS_G, STRIDE_YS_P, stride_counts_e);                  \
        });
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
  Idx_t stride_ys_g = y_s.stride(2);
  Idx_t stride_ys_p = 0;
  if (!cast_scale_ue8m0) {
-    TORCH_CHECK(!is_packed_ue8m0);
+    STD_TORCH_CHECK(!is_packed_ue8m0);
    LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
                false);
    return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
    return;
  }

-  TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kInt32);
+  STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);

  // Int32 packed ue8m0 scales tensor.
  // Let E, T, G be the number to experts, number of tokens and number of groups
@@ -31,7 +31,7 @@

 #include "cutlass/util/packed_stride.hpp"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"

 using namespace cute;
@@ -31,7 +31,7 @@

 #include "cutlass/util/packed_stride.hpp"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"

 using namespace cute;
@@ -19,7 +19,7 @@
 #include "cutlass/gemm/collective/collective_builder.hpp"
 #include "cutlass/util/packed_stride.hpp"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on

@@ -14,7 +14,7 @@
 #include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/collective/collective_builder.hpp"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on

@@ -22,7 +22,7 @@
 #include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/cutlass_extensions/common.hpp"
 // clang-format on

@@ -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
@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(

  const int sf_k_local = local_group_id % kGroupsPerBlockX;
  const int row_local = local_group_id / kGroupsPerBlockX;
-  const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
-  const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
+  // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
+  const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
                  " is not a multiple of 4.");
  const int kx = GetGroupsPerBlockX(padded_groups_per_row);
  const int ry = 16 / kx;
-  const int64_t blocks_x = padded_groups_per_row / kx;
-  const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t sf_k_blocks = padded_groups_per_row / kx;
  const int num_threads = (kx * ry) * THREADS_PER_GROUP;
-  // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
-  STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
-                      blocks_y <= static_cast<int64_t>(INT32_MAX),
+  // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
+  constexpr int64_t kMaxGridDimYZ = 65535;
+  STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
+                      sf_k_blocks <= kMaxGridDimYZ,
                  "per_token_group_quant_8bit_packed grid too large: (",
-                  blocks_x, ", ", blocks_y, ").");
+                  row_blocks, ", ", sf_k_blocks, ").");

  auto dst_type = output_q.scalar_type();

@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
    do {                                                                       \
      cudaLaunchConfig_t config = {};                                          \
-      config.gridDim = dim3(static_cast<unsigned int>(blocks_x),               \
-                            static_cast<unsigned int>(blocks_y));              \
+      config.gridDim = dim3(static_cast<unsigned int>(row_blocks),             \
+                            static_cast<unsigned int>(sf_k_blocks));           \
      config.blockDim = dim3(num_threads);                                     \
      config.dynamicSmemBytes = 0;                                             \
      config.stream = stream;                                                  \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 #else
  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
    do {                                                                       \
-      dim3 grid(static_cast<unsigned int>(blocks_x),                           \
-                static_cast<unsigned int>(blocks_y));                          \
+      dim3 grid(static_cast<unsigned int>(row_blocks),                         \
+                static_cast<unsigned int>(sf_k_blocks));                       \
      dim3 block(num_threads);                                                 \
      per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
                                                        RY>                    \
@@ -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_
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
  // TODO: Remove this once ROCm upgrade to torch 2.11.
  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");

+  // Note about marlin kernel 'workspace' arguments:
+  // Technically these should be mutable since they are modified by the kernel.
+  // But since they are set back to zero once the kernel is finished we can
+  // hand wave and say that they have no net effect.
+  //
+  // The reason to mark 'workspace' as immutable is so that they don't interfere
+  // with using ScalarType arguments in the ops. If they are marked as mutable,
+  // pytorch throws an assert in
+  // 'torch._higher_order_ops._register_effectful_op' that prevents these
+  // kernels from being torch.compile'd.
+  // See the following document for more info on custom types and ops that use
+  // custom types:
+  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
+
  // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
  ops.def(
      "machete_supported_schedules("
@@ -479,7 +493,18 @@ 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! "
+      "y_q, Tensor! y_s, bool use_ue8m0) -> ()");
+  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
+
  // Activation function used in SwiGLU.
  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");

@@ -492,6 +517,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
      "float alpha=1.0, float beta=0.0) -> ()");

+  // SwiGLU activation with FP8 quantization.
+  ops.def(
+      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
+
  // Activation function used in GeGLU with `none` approximation.
  ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");

@@ -688,8 +717,15 @@ 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",
+           TORCH_BOX(&persistent_masked_m_silu_mul_quant));
+  ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
+  ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
  ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
  ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
  ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
@@ -9,28 +9,6 @@

 #include <vector>

-torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
-  // Ensure tensor is on CUDA
-  if (!tensor.is_cuda()) {
-    throw std::runtime_error("Tensor must be on CUDA device");
-  }
-
-  // Get the raw data pointer
-  void* data_ptr = tensor.data_ptr();
-
-  // Get tensor sizes and strides
-  std::vector<int64_t> sizes = tensor.sizes().vec();
-  std::vector<int64_t> strides = tensor.strides().vec();
-
-  // Get tensor options (dtype, device)
-  auto options = tensor.options();
-
-  // Create a new tensor from the raw data pointer
-  auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
-
-  return new_tensor;
-}
-
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
                        double alpha = 1.0, double beta = 0.0);

-void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
-                        torch::Tensor& scale);
-
-void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,   // (E, T, 2*H)
-    const at::Tensor& counts,  // (E)
-    at::Tensor& y_q,           // (E, T, H) [OUT]
-    at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
-    bool use_ue8m0);
-
 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);

 void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
@@ -0,0 +1,5 @@
+#include "core/registration.h"
+
+// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
+// This stub lets Python import vllm._qutlass_C to trigger op registration.
+REGISTER_EXTENSION(_qutlass_C)
@@ -20,17 +20,6 @@

 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // vLLM custom ops
-  //
-
-  ops.def(
-      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
-      "y_q, Tensor! y_s,"
-      "bool use_ue8m0) -> ()");
-  ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
-           &persistent_masked_m_silu_mul_quant);
-
-  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
-  ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);

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

@@ -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 \
@@ -338,7 +338,7 @@ vllm bench serve \
    --model meta-llama/Meta-Llama-3-8B-Instruct \
    --dataset-name spec_bench \
    --dataset-path "<YOUR_DOWNLOADED_PATH>/data/spec_bench/question.jsonl" \
-    --num-prompts -1
+    --num-prompts -1 \
    --spec-bench-category "summarization"
 ```

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

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

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

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

 # Combine pipeline and tensor parallelism
 llm = LLM(
-    model="meta-llama/Llama-3.3-70B-Instruct,
+    model="meta-llama/Llama-3.3-70B-Instruct",
    tensor_parallel_size=4,
    pipeline_parallel_size=2,
 )
@@ -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>

@@ -133,10 +133,10 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
 For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
 The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.

-For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
+For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
 These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).

-For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
+For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`.
 Please follow the same guidelines as case (2) for implementing these models.
 We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention).
 For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`.
@@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles
 Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this.
 It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend.
 Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it.
-Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
+Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this.
 The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended.
@@ -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,8 +160,8 @@ Priority is **1 = highest** (tried first).
 | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
 | `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
-| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
-| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
+| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ✅ | ❌ | ✅ | Decoder | 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 |
 | `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
@@ -231,7 +221,8 @@ MLA decode backends are selected using the standard
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
@@ -248,10 +239,11 @@ DeepSeek V4 sparse MLA uses its own decode backends, selected via
 `--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
 `FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
 pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
-default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
+default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and
+`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.

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

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

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

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

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

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

 ### Eviction (LRU)

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

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

 ## Running a quantized model with vLLM

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

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

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

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

 ## Known Feature Incompatibility

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

 ## Resources for vLLM contributors
@@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup

 Flags: `--tool-call-parser hermes`

-### MiniMax Models (`minimax_m1`)
-
-Supported models:
-
-* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
-* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja))
-
-Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja`
-
 ### DeepSeek-V3 Models (`deepseek_v3`)

 Supported models:
@@ -347,7 +338,7 @@ Supported models:

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

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

 Supported models:

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

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

 ### `LLM.embed`

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

 ### Score task

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

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

 ## Supported Features

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

 ### Score Template

@@ -61,7 +61,7 @@ Models of any architecture can be converted into embedding models using `--conve
 | `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
 | `ColPaliForRetrieval` | ColPali | T / I | `vidore/colpali-v1.3-hf` | | |
 | `ColQwen3` | Qwen3-VL | T / I | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` | | |
-| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3` | | |
+| `ColQwen3_5` | ColQwen3.5 | T + I + V | `athrael-soju/colqwen3.5-4.5B-v3`, `vultr/VultronRetrieverPrime-Qwen3.5-8B` | | |
 | `OpsColQwen3Model` | Qwen3-VL | T / I | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` | | |
 | `Qwen3VLNemotronEmbedModel` | Qwen3-VL | T / I | `nvidia/nemotron-colembed-vl-4b-v2`, `nvidia/nemotron-colembed-vl-8b-v2` | ✅︎ | ✅︎ |
 | `*ForConditionalGeneration`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | \* | N/A | \* | \* |
@@ -441,7 +441,6 @@ th {
 | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
 | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
 | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
-| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
 | `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ |
 | `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ |
 | `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ |
@@ -487,8 +486,6 @@ th {
 | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
 | `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
-| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | |
-| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | |
 | `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | |

 !!! note
@@ -595,10 +592,12 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
 | `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
 | `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
+| `MiniMaxM3SparseForConditionalGeneration` | MiniMax-M3 | T + I<sup>+</sup> + V<sup>+</sup> | `MiniMaxAI/MiniMax-M3`, `MiniMaxAI/MiniMax-M3-MXFP8`, etc. | | |
 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
 | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
 | `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
 | `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. | | ✅︎ |
@@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor
 Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`,
 `Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`).

-Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`).
+Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`).

 Please note that prefix caching is not yet supported for any of the above models.

@@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData:
    return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6")


-def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData:
-    assert modality == "image"
-
-    model_name = "MiniMaxAI/MiniMax-VL-01"
-
-    engine_args = EngineArgs(
-        model=model_name,
-        max_num_seqs=2,
-        limit_mm_per_prompt={modality: 1},
-        trust_remote_code=True,
-        tensor_parallel_size=8,
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    messages = [
-        [
-            {
-                "role": "user",
-                "content": [{"type": "image"}, {"type": "text", "text": question}],
-            }
-        ]
-        for question in questions
-    ]
-    prompts = tokenizer.apply_chat_template(
-        messages, add_generation_prompt=True, tokenize=False
-    )
-
-    return ModelRequestData(
-        engine_args=engine_args,
-        prompts=prompts,
-    )
-
-
 # Mistral-3 HF-format
 def run_mistral3(questions: list[str], modality: str) -> ModelRequestData:
    assert modality == "image"
@@ -2485,7 +2452,6 @@ model_example_map = {
    "mantis": run_mantis,
    "minicpmo": run_minicpmo,
    "minicpmv": run_minicpmv,
-    "minimax_vl_01": run_minimax_vl_01,
    "mistral3": run_mistral3,
    "molmo": run_molmo,
    "molmo2": run_molmo2,
@@ -7,11 +7,27 @@ ColQwen3.5 is a multi-modal ColBERT-style model based on Qwen3.5.
 It produces per-token embeddings and uses MaxSim scoring for retrieval
 and reranking. Supports both text and image inputs.

+Works for any ColQwen3.5 checkpoint, e.g. `athrael-soju/colqwen3.5-4.5B-v3`
+or `vultr/VultronRetrieverPrime-Qwen3.5-8B`.
+
 Start the server with:
-    vllm serve athrael-soju/colqwen3.5-4.5B --max-model-len 4096
+    vllm serve athrael-soju/colqwen3.5-4.5B-v3 --max-model-len 4096 \
+        --mm-processor-kwargs '{"min_pixels": 65536, "max_pixels": 1835008}'

 Then run this script:
    python colqwen3_5_rerank_online.py
+
+Parity note (matching the native colpali ColQwen3_5Processor pipeline):
+  - Visual-token budget: ColQwen3_5Processor uses max_num_visual_tokens=1792,
+    i.e. max_pixels = 1792 * (patch_size*merge_size)^2 = 1792 * 32^2 = 1835008
+    (with min_pixels = shortest_edge = 65536). Pass these via --mm-processor-kwargs
+    as above; the default budget gives fewer visual tokens and lower retrieval ndcg.
+  - When you build prompts yourself (token_embed), reproduce the processor exactly:
+      image (document): wrap in the instruction template
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        "Describe the image.<|im_end|><|endoftext|>"
+      query: append the augmentation suffix  <text> + "<|endoftext|>" * 10
+    Omitting these reproduces a silent ~2.5 ndcg@10 drop vs the native pipeline.
 """

 import requests
@@ -1,91 +0,0 @@
-{{ '<begin_of_document>' -}}
-{%- if custom_tools is defined %}
-    {%- set tools = custom_tools %}
-{%- endif %}
-{%- if not tools is defined %}
-    {%- set tools = none %}
-{%- endif %}
-
-{#- Extract system message #}
-{% set ns = namespace(system_prompt='') -%}
-{%- if messages[0]['role'] == 'system' %}
-    {%- if messages[0]['content'] is string %}
-        {%- set ns.system_prompt = messages[0]['content']|trim %}
-    {%- else %}
-        {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
-    {%- endif %}
-    {%- set messages = messages[1:] %}
-{%- else %}
-    {%- if tools is not none %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- else %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- endif %}
-{%- endif %}
-
-{#- System message #}
-{%- if ns.system_prompt != '' %}
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Tools configuration #}
-{%- if tools is not none %}
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
-{%- for tool in tools %}
-{{ tool | tojson ~ '\n' -}}
-{%- endfor %}
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Process messages #}
-{%- for message in messages %}
-    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
-        {%- if message['role'] == 'user' %}
-{{ '<beginning_of_sentence>user name=user\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ content['text']|trim -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- elif message['role'] == 'assistant' %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
-{{ content['text']|trim -}}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- endif %}
-    {%- elif 'tool_calls' in message %}
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
-{%- for tool_call in message.tool_calls %}
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
-{%- endfor %}
-{{ '</tool_calls><end_of_sentence>\n' -}}
-    {%- elif message.role == "tool" or message.role == "ipython" %}
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
-{%- if message.content is string %}
-{{ 'tool result: ' + message.content + '\n\n' -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
-{%- elif content.get('name') %}
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-    {%- endif %}
-{%- endfor %}
-
-{%- if add_generation_prompt %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- endif %}
@@ -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
@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test.
+schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test.
+schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 bitsandbytes>=0.49.2
 buildkite-test-collector==0.1.9
@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test
+schemathesis>=4.0.0 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
@@ -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}"))
        })
@@ -386,7 +386,6 @@ mod tests {
            tool_chat_template_llama3.2_pythonic.jinja         => String
            tool_chat_template_llama4_json.jinja               => OpenAi
            tool_chat_template_llama4_pythonic.jinja           => OpenAi
-            tool_chat_template_minimax_m1.jinja                => OpenAi
            tool_chat_template_mistral.jinja                   => String
            tool_chat_template_mistral3.jinja                  => OpenAi
            tool_chat_template_mistral_parallel.jinja          => String
@@ -1,91 +0,0 @@
-{{ '<begin_of_document>' -}}
-{%- if custom_tools is defined %}
-    {%- set tools = custom_tools %}
-{%- endif %}
-{%- if not tools is defined %}
-    {%- set tools = none %}
-{%- endif %}
-
-{#- Extract system message #}
-{% set ns = namespace(system_prompt='') -%}
-{%- if messages[0]['role'] == 'system' %}
-    {%- if messages[0]['content'] is string %}
-        {%- set ns.system_prompt = messages[0]['content']|trim %}
-    {%- else %}
-        {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %}
-    {%- endif %}
-    {%- set messages = messages[1:] %}
-{%- else %}
-    {%- if tools is not none %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- else %}
-        {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %}
-    {%- endif %}
-{%- endif %}
-
-{#- System message #}
-{%- if ns.system_prompt != '' %}
-{{ '<beginning_of_sentence>system ai_setting=assistant\n' + ns.system_prompt + '<end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Tools configuration #}
-{%- if tools is not none %}
-{{ '<beginning_of_sentence>system tool_setting=tools\nYou are provided with these tools:\n<tools>\n' -}}
-{%- for tool in tools %}
-{{ tool | tojson ~ '\n' -}}
-{%- endfor %}
-{{ '</tools>\n\nIf you need to call tools, please respond with <tool_calls></tool_calls> XML tags, and provide tool-name and json-object of arguments, following the format below:\n<tool_calls>\n{"name": <tool-name>, "arguments": <args-json-object>}\n...\n</tool_calls><end_of_sentence>\n' -}}
-{%- endif %}
-
-{#- Process messages #}
-{%- for message in messages %}
-    {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %}
-        {%- if message['role'] == 'user' %}
-{{ '<beginning_of_sentence>user name=user\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ content['text']|trim -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- elif message['role'] == 'assistant' %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- if message['content'] is string %}
-{{ message['content']|trim -}}
-{%- else %}
-{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %}
-{{ content['text']|trim -}}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-        {%- endif %}
-    {%- elif 'tool_calls' in message %}
-{{ '<beginning_of_sentence>ai name=assistant\n<tool_calls>\n' -}}
-{%- for tool_call in message.tool_calls %}
-{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}}
-{%- endfor %}
-{{ '</tool_calls><end_of_sentence>\n' -}}
-    {%- elif message.role == "tool" or message.role == "ipython" %}
-{{ '<beginning_of_sentence>tool name=tools\n' -}}
-{%- if message.content is string %}
-{{ 'tool result: ' + message.content + '\n\n' -}}
-{%- else %}
-{%- for content in message['content'] %}
-{%- if content['type'] == 'text' %}
-{{ 'tool result: ' + content['text'] + '\n\n' -}}
-{%- elif content.get('name') %}
-{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}}
-{%- endif %}
-{%- endfor %}
-{%- endif %}
-{{ '<end_of_sentence>\n' -}}
-    {%- endif %}
-{%- endfor %}
-
-{%- if add_generation_prompt %}
-{{ '<beginning_of_sentence>ai name=assistant\n' -}}
-{%- endif %}
@@ -15,6 +15,7 @@ use crate::protocol::handshake::EngineCoreReadyResponse;
 use crate::protocol::lora::LoraRequest;
 use crate::protocol::utility::{EngineCoreUtilityRequest, PauseMode};
 use crate::protocol::{EngineCoreRequest, EngineCoreRequestType, ModelDtype};
+use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime};
 use crate::transport::{self, ConnectedEngine};

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

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

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

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

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

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

    /// Handle an abort request by sending the abort message to the engine.
@@ -434,6 +449,7 @@ mod tests {
        let (send, _) = socket.split();
        ClientInner::new(
            send,
+            Handle::current(),
            "test-model".to_string(),
            &[ConnectedEngine {
                engine_id: EngineId::from(b"engine-0"),
@@ -29,6 +29,8 @@ pub enum Error {
    Io(#[from] std::io::Error),
    #[error("transport error")]
    Transport(#[from] zeromq::ZmqError),
+    #[error("ZMQ runtime task failed")]
+    ZmqRuntimeTask(#[from] tokio::task::JoinError),
    #[error("engine core reported fatal failure")]
    EngineCoreDead,
    #[error("startup handshake timed out while waiting for {stage} after {timeout:?}")]
@@ -4,6 +4,7 @@ mod error;
 mod metrics;
 pub mod mock_engine;
 pub mod protocol;
+pub mod runtime;
 #[cfg(any(test, feature = "test-util"))]
 pub mod test_utils;
 mod transport;
@@ -277,6 +277,12 @@ pub struct EngineCoreSamplingParams {
    pub max_tokens: u32,
    /// Minimum number of tokens to generate before EOS or stop-token handling.
    pub min_tokens: u32,
+    /// Maximum number of reasoning ("thinking") tokens to emit before the
+    /// reasoning section is force-closed. `None` means unlimited; the
+    /// user-facing `-1` sentinel is normalized to `None` by the frontend before
+    /// reaching this DTO, so only non-negative values are sent. Enforced
+    /// engine-side (and only when a reasoning parser is configured).
+    pub thinking_token_budget: Option<u64>,
    /// Number of log probabilities to return per generated token.
    ///
    /// `None` disables sample logprobs. `-1` requests the full vocabulary.
@@ -345,6 +351,7 @@ impl EngineCoreSamplingParams {
            seed: None,
            max_tokens: 65536,
            min_tokens: 0,
+            thinking_token_budget: None,
            logprobs: None,
            prompt_logprobs: None,
            min_p: 0.0,
@@ -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)
+}
@@ -150,6 +150,7 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest {
            top_k: 8,
            max_tokens: 32,
            min_tokens: 1,
+            thinking_token_budget: Some(256),
            stop_token_ids: vec![151643],
            eos_token_id: Some(151645),
            all_stop_token_ids: BTreeSet::from([151643, 151645]),
@@ -2502,6 +2503,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
            seed: None,
            max_tokens: 16,
            min_tokens: 0,
+            thinking_token_budget: None,
            logprobs: None,
            prompt_logprobs: None,
            min_p: 0.0,
@@ -39,6 +39,7 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True):
    seed: int | None = None
    max_tokens: int = 16
    min_tokens: int = 0
+    thinking_token_budget: int | None = None
    min_p: float = 0.0
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0
@@ -122,6 +123,7 @@ request = EngineCoreRequest(
        seed=None,
        max_tokens=32,
        min_tokens=1,
+        thinking_token_budget=256,
        min_p=0.0,
        frequency_penalty=0.0,
        presence_penalty=0.0,
@@ -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
@@ -103,6 +103,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
            | vllm_text::Error::EmptyPromptTokenIds { .. }
            | vllm_text::Error::Logprobs(_)
            | vllm_text::Error::OutOfVocab(_)
+            | vllm_text::Error::InvalidThinkingTokenBudget
            // An empty tokenized prompt detected later, at request prepare
            // time, surfaces through the transparent Llm wrapper.
            | vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
@@ -127,6 +128,18 @@ mod tests {
        assert!(response.error.message.contains("9000"));
    }

+    #[test]
+    fn invalid_thinking_token_budget_maps_to_invalid_request() {
+        let api_error = text_submit_error(
+            "failed to submit completion request",
+            vllm_text::Error::InvalidThinkingTokenBudget,
+        );
+        assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
+        let response = api_error.to_error_response();
+        assert_eq!(response.error.error_type, "invalid_request_error");
+        assert!(response.error.message.contains("thinking_token_budget"));
+    }
+
    #[test]
    fn chat_wrapped_prompt_too_long_maps_to_invalid_request() {
        let error = vllm_chat::Error::Text(vllm_text::Error::PromptTooLong {
@@ -7,6 +7,7 @@ mod listener;
 mod lora;
 mod middleware;
 mod routes;
+mod runtime;
 mod server_info;
 mod state;
 mod utils;
@@ -157,6 +158,9 @@ where
            .with_context(|| format!("failed to bind gRPC listener on {grpc_host}:{grpc_port}"))?;
        let addr = grpc_listener.local_addr()?;
        let svc = grpc::GenerateServer::new(grpc::GenerateServiceImpl::new(state.clone()));
+        let svc = TonicServer::builder()
+            .layer(middleware::request_runtime_layer(state.clone()))
+            .add_service(svc);
        info!(%addr, "starting gRPC server");
        Some((grpc_listener, svc))
    } else {
@@ -238,7 +242,7 @@ where
                shutdown.cancelled().await;
                return Ok(());
            };
-            let server = TonicServer::builder().add_service(svc).serve_with_incoming_shutdown(
+            let server = svc.serve_with_incoming_shutdown(
                TcpListenerStream::new(grpc_listener),
                shutdown.cancelled_owned(),
            );
@@ -2,10 +2,12 @@ mod auth;
 mod cors;
 mod load;
 mod metrics;
+mod offload;
 mod request_id;

 pub use auth::authenticate_api_key;
 pub use cors::{cors_layer, strip_cors_on_no_origin};
 pub use load::track_server_load;
 pub use metrics::track_http_metrics;
+pub(crate) use offload::request_runtime_layer;
 pub use request_id::set_request_id_header;
@@ -0,0 +1,134 @@
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use axum::http::Request;
+use axum::response::{IntoResponse, Response};
+use futures::future::BoxFuture;
+use tokio_util::task::AbortOnDropHandle;
+use tonic::Status;
+use tower::Service;
+use tower::layer::layer_fn;
+use tracing::error;
+
+use crate::error::{ApiError, server_error};
+use crate::state::AppState;
+
+/// Request paths that are run on the request runtime.
+///
+/// These routes can perform CPU-heavy request preparation, including JSON
+/// extraction, validation, chat-template rendering, tokenization, request
+/// lowering, and engine submission. Lightweight operational routes stay on the
+/// HTTP runtime.
+const OFFLOADED_PATHS: &[&str] = &[
+    // HTTP routes:
+    "/v1/chat/completions",
+    "/v1/completions",
+    "/tokenize",
+    "/detokenize",
+    "/inference/v1/generate",
+    // gRPC routes:
+    "/vllm.Generate/Generate",
+    "/vllm.Generate/GenerateStream",
+];
+
+/// Return a Tower layer that runs selected data-plane requests on the request runtime,
+/// so that we can offset heavy request parsing and preprocessing from the HTTP runtime.
+pub(crate) fn request_runtime_layer<S>(
+    state: Arc<AppState>,
+) -> impl tower::Layer<S, Service = RequestRuntimeService<S>> + Clone {
+    layer_fn(move |inner| RequestRuntimeService {
+        inner,
+        state: state.clone(),
+    })
+}
+
+/// Service produced by [`request_runtime_layer`].
+#[derive(Clone)]
+pub(crate) struct RequestRuntimeService<S> {
+    inner: S,
+    state: Arc<AppState>,
+}
+
+impl<S, B> Service<Request<B>> for RequestRuntimeService<S>
+where
+    S: Service<Request<B>> + Clone + Send + 'static,
+    S::Future: Send + 'static,
+    S::Response: RequestRuntimeErrorResponse + Send + 'static,
+    S::Error: Send + 'static,
+    B: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Request<B>) -> Self::Future {
+        if !should_offload(req.uri().path()) {
+            return Box::pin(self.inner.call(req));
+        }
+
+        // Axum extractors and route handlers execute inside the inner service,
+        // so offloading here moves request parsing and preprocessing off the
+        // HTTP runtime without wrapping each handler manually. For streaming
+        // HTTP responses, the response body is still polled on the HTTP runtime.
+        let clone = self.inner.clone();
+        let mut inner = std::mem::replace(&mut self.inner, clone);
+        let task = AbortOnDropHandle::new(self.state.request_runtime().spawn(inner.call(req)));
+
+        Box::pin(async move {
+            match task.await {
+                Ok(result) => result,
+                Err(error) => {
+                    error!(%error, "request runtime task failed");
+                    Ok(S::Response::request_runtime_error_response())
+                }
+            }
+        })
+    }
+}
+
+trait RequestRuntimeErrorResponse {
+    fn request_runtime_error_response() -> Self;
+}
+
+impl RequestRuntimeErrorResponse for Response {
+    fn request_runtime_error_response() -> Self {
+        server_error!("request runtime task failed").into_response()
+    }
+}
+
+impl RequestRuntimeErrorResponse for axum::http::Response<tonic::body::Body> {
+    fn request_runtime_error_response() -> Self {
+        Status::internal("request runtime task failed").into_http()
+    }
+}
+
+fn should_offload(path: &str) -> bool {
+    OFFLOADED_PATHS.contains(&path)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::should_offload;
+
+    #[test]
+    fn offloads_generation_and_tokenization_paths() {
+        assert!(should_offload("/v1/chat/completions"));
+        assert!(should_offload("/v1/completions"));
+        assert!(should_offload("/tokenize"));
+        assert!(should_offload("/detokenize"));
+        assert!(should_offload("/inference/v1/generate"));
+        assert!(should_offload("/vllm.Generate/Generate"));
+        assert!(should_offload("/vllm.Generate/GenerateStream"));
+    }
+
+    #[test]
+    fn passes_through_lightweight_paths() {
+        assert!(!should_offload("/health"));
+        assert!(!should_offload("/metrics"));
+        assert!(!should_offload("/v1/models"));
+    }
+}
--- a/Show More
+++ b/Show More