test: touch tilelang.py for coverage comparison

.buildkite/hardware_tests/cpu.yaml

@@ -62,16 +62,11 @@ steps:
   source_file_dependencies:
   - vllm/v1/worker/cpu/
   - vllm/v1/worker/gpu/
-  - vllm/v1/sample/ops/topk_topp_triton.py
-  - vllm/v1/sample/ops/topk_topp_sampler.py
-  - tests/v1/sample/test_topk_topp_sampler.py
   commands:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
       uv pip install git+https://github.com/triton-lang/triton-cpu.git@270e696d
-      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model
-      # TODO: move to CPU-Kernel Tests once triton-cpu has a pre-built wheel
-      pytest -x -v -s tests/v1/sample/test_topk_topp_sampler.py::TestTritonTopkTopp"
+      VLLM_USE_V2_MODEL_RUNNER=1 pytest -x -v -s tests/models/language/generation/test_granite.py -m cpu_model"
 
 - label: CPU-Quantization Model Tests
   depends_on: []

.buildkite/image_build/image_build.yaml

@@ -6,26 +6,14 @@ steps:
     timeout_in_minutes: 600
     commands:
     - if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
-    retry:
-      automatic:
-        - exit_status: -1  # Agent was lost
-          limit: 2
-        - exit_status: -10  # Agent was lost
-          limit: 2
-
-  - label: ":docker: :smoking: Non-root smoke tests"
-    key: image-build-smoke-test
-    depends_on:
-      - image-build
-    commands:
-    # Smoke 1: the default (root) image must still be importable
+    # Non-root smoke 1: the default (root) image must still be importable
     # under a non-root UID via `--user 2000:0`. Validates the `vllm` passwd
     # entry + group-0-writable /home/vllm + uv path cleanup from #31959.
     # Uses `import vllm` rather than `vllm serve --help` because the latter
     # instantiates `VllmConfig` which requires a GPU attached to the
     # container.
     - docker run --rm --user 2000:0 --entrypoint python3 "$IMAGE_TAG" -c "import vllm; print(vllm.__version__)"
-    # Smoke 2: assert the non-root enabling invariants are baked
+    # Non-root smoke 2: assert the non-root enabling invariants are baked
     # into the image. Runs as UID 2000:0 via a shell so we can verify
     # filesystem perms + passwd/group file state + wrapper presence without
     # triggering vLLM's GPU-requiring config-init path. The opt-in

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

@@ -9,13 +9,6 @@
 # Find <build_number> and <job_uuid> via:
 #   gh pr checks <PR> --repo vllm-project/vllm
 # Each failing row's URL is .../builds/<build_number>#<job_uuid>.
-#
-# Default output path: ci-<build>-<uuid_first_13_chars>.log (e.g.
-# ci-68478-019e6b07-daae.log). Jobs in the same build share the UUID's
-# first 8 chars, so the second segment is needed for uniqueness when
-# fetching multiple jobs in parallel. The script refuses to overwrite an
-# existing output file; pass an explicit path or set CI_FETCH_LOG_FORCE=1
-# to override.
 
 set -euo pipefail
 
@@ -33,12 +26,12 @@ if [ $# -lt 1 ]; then usage; fi
 if [[ "$1" == https://* ]]; then
     BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
     JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
-    OUT="${2:-}"
+    OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
 else
     if [ $# -lt 2 ]; then usage; fi
     BUILD="$1"
     JOB="$2"
-    OUT="${3:-}"
+    OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
 fi
 
 if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
@@ -46,18 +39,6 @@ if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
     usage
 fi
 
-# Jobs in the same build share the UUID's first segment, so include the
-# second segment (chars 9-13, e.g. "019e6b07-daae") to keep default filenames
-# unique when fetching multiple jobs from one build in parallel.
-if [ -z "$OUT" ]; then
-    OUT="ci-${BUILD}-${JOB:0:13}.log"
-fi
-
-if [ -e "$OUT" ] && [ -z "${CI_FETCH_LOG_FORCE:-}" ]; then
-    echo "Refusing to overwrite existing $OUT (set CI_FETCH_LOG_FORCE=1 or pass an explicit output path)." >&2
-    exit 1
-fi
-
 COOKIES=$(mktemp)
 trap 'rm -f "$COOKIES"' EXIT
 

.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh

@@ -49,7 +49,6 @@ for BACK in "${BACKENDS[@]}"; do
     --data-parallel-size 2 \
     --enable-expert-parallel \
     --enable-eplb \
-    --eplb-config '{"use_async": false}' \
     --trust-remote-code \
     --max-model-len 2048 \
     --all2all-backend "$BACK" \

.buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh

@@ -48,7 +48,7 @@ for BACK in "${BACKENDS[@]}"; do
     --enforce-eager \
     --enable-eplb \
     --all2all-backend "$BACK" \
-    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true, "use_async":false}' \
+    --eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
     --tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
     --data-parallel-size "${DATA_PARALLEL_SIZE}" \
     --enable-expert-parallel \

.buildkite/scripts/tool_call/run-bfcl-eval.sh

@@ -70,7 +70,7 @@ echo "============================================"
 # ---- Install bfcl-eval if missing ----
 if ! python3 -c "import bfcl_eval" 2>/dev/null; then
     echo "Installing bfcl-eval..."
-    uv pip install "bfcl-eval>=2025.10.20.1,<2026"
+    pip install "bfcl-eval>=2025.10.20.1,<2026"
 fi
 
 # ---- Cleanup handler ----
@@ -100,7 +100,7 @@ SERVE_ARGS=(
     --tensor-parallel-size "$TP_SIZE"
     --max-model-len "$MAX_MODEL_LEN"
     --enforce-eager
-    --enable-prefix-caching
+    --no-enable-prefix-caching
 )
 
 # Append reasoning parser if specified

.buildkite/test-amd.yaml

@@ -1261,7 +1261,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
 
 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
   timeout_in_minutes: 180
@@ -1484,7 +1484,7 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
-- label: DeepSeek V2-Lite Sync EPLB Accuracy (4xH100-4xMI300) # TBD
+- label: DeepSeek V2-Lite Accuracy (4xH100-4xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -1526,7 +1526,7 @@ steps:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
   - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
 
-- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (4xH100-4xMI300) # TBD
+- label: Qwen3-30B-A3B-FP8-block Accuracy (4xH100-4xMI300) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
   agent_pool: mi300_4
@@ -2768,7 +2768,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
 
 - label: Entrypoints Integration (API Server openai - Part 2) # TBD
   timeout_in_minutes: 180
@@ -2895,7 +2895,7 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx-fp8-and-mixed.txt
 
-- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200-MI355) # TBD
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200-MI355) # TBD
   timeout_in_minutes: 180
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
   agent_pool: mi355_2

.buildkite/test_areas/e2e_integration.yaml

@@ -2,8 +2,8 @@ group: E2E Integration
 depends_on: 
   - image-build
 steps:
-- label: DeepSeek V2-Lite Sync EPLB Accuracy
-  key: deepseek-v2-lite-sync-eplb-accuracy
+- label: DeepSeek V2-Lite Accuracy
+  key: deepseek-v2-lite-accuracy
   timeout_in_minutes: 60
   device: h100
   optional: true
@@ -12,8 +12,8 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
 
-- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy
+- label: Qwen3-30B-A3B-FP8-block Accuracy
+  key: qwen3-30b-a3b-fp8-block-accuracy
   timeout_in_minutes: 60
   device: h100
   optional: true
@@ -22,8 +22,8 @@ steps:
   commands:
   - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
 
-- label: Qwen3-30B-A3B-FP8-block Sync EPLB Accuracy (B200)
-  key: qwen3-30b-a3b-fp8-block-sync-eplb-accuracy-b200
+- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
+  key: qwen3-30b-a3b-fp8-block-accuracy-b200
   timeout_in_minutes: 60
   device: b200-k8s
   optional: true

.buildkite/test_areas/entrypoints.yaml

@@ -43,7 +43,7 @@ steps:
   - tests/entrypoints/test_chat_utils
   commands:
   - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
+  - pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
   mirror:
     amd:
       device: mi325_1

.buildkite/test_areas/model_executor.yaml

@@ -14,12 +14,5 @@ steps:
   commands:
     - apt-get update && apt-get install -y curl libsodium23
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    # Dump tracebacks of all threads if a test hangs, so a wedged GPU/CUDA
-    # init surfaces a stack instead of silently stalling.
-    - export PYTHONFAULTHANDLER=1
-    # Per-test watchdog: a single hung test (e.g. stuck during engine/CUDA
-    # init) fails fast with a traceback instead of running until the global
-    # build timeout. The `thread` method also handles hangs inside C/CUDA
-    # calls that the signal method cannot interrupt.
-    - pytest -v -s model_executor -m '(not slow_test)' --timeout=900 --timeout-method=thread
-    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
+    - pytest -v -s model_executor -m '(not slow_test)'
+    - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py

.buildkite/test_areas/rust_frontend.yaml

@@ -16,7 +16,7 @@ steps:
   - tests/benchmarks/test_serve_cli.py
   - tests/entrypoints/openai/chat_completion/test_chat_completion.py
   # - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
-
+  # - tests/entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
   # - tests/entrypoints/openai/completion/test_prompt_validation.py
   - tests/entrypoints/openai/completion/test_shutdown.py
   # - tests/entrypoints/openai/test_return_token_ids.py
@@ -28,7 +28,7 @@ steps:
   - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
   - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py
   # - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
-
+  # - pytest -v -s entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py
   # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
   - pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
   # - pytest -v -s entrypoints/openai/test_return_token_ids.py

.github/CODEOWNERS

@@ -40,12 +40,6 @@
 /vllm/entrypoints/chat_utils.py @DarkLight1337
 /vllm/entrypoints/llm.py @DarkLight1337
 
-# Rust Frontend
-/rust/ @BugenZhao @njhill
-/build_rust.sh @BugenZhao @njhill
-/rust-toolchain.toml @BugenZhao @njhill
-/.buildkite/test_areas/rust* @BugenZhao @njhill
-
 # Input/Output Processing
 /vllm/sampling_params.py @njhill @NickLucche
 /vllm/pooling_params.py @noooop @DarkLight1337
@@ -78,13 +72,11 @@
 /vllm/v1/worker/gpu/kv_connector.py @orozery
 
 # CI & building
-/.buildkite @Harry-Chen @khluu
-/docker/Dockerfile @Harry-Chen @khluu
-/pyproject.toml @khluu
-/setup.py @khluu
+/.buildkite @Harry-Chen
+/docker/Dockerfile @Harry-Chen
 
 # Test ownership
-/.buildkite/lm-eval-harness @mgoin
+/.buildkite/lm-eval-harness @mgoin 
 /tests/distributed/test_multi_node_assignment.py @youkaichao
 /tests/distributed/test_pipeline_parallel.py @youkaichao
 /tests/distributed/test_same_node.py @youkaichao

.readthedocs.yaml

@@ -9,8 +9,8 @@ build:
     python: "3.12"
   jobs:
     post_checkout:
-      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
       - bash docs/pre_run_check.sh
+      - git fetch origin main --unshallow --no-tags --filter=blob:none || true
     pre_create_environment:
       - pip install uv
     create_environment:

CMakeLists.txt

@@ -144,14 +144,14 @@ endif()
 # Set up GPU language and check the torch version and warn if it isn't
 # what is expected.
 #
-if (NOT HIP_FOUND AND NOT PYTORCH_FOUND_HIP AND CUDA_FOUND)
+if (NOT HIP_FOUND AND CUDA_FOUND)
   set(VLLM_GPU_LANG "CUDA")
 
   if (NOT Torch_VERSION VERSION_EQUAL ${TORCH_SUPPORTED_VERSION_CUDA})
     message(WARNING "Pytorch version ${TORCH_SUPPORTED_VERSION_CUDA} "
       "expected for CUDA build, saw ${Torch_VERSION} instead.")
   endif()
-elseif(HIP_FOUND OR PYTORCH_FOUND_HIP)
+elseif(HIP_FOUND)
   set(VLLM_GPU_LANG "HIP")
 
   # Importing torch recognizes and sets up some HIP/ROCm configuration but does
@@ -305,6 +305,10 @@ endif()
 #
 
 set(VLLM_EXT_SRC
+  "csrc/cache_kernels.cu"
+  "csrc/cache_kernels_fused.cu"
+  "csrc/attention/paged_attention_v1.cu"
+  "csrc/attention/paged_attention_v2.cu"
   "csrc/cuda_view.cu"
   "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu"
   "csrc/quantization/activation_kernels.cu"
@@ -315,8 +319,7 @@ set(VLLM_EXT_SRC
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_EXT_SRC
-    "csrc/minimax_reduce_rms_kernel.cu"
-    "csrc/minimax_m3_build_k2q_csr.cu")
+    "csrc/minimax_reduce_rms_kernel.cu")
 
   SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
 
@@ -638,18 +641,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"
     "csrc/libtorch_stable/pos_encoding_kernels.cu"
     "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu"
-    "csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu"
     "csrc/libtorch_stable/layernorm_kernels.cu"
     "csrc/libtorch_stable/layernorm_quant_kernels.cu"
     "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
     "csrc/libtorch_stable/attention/merge_attn_states.cu"
     "csrc/libtorch_stable/sampler.cu"
     "csrc/libtorch_stable/topk.cu"
-    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v1.cu"
-    "csrc/libtorch_stable/attention/paged_attention_v2.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
-    "csrc/libtorch_stable/cache_kernels_fused.cu")
+    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu")
 
   if(VLLM_GPU_LANG STREQUAL "CUDA")
     list(APPEND VLLM_STABLE_EXT_SRC
@@ -685,22 +683,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
                    "in CUDA target architectures.")
   endif()
 
-  # FP32 router GEMM (H=3072, E=256, M<=32). Requires SM90+ and CUDA >= 12.0.
-  cuda_archs_sm90plus(FP32_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND FP32_ROUTER_GEMM_ARCHS)
-    set(SRCS
-      "csrc/libtorch_stable/fp32_router_gemm_entry.cu"
-      "csrc/libtorch_stable/fp32_router_gemm.cu")
-    set_gencode_flags_for_srcs(
-      SRCS "${SRCS}"
-      CUDA_ARCHS "${FP32_ROUTER_GEMM_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    message(STATUS "Building fp32_router_gemm for archs: ${FP32_ROUTER_GEMM_ARCHS}")
-  else()
-    message(STATUS "Not building fp32_router_gemm as no compatible archs found "
-                   "(requires SM90+ and CUDA >= 12.0).")
-  endif()
-
   # Only build AllSpark kernels if we are building for at least some compatible archs.
   cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
   if (ALLSPARK_ARCHS)
@@ -942,11 +924,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
+    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
     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")
@@ -976,11 +960,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
-    set(NVFP4_KV_SRC "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
     set_gencode_flags_for_srcs(
       SRCS "${NVFP4_KV_SRC}"
       CUDA_ARCHS "${FP4_ARCHS}")
-    list(APPEND VLLM_STABLE_EXT_SRC "${NVFP4_KV_SRC}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
     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")
@@ -1258,22 +1242,24 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
                    " in CUDA target architectures")
   endif()
 
-  # DeepSeek V3 router GEMM kernel requires SM90+ and CUDA >= 12.0.
-  # (fp32_router_gemm has been migrated to _C_stable_libtorch above.)
-  cuda_archs_sm90plus(SM90PLUS_ROUTER_GEMM_ARCHS "${CUDA_ARCHS}")
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND SM90PLUS_ROUTER_GEMM_ARCHS)
+  # DeepSeek V3 router GEMM kernel - requires SM90+
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(DSV3_ROUTER_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_ROUTER_GEMM_ARCHS)
     set(DSV3_ROUTER_GEMM_SRC
       "csrc/moe/dsv3_router_gemm_entry.cu"
       "csrc/moe/dsv3_router_gemm_float_out.cu"
       "csrc/moe/dsv3_router_gemm_bf16_out.cu")
     set_gencode_flags_for_srcs(
       SRCS "${DSV3_ROUTER_GEMM_SRC}"
-      CUDA_ARCHS "${SM90PLUS_ROUTER_GEMM_ARCHS}")
+      CUDA_ARCHS "${DSV3_ROUTER_GEMM_ARCHS}")
     list(APPEND VLLM_MOE_EXT_SRC "${DSV3_ROUTER_GEMM_SRC}")
-
-    message(STATUS "Building DSV3 router GEMM kernels for archs: ${SM90PLUS_ROUTER_GEMM_ARCHS}")
+    message(STATUS "Building DSV3 router GEMM kernel for archs: ${DSV3_ROUTER_GEMM_ARCHS}")
   else()
-    message(STATUS "Not building DSV3 router GEMM kernels as no compatible archs found"
+    message(STATUS "Not building DSV3 router GEMM kernel as no compatible archs found"
                    " (requires SM90+ and CUDA >= 12.0)")
   endif()
 endif()
@@ -1300,14 +1286,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     "csrc/rocm/skinny_gemms.cu"
     "csrc/rocm/attention.cu")
 
-  set(VLLM_ROCM_HAS_GFX1100 OFF)
-  if(VLLM_GPU_ARCHES MATCHES "gfx1100")
-    set(VLLM_ROCM_HAS_GFX1100 ON)
-    list(APPEND VLLM_ROCM_EXT_SRC
-      "csrc/rocm/q_gemm_rdna3.cu"
-      "csrc/rocm/q_gemm_rdna3_wmma.cu")
-  endif()
-
   define_extension_target(
     _rocm_C
     DESTINATION vllm
@@ -1317,10 +1295,6 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
     ARCHITECTURES ${VLLM_GPU_ARCHES}
     USE_SABI 3
     WITH_SOABI)
-
-  if(VLLM_ROCM_HAS_GFX1100)
-    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
-  endif()
 endif()
 
 # Must run after the last HIP `define_extension_target` so every extension

benchmarks/kernels/benchmark_router_gemm.py

@@ -1,154 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import torch
-import torch.nn.functional as F
-
-from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-from vllm.transformers_utils.config import get_config
-from vllm.triton_utils import triton
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-# Dimensions supported by the DSV3 specialized kernel
-DSV3_SUPPORTED_NUM_EXPERTS = [256, 384]
-DSV3_SUPPORTED_HIDDEN_SIZES = [7168]
-
-# Dimensions supported by the gpt-oss specialized kernel
-GPT_OSS_SUPPORTED_NUM_EXPERTS = [32, 128]
-GPT_OSS_SUPPORTED_HIDDEN_SIZES = [2880]
-
-# Dimensions supported by the fp32 specialized kernel (MiniMax-M2)
-FP32_SUPPORTED_NUM_EXPERTS = [256]
-FP32_SUPPORTED_HIDDEN_SIZES = [3072]
-FP32_MAX_TOKENS = 32
-
-
-def get_batch_size_range(max_batch_size):
-    return [2**x for x in range(14) if 2**x <= max_batch_size]
-
-
-def get_model_params(config):
-    if config.architectures[0] in (
-        "DeepseekV2ForCausalLM",
-        "DeepseekV3ForCausalLM",
-        "DeepseekV32ForCausalLM",
-    ):
-        num_experts = config.n_routed_experts
-        hidden_size = config.hidden_size
-    elif config.architectures[0] in ("GptOssForCausalLM",) or config.architectures[
-        0
-    ] in ("MiniMaxM2ForCausalLM",):
-        num_experts = config.num_local_experts
-        hidden_size = config.hidden_size
-    else:
-        raise ValueError(f"Unsupported architecture: {config.architectures}")
-    return num_experts, hidden_size
-
-
-def get_benchmark(model, max_batch_size, trust_remote_code):
-    @triton.testing.perf_report(
-        triton.testing.Benchmark(
-            x_names=["batch_size"],
-            x_vals=get_batch_size_range(max_batch_size),
-            x_log=False,
-            line_arg="provider",
-            line_vals=[
-                "torch",
-                "vllm",
-            ],
-            line_names=["PyTorch", "vLLM"],
-            styles=([("blue", "-"), ("red", "-")]),
-            ylabel="TFLOPs",
-            plot_name=f"{model} router gemm throughput",
-            args={},
-        )
-    )
-    def benchmark(batch_size, provider):
-        config = get_config(model=model, trust_remote_code=trust_remote_code)
-        num_experts, hidden_size = get_model_params(config)
-
-        is_hopper_or_blackwell = current_platform.is_device_capability(
-            90
-        ) or current_platform.is_device_capability_family(100)
-        allow_dsv3_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in DSV3_SUPPORTED_NUM_EXPERTS
-            and hidden_size in DSV3_SUPPORTED_HIDDEN_SIZES
-        )
-        allow_gpt_oss_router_gemm = (
-            is_hopper_or_blackwell
-            and num_experts in GPT_OSS_SUPPORTED_NUM_EXPERTS
-            and hidden_size in GPT_OSS_SUPPORTED_HIDDEN_SIZES
-        )
-        is_fp32_router_model = (
-            is_hopper_or_blackwell
-            and num_experts in FP32_SUPPORTED_NUM_EXPERTS
-            and hidden_size in FP32_SUPPORTED_HIDDEN_SIZES
-        )
-        allow_fp32_router_gemm = is_fp32_router_model and batch_size <= FP32_MAX_TOKENS
-
-        # Weight dtype: fp32 kernel requires fp32 weights; others use bf16.
-        weight_dtype = torch.float32 if is_fp32_router_model else torch.bfloat16
-        mat_a = torch.randn(
-            (batch_size, hidden_size), dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-        mat_b = torch.randn(
-            (num_experts, hidden_size), dtype=weight_dtype, device="cuda"
-        ).contiguous()
-        bias = torch.randn(
-            num_experts, dtype=torch.bfloat16, device="cuda"
-        ).contiguous()
-
-        has_bias = allow_gpt_oss_router_gemm
-
-        quantiles = [0.5, 0.2, 0.8]
-
-        if provider == "torch":
-
-            def runner():
-                if allow_fp32_router_gemm:
-                    F.linear(mat_a.float(), mat_b)
-                elif has_bias:
-                    F.linear(mat_a, mat_b, bias)
-                else:
-                    F.linear(mat_a, mat_b)
-        elif provider == "vllm":
-
-            def runner():
-                if allow_dsv3_router_gemm:
-                    ops.dsv3_router_gemm(mat_a, mat_b, torch.bfloat16)
-                elif allow_fp32_router_gemm:
-                    ops.fp32_router_gemm(mat_a, mat_b)
-                elif allow_gpt_oss_router_gemm:
-                    ops.gpt_oss_router_gemm(mat_a, mat_b, bias)
-                elif is_fp32_router_model:
-                    # batch_size > FP32_MAX_TOKENS: fall back to F.linear
-                    F.linear(mat_a.float(), mat_b)
-                else:
-                    F.linear(mat_a, mat_b)
-
-        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            runner, quantiles=quantiles
-        )
-
-        def tflops(t_ms):
-            flops = 2 * batch_size * hidden_size * num_experts
-            return flops / (t_ms * 1e-3) / 1e12
-
-        return tflops(ms), tflops(max_ms), tflops(min_ms)
-
-    return benchmark
-
-
-if __name__ == "__main__":
-    parser = FlexibleArgumentParser()
-    parser.add_argument("--model", type=str, default="openai/gpt-oss-20b")
-    parser.add_argument("--max-batch-size", default=16, type=int)
-    parser.add_argument("--trust-remote-code", action="store_true")
-    args = parser.parse_args()
-
-    # Get the benchmark function
-    benchmark = get_benchmark(args.model, args.max_batch_size, args.trust_remote_code)
-    # Run performance benchmark
-    benchmark.run(print_data=True)

cmake/utils.cmake

@@ -476,16 +476,6 @@ function(cuda_archs_loose_intersection OUT_CUDA_ARCHS SRC_CUDA_ARCHS TGT_CUDA_AR
   set(${OUT_CUDA_ARCHS} ${_CUDA_ARCHS} PARENT_SCOPE)
 endfunction()
 
-
-function(cuda_archs_sm90plus OUT_CUDA_ARCHS TGT_CUDA_ARCHS)
-  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
-    cuda_archs_loose_intersection(_archs "9.0a;10.0f;11.0f" "${TGT_CUDA_ARCHS}")
-  else()
-    cuda_archs_loose_intersection(_archs "9.0a;10.0a;10.1a;10.3a" "${TGT_CUDA_ARCHS}")
-  endif()
-  set(${OUT_CUDA_ARCHS} ${_archs} PARENT_SCOPE)
-endfunction()
-
 #
 # Override the GPU architectures detected by cmake/torch and filter them by
 # `GPU_SUPPORTED_ARCHES`. Sets the final set of architectures in

csrc/attention/attention_kernels.cuh

@@ -17,18 +17,21 @@
  * limitations under the License.
  */
 
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
 #include <algorithm>
 
-#include "../../attention/attention_dtypes.h"
+#include "attention_dtypes.h"
 #include "attention_utils.cuh"
-#include "../../cuda_compat.h"
+#include "../cuda_compat.h"
 
 #ifdef USE_ROCM
   #include <hip/hip_bf16.h>
-  #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
 typedef __hip_bfloat16 __nv_bfloat16;
 #else
-  #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))

csrc/attention/attention_utils.cuh

@@ -18,8 +18,8 @@
  */
 #pragma once
 
-#include "../../cuda_compat.h"
-#include "../../attention/attention_dtypes.h"
+#include "../cuda_compat.h"
+#include "attention_dtypes.h"
 
 #include <float.h>
 #include <type_traits>

csrc/attention/paged_attention_v1.cu

@@ -16,9 +16,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "../torch_utils.h"
 #include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
+#include "../cuda_compat.h"
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -45,15 +44,13 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
           vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
           int NUM_THREADS = 128>
 void paged_attention_v1_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int num_kv_heads, float scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
+    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
   int num_seqs = query.size(0);
   int num_heads = query.size(1);
   int head_size = query.size(2);
@@ -72,8 +69,8 @@ void paged_attention_v1_launcher(
   T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
   CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
   CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
+  int* block_tables_ptr = block_tables.data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.data_ptr<int>();
   const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
   const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
 
@@ -88,9 +85,8 @@ void paged_attention_v1_launcher(
 
   dim3 grid(num_heads, num_seqs, 1);
   dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   switch (head_size) {
     // NOTE(woosuk): To reduce the compilation time, we only compile for the
     // head sizes that we use in the model. However, we can easily extend this
@@ -123,7 +119,7 @@ void paged_attention_v1_launcher(
       LAUNCH_PAGED_ATTENTION_V1(256);
       break;
     default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      TORCH_CHECK(false, "Unsupported head size: ", head_size);
       break;
   }
 }
@@ -145,43 +141,43 @@ void paged_attention_v1_launcher(
 
 // NOTE(woosuk): To reduce the compilation time, we omitted block sizes
 // 1, 2, 4, 64, 128, 256.
-#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
+#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
+  switch (block_size) {                                           \
+    case 8:                                                       \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
+      break;                                                      \
+    case 16:                                                      \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
   }
 
 void paged_attention_v1(
-    torch::stable::Tensor& out,    // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
+    torch::Tensor& out,    // [num_seqs, num_heads, head_size]
+    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::Tensor&
         key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
+    torch::Tensor&
         value_cache,       // [num_blocks, num_heads, head_size, block_size]
     int64_t num_kv_heads,  // [num_heads]
     double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
+    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor& seq_lens,      // [num_seqs]
     int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
     const int64_t blocksparse_local_blocks,
     const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
     const int64_t blocksparse_head_sliding_step) {
   const bool is_block_sparse = (blocksparse_vert_stride > 1);
 
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                              CALL_V1_LAUNCHER_BLOCK_SIZE)
 }
 

csrc/attention/paged_attention_v2.cu

@@ -16,9 +16,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#include "../torch_utils.h"
 #include "attention_kernels.cuh"
-#include "../../cuda_compat.h"
+#include "../cuda_compat.h"
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -45,16 +44,14 @@ template <typename T, typename CACHE_T, int BLOCK_SIZE,
           vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
           int NUM_THREADS = 128, int PARTITION_SIZE = 512>
 void paged_attention_v2_launcher(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int num_kv_heads, float scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int max_seq_len, const std::optional<torch::stable::Tensor>& alibi_slopes,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale,
-    const int tp_rank, const int blocksparse_local_blocks,
-    const int blocksparse_vert_stride, const int blocksparse_block_size,
-    const int blocksparse_head_sliding_step) {
+    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
+    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
   int num_seqs = query.size(0);
   int num_heads = query.size(1);
   int head_size = query.size(2);
@@ -76,8 +73,8 @@ void paged_attention_v2_launcher(
   T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
   CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
   CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
-  int* block_tables_ptr = block_tables.mutable_data_ptr<int>();
-  int* seq_lens_ptr = seq_lens.mutable_data_ptr<int>();
+  int* block_tables_ptr = block_tables.data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.data_ptr<int>();
   const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
   const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
 
@@ -94,9 +91,8 @@ void paged_attention_v2_launcher(
   int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
 
   dim3 block(NUM_THREADS);
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      query.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   switch (head_size) {
     // NOTE(woosuk): To reduce the compilation time, we only compile for the
     // head sizes that we use in the model. However, we can easily extend this
@@ -129,7 +125,7 @@ void paged_attention_v2_launcher(
       LAUNCH_PAGED_ATTENTION_V2(256);
       break;
     default:
-      STD_TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      TORCH_CHECK(false, "Unsupported head size: ", head_size);
       break;
   }
 }
@@ -152,48 +148,46 @@ void paged_attention_v2_launcher(
 
 // NOTE(woosuk): To reduce the compilation time, we omitted block sizes
 // 1, 2, 4, 64, 128, 256.
-#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)             \
-  switch (block_size) {                                               \
-    case 8:                                                           \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);             \
-      break;                                                          \
-    case 16:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);            \
-      break;                                                          \
-    case 32:                                                          \
-      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);            \
-      break;                                                          \
-    default:                                                          \
-      STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \
-      break;                                                          \
+#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
+  switch (block_size) {                                           \
+    case 8:                                                       \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
+      break;                                                      \
+    case 16:                                                      \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
   }
 
 void paged_attention_v2(
-    torch::stable::Tensor& out,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
-        exp_sums,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
-        max_logits,  // [num_seqs, num_heads, max_num_partitions]
-    torch::stable::Tensor&
+    torch::Tensor& out,         // [num_seqs, num_heads, head_size]
+    torch::Tensor& exp_sums,    // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor& max_logits,  // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor&
         tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
-    torch::stable::Tensor& query,  // [num_seqs, num_heads, head_size]
-    torch::stable::Tensor&
+    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::Tensor&
         key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
-    torch::stable::Tensor&
+    torch::Tensor&
         value_cache,       // [num_blocks, num_heads, head_size, block_size]
     int64_t num_kv_heads,  // [num_heads]
     double scale,
-    torch::stable::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
-    torch::stable::Tensor& seq_lens,      // [num_seqs]
+    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor& seq_lens,      // [num_seqs]
     int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
     const int64_t blocksparse_local_blocks,
     const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
     const int64_t blocksparse_head_sliding_step) {
   const bool is_block_sparse = (blocksparse_vert_stride > 1);
-  DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
                              CALL_V2_LAUNCHER_BLOCK_SIZE)
 }
 

csrc/cache_kernels_fused.cu

@@ -1,13 +1,15 @@
-#include "torch_utils.h"
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+
+#include "cuda_compat.h"
 #include "dispatch_utils.h"
 
-#include "../cuda_compat.h"
-
-#include "../quantization/w8a8/fp8/common.cuh"
+#include "quantization/w8a8/fp8/common.cuh"
 #ifdef USE_ROCM
-  #include "../quantization/w8a8/fp8/amd/quant_utils.cuh"
+  #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
 #else
-  #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh"
+  #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
 #endif
 
 #ifdef USE_ROCM
@@ -162,52 +164,43 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
 
 }  // namespace vllm
 
-#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE)  \
-  do {                                                                     \
-    VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                   \
-        q_pe.scalar_type(), "qk_scalar_type", [&] {                        \
-          using qk_t = scalar_t;                                           \
-          VLLM_STABLE_DISPATCH_FLOATING_TYPES(                             \
-              rope_cos_sin_cache.scalar_type(),                            \
-              "rope_cos_sin_cache_scalar_type", [&] {                      \
-                using cos_sin_t = scalar_t;                                \
-                if (rope_is_neox) {                                        \
-                  vllm::concat_and_cache_mla_rope_fused_kernel<            \
-                      qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE>  \
-                      <<<grid, block, 0, stream>>>(                        \
-                          positions.const_data_ptr<int64_t>(),             \
-                          q_pe.mutable_data_ptr<qk_t>(),                   \
-                          k_pe.mutable_data_ptr<qk_t>(),                   \
-                          kv_c.const_data_ptr<qk_t>(),                     \
-                          rope_cos_sin_cache.const_data_ptr<cos_sin_t>(),  \
-                          rot_dim, q_pe_stride_token, q_pe_stride_head,    \
-                          k_pe_stride, kv_c_stride, num_q_heads,           \
-                          reinterpret_cast<CACHE_T*>(                      \
-                              kv_cache.mutable_data_ptr()),                \
-                          slot_mapping.const_data_ptr<int64_t>(),          \
-                          block_stride, entry_stride, kv_lora_rank,        \
-                          block_size,                                      \
-                          kv_cache_quant_scale.const_data_ptr<float>());   \
-                } else {                                                   \
-                  vllm::concat_and_cache_mla_rope_fused_kernel<            \
-                      qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
-                      <<<grid, block, 0, stream>>>(                        \
-                          positions.const_data_ptr<int64_t>(),             \
-                          q_pe.mutable_data_ptr<qk_t>(),                   \
-                          k_pe.mutable_data_ptr<qk_t>(),                   \
-                          kv_c.const_data_ptr<qk_t>(),                     \
-                          rope_cos_sin_cache.const_data_ptr<cos_sin_t>(),  \
-                          rot_dim, q_pe_stride_token, q_pe_stride_head,    \
-                          k_pe_stride, kv_c_stride, num_q_heads,           \
-                          reinterpret_cast<CACHE_T*>(                      \
-                              kv_cache.mutable_data_ptr()),                \
-                          slot_mapping.const_data_ptr<int64_t>(),          \
-                          block_stride, entry_stride, kv_lora_rank,        \
-                          block_size,                                      \
-                          kv_cache_quant_scale.const_data_ptr<float>());   \
-                }                                                          \
-              });                                                          \
-        });                                                                \
+#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE)     \
+  do {                                                                        \
+    VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] {  \
+      using qk_t = scalar_t;                                                  \
+      VLLM_DISPATCH_FLOATING_TYPES(                                           \
+          rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
+          [&] {                                                               \
+            using cos_sin_t = scalar_t;                                       \
+            if (rope_is_neox) {                                               \
+              vllm::concat_and_cache_mla_rope_fused_kernel<                   \
+                  qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE>         \
+                  <<<grid, block, 0, stream>>>(                               \
+                      positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),   \
+                      k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),           \
+                      rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim,      \
+                      q_pe_stride_token, q_pe_stride_head, k_pe_stride,       \
+                      kv_c_stride, num_q_heads,                               \
+                      reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),        \
+                      slot_mapping.data_ptr<int64_t>(), block_stride,         \
+                      entry_stride, kv_lora_rank, block_size,                 \
+                      kv_cache_quant_scale.data_ptr<float>());                \
+            } else {                                                          \
+              vllm::concat_and_cache_mla_rope_fused_kernel<                   \
+                  qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE>        \
+                  <<<grid, block, 0, stream>>>(                               \
+                      positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(),   \
+                      k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(),           \
+                      rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim,      \
+                      q_pe_stride_token, q_pe_stride_head, k_pe_stride,       \
+                      kv_c_stride, num_q_heads,                               \
+                      reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),        \
+                      slot_mapping.data_ptr<int64_t>(), block_stride,         \
+                      entry_stride, kv_lora_rank, block_size,                 \
+                      kv_cache_quant_scale.data_ptr<float>());                \
+            }                                                                 \
+          });                                                                 \
+    });                                                                       \
   } while (false)
 
 // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
@@ -215,69 +208,64 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
 // Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and
 // concat_and_cache_mla.
 void concat_and_cache_mla_rope_fused(
-    torch::stable::Tensor& positions,  // [num_tokens]
-    torch::stable::Tensor& q_pe,       // [num_tokens, num_q_heads, rot_dim]
-    torch::stable::Tensor& k_pe,       // [num_tokens, rot_dim]
-    torch::stable::Tensor& kv_c,       // [num_tokens, kv_lora_rank]
-    torch::stable::Tensor& rope_cos_sin_cache,  // [max_position, rot_dim]
+    torch::Tensor& positions,           // [num_tokens]
+    torch::Tensor& q_pe,                // [num_tokens, num_q_heads, rot_dim]
+    torch::Tensor& k_pe,                // [num_tokens, rot_dim]
+    torch::Tensor& kv_c,                // [num_tokens, kv_lora_rank]
+    torch::Tensor& rope_cos_sin_cache,  // [max_position, rot_dim]
     bool rope_is_neox,
-    torch::stable::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
-    torch::stable::Tensor&
+    torch::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
+    torch::Tensor&
         kv_cache,  // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
-    const std::string& kv_cache_dtype,
-    torch::stable::Tensor& kv_cache_quant_scale) {
+    const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
   // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
   // slot_mapping.size(0) because of padding for CUDA graphs.
-  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0)
-  // because both include padding.
-  // In vLLM V1, however, key.size(0) can be larger than
-  // slot_mapping.size(0) since key includes padding for CUDA graphs,
-  // while slot_mapping does not. In this case,
-  // slot_mapping.size(0) represents the actual number of tokens
+  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
+  // both include padding.
+  // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
+  // since key includes padding for CUDA graphs, while slot_mapping does not.
+  // In this case, slot_mapping.size(0) represents the actual number of tokens
   // before padding.
-  // For compatibility with both cases, we use slot_mapping.size(0) as
-  // the number of tokens.
-  const int64_t num_tokens = slot_mapping.size(0);
-  const int64_t num_padded_tokens = q_pe.size(0);
-  STD_TORCH_CHECK(num_padded_tokens >= num_tokens);
+  // For compatibility with both cases, we use slot_mapping.size(0) as the
+  // number of tokens.
+  int num_tokens = slot_mapping.size(0);
+  int num_padded_tokens = q_pe.size(0);
+  TORCH_CHECK_GE(num_padded_tokens, num_tokens);
 
   const int num_q_heads = q_pe.size(1);
   const int rot_dim = q_pe.size(2);
   const int kv_lora_rank = kv_c.size(1);
 
-  STD_TORCH_CHECK(positions.size(0) == num_padded_tokens);
-  STD_TORCH_CHECK(positions.dim() == 1);
-  STD_TORCH_CHECK(positions.scalar_type() ==
-                  torch::headeronly::ScalarType::Long);
+  TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(positions.dim(), 1);
+  TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
 
-  STD_TORCH_CHECK(q_pe.dim() == 3);
-  STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens);
-  STD_TORCH_CHECK(q_pe.size(1) == num_q_heads);
-  STD_TORCH_CHECK(q_pe.size(2) == rot_dim);
+  TORCH_CHECK_EQ(q_pe.dim(), 3);
+  TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
+  TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
 
-  STD_TORCH_CHECK(k_pe.dim() == 2);
-  STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens);
-  STD_TORCH_CHECK(k_pe.size(1) == rot_dim);
-  STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type());
+  TORCH_CHECK_EQ(k_pe.dim(), 2);
+  TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
+  TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
 
-  STD_TORCH_CHECK(kv_c.dim() == 2);
-  STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens);
-  STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank);
-  STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type());
+  TORCH_CHECK_EQ(kv_c.dim(), 2);
+  TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
+  TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
+  TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
+  TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
 
-  STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim);
-  STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type());
+  TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
 
-  STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens);
-  STD_TORCH_CHECK(slot_mapping.scalar_type() ==
-                  torch::headeronly::ScalarType::Long);
+  TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
+  TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
 
-  STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim);
-  STD_TORCH_CHECK(kv_cache.dim() == 3);
+  TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
+  TORCH_CHECK_EQ(kv_cache.dim(), 3);
 
-  STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1);
-  STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() ==
-                  torch::headeronly::ScalarType::Float);
+  TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1);
+  TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float);
 
   int64_t q_pe_stride_token = q_pe.stride(0);
   int64_t q_pe_stride_head = q_pe.stride(1);
@@ -298,10 +286,9 @@ void concat_and_cache_mla_rope_fused(
   dim3 grid(num_tokens, 1, 1);
   dim3 block(thread_block_size, 1, 1);
 
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      positions.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(positions));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype,
+  DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
                              CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED);
 }

csrc/libtorch_stable/activation_kernels.cu

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

csrc/libtorch_stable/attention/merge_attn_states.cu

@@ -7,7 +7,7 @@
 #include <torch/headeronly/core/ScalarType.h>
 
 #include "../../attention/attention_dtypes.h"
-#include "attention_utils.cuh"
+#include "../../attention/attention_utils.cuh"
 #include "../../quantization/w8a8/fp8/common.cuh"
 
 namespace vllm {

csrc/libtorch_stable/fp32_router_gemm.cu

@@ -1,223 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// Router GEMM: activation(T) x weight(fp32) -> fp32, H=3072, E=256, M<=32.
-// Supports bf16 or fp32 activation; weight is always fp32.
-// Adapted from dsv3_router_gemm_float_out.cu.
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-// ---------------------------------------------------------------------------
-// Load helpers
-// ---------------------------------------------------------------------------
-
-// Load VPT fp32 values from the weight matrix (always fp32).
-//   VPT=4 when activation is fp32 (one float4 load)
-//   VPT=8 when activation is bf16 (two float4 loads)
-template <int VPT>
-__device__ __forceinline__ void load_weight(float const* ptr, float* dst);
-
-template <>
-__device__ __forceinline__ void load_weight<4>(float const* ptr, float* dst) {
-  float4 v = *reinterpret_cast<float4 const*>(ptr);
-  dst[0] = v.x;
-  dst[1] = v.y;
-  dst[2] = v.z;
-  dst[3] = v.w;
-}
-
-template <>
-__device__ __forceinline__ void load_weight<8>(float const* ptr, float* dst) {
-  float4 v0 = *reinterpret_cast<float4 const*>(ptr);
-  float4 v1 = *reinterpret_cast<float4 const*>(ptr + 4);
-  dst[0] = v0.x;
-  dst[1] = v0.y;
-  dst[2] = v0.z;
-  dst[3] = v0.w;
-  dst[4] = v1.x;
-  dst[5] = v1.y;
-  dst[6] = v1.z;
-  dst[7] = v1.w;
-}
-
-// Load VPT activation values and convert to fp32.
-template <typename T, int VPT>
-__device__ __forceinline__ void load_activation(T const* ptr, float* dst);
-
-// fp32 activation: one float4 load, no conversion needed.
-template <>
-__device__ __forceinline__ void load_activation<float, 4>(float const* ptr,
-                                                          float* dst) {
-  float4 v = *reinterpret_cast<float4 const*>(ptr);
-  dst[0] = v.x;
-  dst[1] = v.y;
-  dst[2] = v.z;
-  dst[3] = v.w;
-}
-
-// bf16 activation: one uint4 load (8 × bf16) + element-wise conversion.
-template <>
-__device__ __forceinline__ void load_activation<__nv_bfloat16, 8>(
-    __nv_bfloat16 const* ptr, float* dst) {
-  uint4 v = *reinterpret_cast<uint4 const*>(ptr);
-  __nv_bfloat16 const* bf16_ptr = reinterpret_cast<__nv_bfloat16 const*>(&v);
-#pragma unroll
-  for (int i = 0; i < 8; i++) dst[i] = __bfloat162float(bf16_ptr[i]);
-}
-
-// ---------------------------------------------------------------------------
-// Kernel
-// ---------------------------------------------------------------------------
-
-// InputT : type of activation (float or __nv_bfloat16)
-// Weight is always fp32; output is always fp32.
-// VPT = 16 / sizeof(InputT):  4 for fp32, 8 for bf16
-template <typename InputT, int kBlockSize, int kNumTokens, int kNumExperts,
-          int kHiddenDim>
-__global__ __launch_bounds__(128, 1) void fp32_router_gemm_kernel(
-    float* out, InputT const* mat_a, float const* mat_b) {
-  constexpr int VPT = 16 / sizeof(InputT);
-  constexpr int k_elems_per_k_iteration = VPT * kBlockSize;
-  constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration;
-  constexpr int kWarpSize = 32;
-  constexpr int kNumWarps = kBlockSize / kWarpSize;
-
-  int const n_idx = blockIdx.x;
-  int const tid = threadIdx.x;
-  int const warpId = tid / kWarpSize;
-  int const laneId = tid % kWarpSize;
-
-  float acc[kNumTokens] = {};
-  __shared__ float sm_reduction[kNumTokens][kNumWarps];
-
-  float const* b_col = mat_b + n_idx * kHiddenDim;
-
-  int k_bases[k_iterations];
-#pragma unroll
-  for (int ki = 0; ki < k_iterations; ki++) {
-    k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT;
-  }
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("griddepcontrol.wait;");
-#endif
-
-  for (int ki = 0; ki < k_iterations; ki++) {
-    int const k_base = k_bases[ki];
-
-    float b_float[VPT];
-    load_weight<VPT>(b_col + k_base, b_float);
-
-#pragma unroll
-    for (int m_idx = 0; m_idx < kNumTokens; m_idx++) {
-      float a_float[VPT];
-      load_activation<InputT, VPT>(mat_a + m_idx * kHiddenDim + k_base,
-                                   a_float);
-#pragma unroll
-      for (int k = 0; k < VPT; k++) {
-        acc[m_idx] += a_float[k] * b_float[k];
-      }
-    }
-  }
-
-  // Warp-level butterfly reduction
-#pragma unroll
-  for (int m = 0; m < kNumTokens; m++) {
-    float sum = acc[m];
-    sum += __shfl_xor_sync(0xffffffff, sum, 16);
-    sum += __shfl_xor_sync(0xffffffff, sum, 8);
-    sum += __shfl_xor_sync(0xffffffff, sum, 4);
-    sum += __shfl_xor_sync(0xffffffff, sum, 2);
-    sum += __shfl_xor_sync(0xffffffff, sum, 1);
-    if (laneId == 0) sm_reduction[m][warpId] = sum;
-  }
-
-  __syncthreads();
-
-  if (tid == 0) {
-#pragma unroll
-    for (int m = 0; m < kNumTokens; m++) {
-      float final_sum = 0.0f;
-#pragma unroll
-      for (int w = 0; w < kNumWarps; w++) final_sum += sm_reduction[m][w];
-      out[m * kNumExperts + n_idx] = final_sum;
-    }
-  }
-
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("griddepcontrol.launch_dependents;");
-#endif
-}
-
-// ---------------------------------------------------------------------------
-// Launcher
-// ---------------------------------------------------------------------------
-
-template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeFp32RouterGemm(float* output, InputT const* mat_a,
-                          float const* mat_b, cudaStream_t stream) {
-  constexpr int kBlockSize = 128;
-  cudaLaunchConfig_t config;
-  config.gridDim = kNumExperts;
-  config.blockDim = kBlockSize;
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.numAttrs = 1;
-  config.attrs = attrs;
-  cudaLaunchKernelEx(&config,
-                     fp32_router_gemm_kernel<InputT, kBlockSize, kNumTokens,
-                                             kNumExperts, kHiddenDim>,
-                     output, mat_a, mat_b);
-}
-
-// ---------------------------------------------------------------------------
-// Explicit instantiations: M=1..32, E=256, H=3072, for both input types
-// ---------------------------------------------------------------------------
-
-#define INSTANTIATE(T, M)                              \
-  template void invokeFp32RouterGemm<T, M, 256, 3072>( \
-      float*, T const*, float const*, cudaStream_t);
-
-#define INSTANTIATE_ALL(T) \
-  INSTANTIATE(T, 1)        \
-  INSTANTIATE(T, 2)        \
-  INSTANTIATE(T, 3)        \
-  INSTANTIATE(T, 4)        \
-  INSTANTIATE(T, 5)        \
-  INSTANTIATE(T, 6)        \
-  INSTANTIATE(T, 7)        \
-  INSTANTIATE(T, 8)        \
-  INSTANTIATE(T, 9)        \
-  INSTANTIATE(T, 10)       \
-  INSTANTIATE(T, 11)       \
-  INSTANTIATE(T, 12)       \
-  INSTANTIATE(T, 13)       \
-  INSTANTIATE(T, 14)       \
-  INSTANTIATE(T, 15)       \
-  INSTANTIATE(T, 16)       \
-  INSTANTIATE(T, 17)       \
-  INSTANTIATE(T, 18)       \
-  INSTANTIATE(T, 19)       \
-  INSTANTIATE(T, 20)       \
-  INSTANTIATE(T, 21)       \
-  INSTANTIATE(T, 22)       \
-  INSTANTIATE(T, 23)       \
-  INSTANTIATE(T, 24)       \
-  INSTANTIATE(T, 25)       \
-  INSTANTIATE(T, 26)       \
-  INSTANTIATE(T, 27)       \
-  INSTANTIATE(T, 28)       \
-  INSTANTIATE(T, 29)       \
-  INSTANTIATE(T, 30)       \
-  INSTANTIATE(T, 31)       \
-  INSTANTIATE(T, 32)
-
-INSTANTIATE_ALL(float)
-INSTANTIATE_ALL(__nv_bfloat16)
-
-#undef INSTANTIATE_ALL
-#undef INSTANTIATE

csrc/libtorch_stable/fp32_router_gemm_entry.cu

@@ -1,127 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor.h>
-#include <torch/headeronly/core/ScalarType.h>
-
-#include "core/registration.h"
-#include "libtorch_stable/torch_utils.h"
-
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-
-#include <stdexcept>
-
-namespace {
-
-inline int getSMVersion() {
-  auto* props = get_device_prop();
-  return props->major * 10 + props->minor;
-}
-
-}  // namespace
-
-static constexpr int FP32_NUM_EXPERTS = 256;
-static constexpr int FP32_HIDDEN_DIM = 3072;
-static constexpr int FP32_MAX_TOKENS = 32;
-
-// Forward declarations — 4 template params must match fp32_router_gemm.cu
-template <typename InputT, int kNumTokens, int kNumExperts, int kHiddenDim>
-void invokeFp32RouterGemm(float* output, InputT const* mat_a,
-                          float const* mat_b, cudaStream_t stream);
-
-// LoopUnroller templated on InputT
-template <typename InputT, int kBegin, int kEnd>
-struct Fp32LoopUnroller {
-  static void unroll(int num_tokens, float* output, InputT const* mat_a,
-                     float const* mat_b, cudaStream_t stream) {
-    if (num_tokens == kBegin) {
-      invokeFp32RouterGemm<InputT, kBegin, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
-          output, mat_a, mat_b, stream);
-    } else {
-      Fp32LoopUnroller<InputT, kBegin + 1, kEnd>::unroll(num_tokens, output,
-                                                         mat_a, mat_b, stream);
-    }
-  }
-};
-
-template <typename InputT, int kEnd>
-struct Fp32LoopUnroller<InputT, kEnd, kEnd> {
-  static void unroll(int num_tokens, float* output, InputT const* mat_a,
-                     float const* mat_b, cudaStream_t stream) {
-    if (num_tokens == kEnd) {
-      invokeFp32RouterGemm<InputT, kEnd, FP32_NUM_EXPERTS, FP32_HIDDEN_DIM>(
-          output, mat_a, mat_b, stream);
-    } else {
-      throw std::invalid_argument(
-          "fp32_router_gemm: num_tokens must be in [1, 32]");
-    }
-  }
-};
-
-void fp32_router_gemm(
-    torch::stable::Tensor& output,       // [num_tokens, num_experts]
-    torch::stable::Tensor const& mat_a,  // [num_tokens, hidden_dim]
-    torch::stable::Tensor const& mat_b   // [num_experts, hidden_dim]
-) {
-  STD_TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2);
-  STD_TORCH_CHECK(output.is_cuda() && mat_a.is_cuda() && mat_b.is_cuda(),
-                  "fp32_router_gemm: all tensors must be CUDA tensors");
-  STD_TORCH_CHECK(output.get_device_index() == mat_a.get_device_index() &&
-                      output.get_device_index() == mat_b.get_device_index(),
-                  "fp32_router_gemm: all tensors must be on the same device");
-  STD_TORCH_CHECK(
-      output.is_contiguous() && mat_a.is_contiguous() && mat_b.is_contiguous(),
-      "fp32_router_gemm: all tensors must be contiguous");
-
-  const int num_tokens = mat_a.size(0);
-  const int num_experts = mat_b.size(0);
-  const int hidden_dim = mat_a.size(1);
-
-  STD_TORCH_CHECK(output.size(0) == num_tokens && output.size(1) == num_experts,
-                  "fp32_router_gemm: output must have shape [num_tokens, "
-                  "num_experts]");
-  STD_TORCH_CHECK(
-      mat_a.size(1) == mat_b.size(1),
-      "fp32_router_gemm: mat_a and mat_b must have the same hidden_dim");
-  STD_TORCH_CHECK(hidden_dim == FP32_HIDDEN_DIM,
-                  "fp32_router_gemm: expected hidden_dim=3072");
-  STD_TORCH_CHECK(num_experts == FP32_NUM_EXPERTS,
-                  "fp32_router_gemm: expected num_experts=256");
-  STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS,
-                  "fp32_router_gemm: num_tokens must be in [0, 32]");
-  STD_TORCH_CHECK(
-      mat_a.scalar_type() == torch::headeronly::ScalarType::Float ||
-          mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16,
-      "fp32_router_gemm: mat_a must be float32 or bfloat16");
-  STD_TORCH_CHECK(mat_b.scalar_type() == torch::headeronly::ScalarType::Float,
-                  "fp32_router_gemm: mat_b (weight) must be float32");
-  STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Float,
-                  "fp32_router_gemm: output must be float32");
-
-  if (num_tokens == 0) {
-    return;
-  }
-
-  STD_TORCH_CHECK(getSMVersion() >= 90, "fp32_router_gemm: requires SM90+");
-
-  auto stream = get_current_cuda_stream(mat_a.get_device_index());
-  float* out_ptr = reinterpret_cast<float*>(output.mutable_data_ptr());
-  float const* mat_b_ptr = reinterpret_cast<float const*>(mat_b.data_ptr());
-
-  if (mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
-    auto const* mat_a_ptr =
-        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr());
-    Fp32LoopUnroller<__nv_bfloat16, 1, FP32_MAX_TOKENS>::unroll(
-        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
-  } else {
-    auto const* mat_a_ptr = reinterpret_cast<float const*>(mat_a.data_ptr());
-    Fp32LoopUnroller<float, 1, FP32_MAX_TOKENS>::unroll(
-        num_tokens, out_ptr, mat_a_ptr, mat_b_ptr, stream);
-  }
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
-  m.impl("fp32_router_gemm", TORCH_BOX(&fp32_router_gemm));
-}

csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu

@@ -1,615 +0,0 @@
-/*
- * SPDX-License-Identifier: Apache-2.0
- * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
- *
- * Horizontally-fused MiniMax-M3 attention pre-processing kernel.
- *
- * Replaces the per-token Python sequence in
- * ``MiniMaxM3SparseAttention.forward`` / ``MiniMaxM3Attention.forward``:
- *
- *     q  = q_norm(q);  k = k_norm(k);  q, k = rotary_emb(pos, q, k)
- *     index_q = index_q_norm(index_q);  index_k = index_k_norm(index_k)
- *     index_q, index_k = rotary_emb(pos, index_q, index_k)
- *     _insert_kv(k, v, index_k)
- *
- * All branches share head_dim=128 and the *same* partial-NeoX RoPE table
- * (``rotary_dim`` rotated, the trailing dims pass through).  The four norms
- * are Gemma-style RMSNorm (``x * rsqrt(mean(x^2)+eps) * (1 + weight)``) with
- * independent weights.
- *
- * Everything lives in a single fused ``qkv`` tensor.  The sparse layer's
- * fused projection (MinimaxM3QKVParallelLinearWithIndexer) emits, per token::
- *
- *     [ q | k | v | index_q | index_k ]   (the "5 results")
- *
- * while the dense layer emits just ``[ q | k | v ]``.  The kernel reads the
- * index branch straight out of that packed row -- no separate index tensors.
- *
- * One kernel, one grid; each warp owns one (token, head-slot) pair.  Slot
- * enumeration per token:
- *     [0, nq)                         Q  heads   -> norm(q_w)  + RoPE, write
- * qkv [nq, nq+nkv)                    K  heads   -> norm(k_w)  + RoPE, write
- * qkv
- *                                                  (+ insert into key cache)
- *     [nq+nkv, nq+2*nkv)              V  heads   -> insert into value cache
- *     IQ heads (niq)                             -> norm(iq_w) + RoPE, write iq
- *     IK       (1)                               -> norm(ik_w) + RoPE
- *                                                  (+ insert into index cache)
- *
- * The IQ/IK warps address the index_q/index_k sub-blocks *inside* qkv at the
- * fixed physical offsets (nq+2*nkv)*128 and (nq+2*nkv+niq)*128.
- *
- * Dense vs sparse is a compile-time choice via the ``kIsSparse``/``kInsertKV``
- * template bools (3 instantiations: dense <false,false>, sparse-profiling
- * <true,false>, sparse-serving <true,true>), so the index slots, the V slots
- * and the cache inserts fold away entirely on paths that don't use them. The
- * dense layer passes no caches/index: norm+RoPE happens in place and the
- * generic ``Attention`` layer owns the cache write.
- *
- * Q/K and (sparse) index_q/index_k are all rewritten in place inside the fused
- * ``qkv`` tensor.  Caches (bf16) are scatter-written by slot.
- */
-
-#include <cmath>
-#include <cuda_runtime.h>
-#include <type_traits>
-
-#include "torch_utils.h"
-
-#include "../cuda_compat.h"
-#include "../type_convert.cuh"
-#include "dispatch_utils.h"
-
-#ifndef FINAL_MASK
-  #ifdef USE_ROCM
-    #define FINAL_MASK 0xffffffffffffffffULL
-  #else
-    #define FINAL_MASK 0xffffffffu
-  #endif
-#endif
-
-namespace vllm {
-namespace minimax_m3_fused_ops {
-
-namespace {
-inline int getSMVersion() {
-  auto* props = get_device_prop();
-  return props->major * 10 + props->minor;
-}
-}  // namespace
-
-// ────────────────────────────────────────────────────────────────────────────
-// Constants (hard-coded for MiniMax-M3-preview).
-// ────────────────────────────────────────────────────────────────────────────
-constexpr int kHeadDim = 128;
-constexpr int kNumLanes = 32;
-constexpr int kElemsPerLane = kHeadDim / kNumLanes;  // 4
-
-// ────────────────────────────────────────────────────────────────────────────
-// Helpers
-// ────────────────────────────────────────────────────────────────────────────
-__device__ __forceinline__ float warpReduceSum(float val) {
-#pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1) {
-    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
-  }
-  return val;
-}
-
-// Gemma RMSNorm over the full head (no-op when ``weight == nullptr``) followed
-// by partial NeoX RoPE on the leading ``rotary_dim`` dims, all in fp32.  Each
-// lane owns ``kElemsPerLane`` contiguous dims [laneId*4, laneId*4+4).
-template <typename scalar_t>
-__device__ __forceinline__ void normAndRope(
-    float (&elems)[kElemsPerLane], int const laneId, float const eps,
-    scalar_t const* __restrict__ weight,  // [kHeadDim] or nullptr (no norm)
-    bool const do_rope, int const rotary_dim,
-    scalar_t const* __restrict__ cos_ptr,  // cos_sin_cache + pos*rotary_dim
-    bool const apply_norm) {
-  // ── Gemma RMSNorm: x * rsqrt(mean(x^2)+eps) * (1 + w) ──────────────────
-  if (apply_norm) {
-    float sumsq = 0.0f;
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) sumsq += elems[i] * elems[i];
-    sumsq = warpReduceSum(sumsq);
-    float const rms_rcp = rsqrtf(sumsq / static_cast<float>(kHeadDim) + eps);
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      int const dim = laneId * kElemsPerLane + i;
-      float const w = 1.0f + static_cast<float>(weight[dim]);
-      elems[i] = elems[i] * rms_rcp * w;
-    }
-  }
-
-  // ── Partial NeoX RoPE on dims [0, rotary_dim) ──────────────────────────
-  // half = rotary_dim/2.  Pair (i, i+half) for i in [0, half).  Lane L owns
-  // dims [4L, 4L+4); since half is a multiple of 4, a lane lies wholly in the
-  // first half (own=x[i]) or second half (own=x[i+half]); its partner lives
-  // ``half/4`` lanes away (XOR with that distance).
-  if (do_rope) {
-    int const half = rotary_dim / 2;
-    int const dim0 = laneId * kElemsPerLane;
-    bool const in_rope = dim0 < rotary_dim;
-    int const lane_xor = half / kElemsPerLane;  // partner-lane distance
-
-    float partner[kElemsPerLane];
-#pragma unroll
-    for (int i = 0; i < kElemsPerLane; i++) {
-      partner[i] = __shfl_xor_sync(FINAL_MASK, elems[i], lane_xor, 32);
-    }
-    if (in_rope) {
-      bool const first_half = dim0 < half;
-      int const i_base = first_half ? dim0 : (dim0 - half);  // cos/sin index
-      scalar_t const* sin_ptr = cos_ptr + half;
-#pragma unroll
-      for (int i = 0; i < kElemsPerLane; i++) {
-        float const c = static_cast<float>(cos_ptr[i_base + i]);
-        float const s = static_cast<float>(sin_ptr[i_base + i]);
-        if (first_half) {
-          elems[i] = elems[i] * c - partner[i] * s;
-        } else {
-          elems[i] = elems[i] * c + partner[i] * s;
-        }
-      }
-    }
-  }
-}
-
-// Load 4 contiguous bf16 -> 4 fp32 registers.
-template <typename scalar_t>
-__device__ __forceinline__ void loadElems(scalar_t const* __restrict__ src,
-                                          float (&elems)[kElemsPerLane]) {
-  using Converter = vllm::_typeConvert<scalar_t>;
-  uint2 v = *reinterpret_cast<uint2 const*>(src);
-  auto const* p =
-      reinterpret_cast<typename Converter::packed_hip_type const*>(&v);
-#pragma unroll
-  for (int i = 0; i < kElemsPerLane / 2; i++) {
-    float2 f2 = Converter::convert(p[i]);
-    elems[2 * i] = f2.x;
-    elems[2 * i + 1] = f2.y;
-  }
-}
-
-// Store 4 fp32 registers -> 4 contiguous bf16.
-template <typename scalar_t>
-__device__ __forceinline__ void storeElems(
-    scalar_t* __restrict__ dst, float const (&elems)[kElemsPerLane]) {
-  using Converter = vllm::_typeConvert<scalar_t>;
-  uint2 v;
-  auto* p = reinterpret_cast<typename Converter::packed_hip_type*>(&v);
-#pragma unroll
-  for (int i = 0; i < kElemsPerLane / 2; i++) {
-    p[i] = Converter::convert(make_float2(elems[2 * i], elems[2 * i + 1]));
-  }
-  *reinterpret_cast<uint2*>(dst) = v;
-}
-
-// ────────────────────────────────────────────────────────────────────────────
-// Kernel
-// ────────────────────────────────────────────────────────────────────────────
-// Grid: 1D, ceil(num_tokens * slots_per_token / warps_per_block).
-// Each warp = one (token, slot).
-//
-// `kIsSparse` and `kInsertKV` are compile-time template bools, so all the
-// branch decisions that distinguish the dense layer from the sparse layer
-// (index slots, KV/index inserts, V slots) fold away per instantiation.
-// Three instantiations are built: dense <false,false>, sparse-profiling
-// <true,false> and sparse-serving <true,true>.  Slots per token:
-//     Q : nq                            (always — norm+RoPE)
-//     K : nkv                           (always — norm+RoPE; +K-cache insert)
-//     V : nkv  only if kInsertKV        (V-cache insert; no warps in dense)
-//     IQ: niq  only if kIsSparse        (norm+RoPE)
-//     IK: 1    only if kIsSparse        (norm+RoPE; +index-cache insert)
-template <typename scalar_t, bool kIsSparse, bool kInsertKV>
-__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 const* __restrict__ q_norm_w,
-    scalar_t const* __restrict__ k_norm_w,
-    scalar_t const* __restrict__ iq_norm_w,
-    scalar_t const* __restrict__ ik_norm_w,
-    scalar_t const* __restrict__ cos_sin_cache,  // [max_pos, rotary_dim]
-    int64_t const* __restrict__ positions,       // [N] i64
-    int64_t const* __restrict__ slot_mapping,    // [N] i64 or nullptr
-    scalar_t* __restrict__ kv_cache,             // [nb,2,bs,nkv,128] or nullptr
-    scalar_t* __restrict__ index_cache,          // [nb*bs, 128] or nullptr
-    float const eps, int const rotary_dim, int const num_tokens, int const nq,
-    int const nkv, int const niq, int const block_size,
-    // kv_cache strides (in elements) for logical shape [nb, 2, bs, nkv, 128].
-    // The head_dim (last) dim is always innermost-contiguous (stride 1), so the
-    // NHD/HND layout choice is fully captured by these four strides: NHD keeps
-    // s_token < s_head, HND swaps them. dim_base addresses head_dim directly.
-    int64_t const kv_s_block, int64_t const kv_s_kv, int64_t const kv_s_token,
-    int64_t const kv_s_head) {
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  // _typeConvert<BFloat16> is unavailable on pre-Ampere; the M3 kernel only
-  // runs with bf16/fp16 inputs in practice.  Discard the bf16 body there.
-  if constexpr (std::is_same_v<scalar_t, c10::BFloat16>) {
-    return;
-  } else {
-#endif
-    int const warpsPerBlock = blockDim.x / 32;
-    int const laneId = threadIdx.x % 32;
-    int const globalWarpIdx = blockIdx.x * warpsPerBlock + (threadIdx.x / 32);
-
-    // Slot layout (compile-time gated: dense has neither V nor index slots).
-    int const v_slots = kInsertKV ? nkv : 0;
-    int const idx_slots = kIsSparse ? niq + 1 : 0;
-    int const slots_per_token = nq + nkv + v_slots + idx_slots;
-
-    int const tokenIdx = globalWarpIdx / slots_per_token;
-    int const slot = globalWarpIdx % slots_per_token;
-    if (tokenIdx >= num_tokens) return;
-
-    // Slot boundaries.
-    int const k_begin = nq;
-    int const v_begin = nq + nkv;             // valid only when kInsertKV
-    int const iq_begin = nq + nkv + v_slots;  // index block start
-    int const ik_slot = iq_begin + niq;       // valid only when kIsSparse
-
-    bool const isQ = slot < k_begin;
-    bool const isK = slot >= k_begin && slot < v_begin;
-    bool isV = false;
-    if constexpr (kInsertKV) isV = slot >= v_begin && slot < v_begin + nkv;
-    bool isIQ = false, isIK = false;
-    if constexpr (kIsSparse) {
-      isIQ = slot >= iq_begin && slot < ik_slot;
-      isIK = slot == ik_slot;
-    }
-
-    int const dim_base = laneId * kElemsPerLane;
-    // Physical row width of qkv: the dense layer packs [q|k|v]; the sparse
-    // layer additionally packs [index_q (niq heads) | index_k (1 head)].
-    int const qkv_row = (nq + 2 * nkv + (kIsSparse ? (niq + 1) : 0)) * kHeadDim;
-
-    // ── Resolve source pointer + per-branch parameters. ────────────────────
-    scalar_t* row_ptr = nullptr;       // in-place output location
-    scalar_t const* norm_w = nullptr;  // nullptr -> skip norm (V)
-    bool do_rope = true;
-    int head = 0;  // kv head index for inserts
-
-    if (isQ) {
-      row_ptr =
-          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
-      norm_w = q_norm_w;
-    } else if (isK) {
-      head = slot - k_begin;
-      row_ptr =
-          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
-      norm_w = k_norm_w;
-    } else if (isV) {
-      // qkv V section starts at slot index (nq + nkv): slot * kHeadDim is the
-      // correct in-tensor offset.
-      head = slot - v_begin;
-      row_ptr =
-          qkv + static_cast<int64_t>(tokenIdx) * qkv_row + slot * kHeadDim;
-      norm_w = nullptr;  // V: no norm, no rope
-      do_rope = false;
-    } else if (isIQ) {
-      // index_q sub-block lives at physical offset (nq+2*nkv)*128 in qkv.
-      int const ih = slot - iq_begin;
-      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
-                (nq + 2 * nkv + ih) * kHeadDim;
-      norm_w = iq_norm_w;
-    } else {  // isIK -- single shared index key at (nq+2*nkv+niq)*128.
-      row_ptr = qkv + static_cast<int64_t>(tokenIdx) * qkv_row +
-                (nq + 2 * nkv + niq) * kHeadDim;
-      norm_w = ik_norm_w;
-    }
-
-    // Store destination.  Q and index_q are gathered into dedicated contiguous
-    // output buffers (when provided) so the downstream SM100 sparse kernel's
-    // flat TMA descriptor can address them as [tokens*heads, head_dim]; this
-    // folds the de-interleaving into the store the kernel already does, instead
-    // of a separate q.contiguous() copy.  Everything else stays in place.
-    scalar_t* store_ptr = row_ptr;
-    if (isQ && q_out != nullptr) {
-      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;
-    }
-
-    // PDL: wait for the predecessor kernel (the qkv-projection GEMM that
-    // produces ``qkv``) to finish before touching any global memory.  No-op
-    // when PDL is not enabled on the launch.  The CUDA runtime wrapper emits
-    // the griddepcontrol.wait PTX with the required memory clobber internally.
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaGridDependencySynchronize();
-#endif
-
-    // ── Load -> norm+rope (fp32) -> store back in place. ───────────────────
-    float elems[kElemsPerLane];
-    loadElems<scalar_t>(row_ptr + dim_base, elems);
-
-    if (!isV) {
-      int64_t const pos = positions[tokenIdx];
-      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);
-    }
-
-    // ── Cache inserts (sparse serving only). ───────────────────────────────
-    if constexpr (kInsertKV) {
-      // Guard (not early-return) so every thread reaches the PDL trigger below.
-      int64_t const sm = (isK || isV || isIK) ? 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);
-        } 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);
-          // the physical NHD/HND layout is honoured via the passed strides.
-          int64_t const b = sm / block_size;
-          int64_t const t = sm % block_size;
-          int const kv = isK ? 0 : 1;
-          int64_t const off =
-              b * kv_s_block + kv * kv_s_kv + t * kv_s_token + head * kv_s_head;
-          storeElems<scalar_t>(kv_cache + off + dim_base, elems);
-        }
-      }
-    }
-
-    // PDL: signal that this kernel is done so a dependent successor may launch
-    // early.  No-op when PDL is not enabled on the launch.
-#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-    cudaTriggerProgrammaticLaunchCompletion();
-#endif
-#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
-  }
-#endif
-}
-
-// ────────────────────────────────────────────────────────────────────────────
-// Launch wrapper
-// ────────────────────────────────────────────────────────────────────────────
-template <typename scalar_t>
-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,
-                          scalar_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) {
-  // 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;
-  int const slots_per_token = nq + nkv + v_slots + idx_slots;
-
-  constexpr int kBlockSize = 256;
-  constexpr int kWarpsPerBlock = kBlockSize / 32;
-  int64_t const total_warps =
-      static_cast<int64_t>(num_tokens) * slots_per_token;
-  int const grid =
-      static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
-  if (grid == 0) return;
-
-#ifndef USE_ROCM
-  // PDL: enable programmatic stream serialization whenever the hardware
-  // supports it (SM90+).  On pre-Hopper GPUs the attribute is unavailable, so
-  // leave numAttrs = 0 and launch as a regular kernel via cudaLaunchKernelEx.
-  static int const sm_version = getSMVersion();
-  cudaLaunchConfig_t config;
-  config.gridDim = dim3(grid);
-  config.blockDim = dim3(kBlockSize);
-  config.dynamicSmemBytes = 0;
-  config.stream = stream;
-  cudaLaunchAttribute attrs[1];
-  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
-  attrs[0].val.programmaticStreamSerializationAllowed = 1;
-  config.attrs = attrs;
-  config.numAttrs = (sm_version >= 90) ? 1 : 0;
-
-  #define LAUNCH(IS_SPARSE, INSERT)                                            \
-    cudaLaunchKernelEx(                                                        \
-        &config,                                                               \
-        fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, 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, 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)
-#else
-  // ROCm: standard kernel launch syntax (no PDL/stream serialization).
-  // clang-format off
-  #define LAUNCH(IS_SPARSE, INSERT)                                          \
-    fusedMiniMaxM3QNormRopeKVInsertKernel<scalar_t, IS_SPARSE, INSERT>       \
-        <<<grid, kBlockSize, 0, stream>>>(                                   \
-            qkv, q_out, index_q_out, q_norm_w, k_norm_w, iq_norm_w,          \
-            ik_norm_w, cos_sin_cache, positions, slot_mapping, 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)
-  // clang-format on
-#endif
-
-  if (has_index) {
-    if (insert_kv) {
-      LAUNCH(true, true);  // sparse serving
-    } else {
-      LAUNCH(true, false);  // sparse profiling
-    }
-  } else {
-    // Dense layer: never has an index branch and never inserts here (the
-    // generic Attention layer owns the KV insert).
-    LAUNCH(false, false);
-  }
-#undef LAUNCH
-}
-
-}  // namespace minimax_m3_fused_ops
-}  // namespace vllm
-
-// ────────────────────────────────────────────────────────────────────────────
-// Torch op wrapper
-// ────────────────────────────────────────────────────────────────────────────
-void fused_minimax_m3_qknorm_rope_kv_insert(
-    torch::stable::Tensor& qkv,  // [N, qkv_row] (packs index if sparse)
-    torch::stable::Tensor const& q_norm_weight,  // [128]
-    torch::stable::Tensor const& k_norm_weight,  // [128]
-    torch::stable::Tensor const& cos_sin_cache,  // [max_pos, rotary_dim]
-    torch::stable::Tensor const& positions,      // [N] i64
-    int64_t num_heads, int64_t num_kv_heads, int64_t rotary_dim, double eps,
-    std::optional<torch::stable::Tensor> index_q_norm_weight,  // [128]
-    std::optional<torch::stable::Tensor> index_k_norm_weight,  // [128]
-    int64_t num_index_heads,                            // niq; 0 => dense
-    std::optional<torch::stable::Tensor> slot_mapping,  // [N] i64
-    std::optional<torch::stable::Tensor> kv_cache,      // [nb,2,bs,nkv,128]
-    std::optional<torch::stable::Tensor> index_cache,   // [nb,bs,128]
-    int64_t block_size,
-    std::optional<torch::stable::Tensor> q_out,  // [N, nq*128] contiguous
-    std::optional<torch::stable::Tensor>
-        index_q_out) {  // [N, niq*128] contiguous
-  STD_TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous(),
-                  "qkv must be contiguous CUDA");
-  STD_TORCH_CHECK(
-      positions.is_cuda() &&
-          positions.scalar_type() == torch::headeronly::ScalarType::Long,
-      "positions must be int64 CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.is_cuda() && cos_sin_cache.is_contiguous(),
-                  "cos_sin_cache must be contiguous CUDA");
-  STD_TORCH_CHECK(cos_sin_cache.scalar_type() == qkv.scalar_type(),
-                  "cos_sin_cache dtype must match qkv");
-  STD_TORCH_CHECK(
-      cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == rotary_dim,
-      "cos_sin_cache shape [max_pos, rotary_dim]");
-
-  STD_TORCH_CHECK(q_norm_weight.scalar_type() == qkv.scalar_type() &&
-                      k_norm_weight.scalar_type() == qkv.scalar_type(),
-                  "q/k norm weight dtype must match qkv");
-  STD_TORCH_CHECK(
-      q_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim &&
-          k_norm_weight.numel() == vllm::minimax_m3_fused_ops::kHeadDim,
-      "q/k norm weight must have 128 elements");
-  STD_TORCH_CHECK(rotary_dim > 0 && rotary_dim % 8 == 0 &&
-                      rotary_dim <= vllm::minimax_m3_fused_ops::kHeadDim,
-                  "rotary_dim must be a positive multiple of 8 and <= 128");
-
-  int const num_tokens = static_cast<int>(qkv.size(0));
-  int const nq = static_cast<int>(num_heads);
-  int const nkv = static_cast<int>(num_kv_heads);
-  int const niq = static_cast<int>(num_index_heads);
-
-  // The sparse layer packs the index branch ([index_q (niq heads) | index_k
-  // (1 head)]) right after [q|k|v] in the same row; the dense layer does not.
-  bool const has_index = niq > 0;
-  bool const insert_kv = kv_cache.has_value();
-  int const kHeadDim = vllm::minimax_m3_fused_ops::kHeadDim;
-  int const expected_row =
-      (nq + 2 * nkv + (has_index ? niq + 1 : 0)) * kHeadDim;
-  STD_TORCH_CHECK(qkv.size(1) == expected_row,
-                  "qkv last dim must be (num_heads + 2*num_kv_heads"
-                  " + num_index_heads + 1) * 128 for sparse, "
-                  "(num_heads + 2*num_kv_heads) * 128 for dense");
-
-  // Only the sparse layer inserts here (dense lets the generic Attention layer
-  // own the KV write); there is no dense+insert kernel instantiation.
-  STD_TORCH_CHECK(
-      !insert_kv || has_index,
-      "insert mode (kv_cache) requires the index branch (sparse layer)");
-  if (has_index) {
-    STD_TORCH_CHECK(
-        index_q_norm_weight.has_value() && index_k_norm_weight.has_value(),
-        "index branch requires both index norm weights");
-    STD_TORCH_CHECK(index_q_norm_weight->scalar_type() == qkv.scalar_type() &&
-                        index_k_norm_weight->scalar_type() == qkv.scalar_type(),
-                    "index norm weights dtype must match qkv");
-    STD_TORCH_CHECK(index_q_norm_weight->numel() == kHeadDim &&
-                        index_k_norm_weight->numel() == kHeadDim,
-                    "index norm weights must have 128 elements");
-  }
-  // kv_cache strides (logical shape [nb, 2, bs, nkv, head_dim]). Read straight
-  // off the tensor so the kernel honours whatever physical layout the attention
-  // backend allocated (NHD: stride order (0,1,2,3,4); HND: (0,1,3,2,4)). No new
-  // op argument is needed -- the strides ride along with the tensor itself.
-  int64_t kv_s_block = 0, kv_s_kv = 0, kv_s_token = 0, kv_s_head = 0;
-  if (insert_kv) {
-    STD_TORCH_CHECK(
-        slot_mapping.has_value() &&
-            slot_mapping->scalar_type() == torch::headeronly::ScalarType::Long,
-        "insert mode requires int64 slot_mapping");
-    STD_TORCH_CHECK(kv_cache->scalar_type() == qkv.scalar_type(),
-                    "kv_cache dtype must match qkv (bf16 cache only)");
-    STD_TORCH_CHECK(index_cache.has_value() &&
-                        index_cache->scalar_type() == qkv.scalar_type(),
-                    "insert mode requires matching index_cache");
-    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)");
-    kv_s_block = kv_cache->stride(0);
-    kv_s_kv = kv_cache->stride(1);
-    kv_s_token = kv_cache->stride(2);
-    kv_s_head = kv_cache->stride(3);
-  }
-  // Optional contiguous gather targets: when given, the normed/roped q (and
-  // index_q) are written here instead of in place, so callers avoid a separate
-  // .contiguous() copy.  index_q_out only makes sense on the sparse path.
-  if (q_out.has_value()) {
-    STD_TORCH_CHECK(
-        q_out->is_cuda() && q_out->is_contiguous() &&
-            q_out->scalar_type() == qkv.scalar_type(),
-        "q_out must be a contiguous CUDA tensor matching qkv dtype");
-    STD_TORCH_CHECK(
-        q_out->numel() == static_cast<int64_t>(num_tokens) * nq * kHeadDim,
-        "q_out must have num_tokens * num_heads * 128 elements");
-  }
-  if (index_q_out.has_value()) {
-    STD_TORCH_CHECK(
-        has_index,
-        "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");
-    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");
-  }
-
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      qkv.get_device_index());
-  auto stream = get_current_cuda_stream(qkv.get_device_index());
-
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
-      qkv.scalar_type(), "fused_minimax_m3_qknorm_rope_kv_insert", [&] {
-        using st = scalar_t;
-        vllm::minimax_m3_fused_ops::launchFusedMiniMaxM3<st>(
-            reinterpret_cast<st*>(qkv.data_ptr()),
-            q_out.has_value() ? reinterpret_cast<st*>(q_out->data_ptr())
-                              : nullptr,
-            index_q_out.has_value()
-                ? reinterpret_cast<st*>(index_q_out->data_ptr())
-                : nullptr,
-            reinterpret_cast<st const*>(q_norm_weight.data_ptr()),
-            reinterpret_cast<st const*>(k_norm_weight.data_ptr()),
-            has_index
-                ? reinterpret_cast<st const*>(index_q_norm_weight->data_ptr())
-                : nullptr,
-            has_index
-                ? reinterpret_cast<st const*>(index_k_norm_weight->data_ptr())
-                : nullptr,
-            reinterpret_cast<st const*>(cos_sin_cache.data_ptr()),
-            reinterpret_cast<int64_t const*>(positions.data_ptr()),
-            insert_kv
-                ? reinterpret_cast<int64_t const*>(slot_mapping->data_ptr())
-                : nullptr,
-            insert_kv ? reinterpret_cast<st*>(kv_cache->data_ptr()) : nullptr,
-            (insert_kv && has_index)
-                ? reinterpret_cast<st*>(index_cache->data_ptr())
-                : nullptr,
-            static_cast<float>(eps), static_cast<int>(rotary_dim), num_tokens,
-            nq, nkv, niq, static_cast<int>(block_size), kv_s_block, kv_s_kv,
-            kv_s_token, kv_s_head, has_index, insert_kv, stream);
-      });
-}

csrc/libtorch_stable/layernorm_kernels.cu

@@ -78,7 +78,8 @@ __global__ void rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < VEC_SIZE; j++) {
       float x = static_cast<float>(src1.val[j]);
-      dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
+      float w = static_cast<float>(src2.val[j]);
+      dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
     }
     v_out[i] = dst;
   }
@@ -142,7 +143,8 @@ fused_add_rms_norm_kernel(
 #pragma unroll
     for (int j = 0; j < width; ++j) {
       float x = Converter::convert(res.data[j]);
-      out.data[j] = Converter::convert(x * s_variance) * w.data[j];
+      float wf = Converter::convert(w.data[j]);
+      out.data[j] = Converter::convert(x * s_variance * wf);
     }
     input_v[strided_id] = out;
   }
@@ -181,8 +183,8 @@ fused_add_rms_norm_kernel(
 
   for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
     float x = (float)residual[blockIdx.x * hidden_size + idx];
-    input[blockIdx.x * input_stride + idx] =
-        (scalar_t)(x * s_variance) * weight[idx];
+    float w = (float)weight[idx];
+    input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
   }
 }
 

csrc/libtorch_stable/layernorm_quant_kernels.cu

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

csrc/libtorch_stable/ops.h

@@ -231,23 +231,6 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
                         torch::stable::Tensor& position_ids,
                         int64_t forced_token_heads_per_warp);
 
-// Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
-// index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
-// the index branch and scatters k/v/index_k into their paged caches.
-void fused_minimax_m3_qknorm_rope_kv_insert(
-    torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
-    torch::stable::Tensor const& k_norm_weight,
-    torch::stable::Tensor const& cos_sin_cache,
-    torch::stable::Tensor const& positions, int64_t num_heads,
-    int64_t num_kv_heads, int64_t rotary_dim, double eps,
-    std::optional<torch::stable::Tensor> index_q_norm_weight,
-    std::optional<torch::stable::Tensor> index_k_norm_weight,
-    int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
-    std::optional<torch::stable::Tensor> kv_cache,
-    std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
-    std::optional<torch::stable::Tensor> q_out,
-    std::optional<torch::stable::Tensor> index_q_out);
-
 // Sampler kernels (shared CUDA/ROCm)
 void apply_repetition_penalties_(
     torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
@@ -293,8 +276,7 @@ void selective_scan_fwd(
 // Activation kernels (shared CUDA/ROCm)
 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);
+                        torch::stable::Tensor& input, double limit);
 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,
@@ -373,132 +355,3 @@ torch::stable::Tensor ggml_moe_a8_vec(torch::stable::Tensor X,
                                       int64_t tokens);
 
 int64_t ggml_moe_get_block_size(int64_t type);
-
-void paged_attention_v1(
-    torch::stable::Tensor& out, torch::stable::Tensor& query,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables,
-    torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-void paged_attention_v2(
-    torch::stable::Tensor& out, torch::stable::Tensor& exp_sums,
-    torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out,
-    torch::stable::Tensor& query, torch::stable::Tensor& key_cache,
-    torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale,
-    torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens,
-    int64_t block_size, int64_t max_seq_len,
-    const std::optional<torch::stable::Tensor>& alibi_slopes,
-    const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale,
-    torch::stable::Tensor& v_scale, const int64_t tp_rank,
-    const int64_t blocksparse_local_blocks,
-    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
-    const int64_t blocksparse_head_sliding_step);
-
-// Cache ops (shared CUDA/ROCm)
-void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
-                 int64_t block_size_in_bytes,
-                 const torch::stable::Tensor& block_mapping);
-
-// Batch swap: submit all block copies in a single driver call.
-void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
-                       const torch::stable::Tensor& dst_ptrs,
-                       const torch::stable::Tensor& sizes,
-                       bool is_src_access_order_any);
-
-void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
-                       torch::stable::Tensor& key_cache,
-                       torch::stable::Tensor& value_cache,
-                       torch::stable::Tensor& slot_mapping,
-                       const std::string& kv_cache_dtype,
-                       torch::stable::Tensor& k_scale,
-                       torch::stable::Tensor& v_scale);
-
-void reshape_and_cache_flash(
-    torch::stable::Tensor& key, torch::stable::Tensor& value,
-    torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
-    torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype,
-    torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale);
-
-void concat_and_cache_mla(torch::stable::Tensor& kv_c,
-                          torch::stable::Tensor& k_pe,
-                          torch::stable::Tensor& kv_cache,
-                          torch::stable::Tensor& slot_mapping,
-                          const std::string& kv_cache_dtype,
-                          torch::stable::Tensor& scale);
-
-// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
-void concat_and_cache_mla_rope_fused(
-    torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
-    torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c,
-    torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox,
-    torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache,
-    const std::string& kv_cache_dtype,
-    torch::stable::Tensor& kv_cache_quant_scale);
-
-// Just for unittest
-void convert_fp8(torch::stable::Tensor& dst_cache,
-                 torch::stable::Tensor& src_cache, const double scale,
-                 const std::string& kv_cache_dtype);
-
-void gather_and_maybe_dequant_cache(
-    torch::stable::Tensor const& src_cache,     // [NUM_BLOCKS, BLOCK_SIZE,
-                                                // ENTRIES...]
-    torch::stable::Tensor const& dst,           // [TOT_TOKENS, ENTRIES...]
-    torch::stable::Tensor const& block_table,   // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& cu_seq_lens,   // [BATCH+1]
-    torch::stable::Tensor const& token_to_seq,  // [MAX_TOKEN_ACROSS_CHUNKS]
-    int64_t num_tokens, const std::string& kv_cache_dtype,
-    torch::stable::Tensor const& scale,
-    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
-
-// TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
-void cp_gather_cache(
-    torch::stable::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE,
-                                               // ENTRIES...]
-    torch::stable::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
-    torch::stable::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& cu_seq_lens,  // [BATCH+1]
-    int64_t batch_size,
-    std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
-
-// Gather and upconvert FP8 KV cache to BF16 workspace
-void cp_gather_and_upconvert_fp8_kv_cache(
-    torch::stable::Tensor const& src_cache,         // [NUM_BLOCKS, BLOCK_SIZE,
-                                                    // 656]
-    torch::stable::Tensor const& dst,               // [TOT_TOKENS, 576]
-    torch::stable::Tensor const& block_table,       // [BATCH, BLOCK_INDICES]
-    torch::stable::Tensor const& seq_lens,          // [BATCH]
-    torch::stable::Tensor const& workspace_starts,  // [BATCH]
-    int64_t batch_size);
-
-// Indexer K quantization and cache function
-void indexer_k_quant_and_cache(
-    torch::stable::Tensor& k,             // [num_tokens, head_dim]
-    torch::stable::Tensor& kv_cache,      // [num_blocks, block_size,
-                                          // cache_stride]
-    torch::stable::Tensor& slot_mapping,  // [num_tokens]
-    int64_t quant_block_size,             // quantization block size
-    const std::string& scale_fmt);
-
-// Concatenate query nope and rope for MLA/DSA attention
-void concat_mla_q(
-    torch::stable::Tensor& ql_nope,  // [num_tokens, num_heads, nope_dim]
-    torch::stable::Tensor& q_pe,     // [num_tokens, num_heads, rope_dim]
-    torch::stable::Tensor& q_out);   // [num_tokens, num_heads, nope_dim +
-                                     // rope_dim]
-
-// Extract function to gather quantized K cache
-void cp_gather_indexer_k_quant_cache(
-    const torch::stable::Tensor& kv_cache,      // [num_blocks, block_size,
-                                                // cache_stride]
-    torch::stable::Tensor& dst_k,               // [num_tokens, head_dim]
-    torch::stable::Tensor& dst_scale,           // [num_tokens, head_dim /
-                                                // quant_block_size * 4]
-    const torch::stable::Tensor& block_table,   // [batch_size, num_blocks]
-    const torch::stable::Tensor& cu_seq_lens);  // [batch_size + 1]

csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu

@@ -17,7 +17,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "../../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -27,7 +27,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "../../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 #include "cuda_utils.h"
 
 #include "nvfp4_utils.cuh"

csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu

@@ -17,7 +17,7 @@
 #include <torch/csrc/stable/tensor.h>
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "../../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 
 #include <cuda_runtime_api.h>
 #include <cuda_runtime.h>

csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu

@@ -23,7 +23,7 @@
 
 #include "libtorch_stable/torch_utils.h"
 #include "libtorch_stable/dispatch_utils.h"
-#include "../../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 
 #include "cuda_utils.h"
 #include "launch_bounds_utils.h"

csrc/libtorch_stable/quantization/fp4/nvfp4_utils.cuh

@@ -20,7 +20,7 @@
 #include <cuda_fp8.h>
 #include <utility>
 
-#include "../../cuda_vec_utils.cuh"
+#include "cuda_vec_utils.cuh"
 
 #if defined(NVFP4_ENABLE_ELTS16) && defined(CUDA_VERSION) && \
     CUDA_VERSION >= 12090

csrc/libtorch_stable/torch_bindings.cpp

@@ -247,10 +247,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def(
       "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
 
-  // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+).
-  // conditionally compiled so impl registration is in source file
-  ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
-
   // reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
   ops.def(
       "rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
@@ -337,17 +333,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "bool is_neox, Tensor position_ids, "
       "int forced_token_heads_per_warp=-1) -> ()");
 
-  // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE + KV-insert.
-  ops.def(
-      "fused_minimax_m3_qknorm_rope_kv_insert("
-      "Tensor! qkv, Tensor q_norm_weight, Tensor k_norm_weight, "
-      "Tensor cos_sin_cache, Tensor positions, int num_heads, "
-      "int num_kv_heads, int rotary_dim, float eps, "
-      "Tensor? index_q_norm_weight, Tensor? index_k_norm_weight, "
-      "int num_index_heads, "
-      "Tensor? slot_mapping, Tensor!? kv_cache, Tensor!? index_cache, "
-      "int block_size, Tensor!? q_out, Tensor!? index_q_out) -> ()");
-
   // Apply repetition penalties to logits in-place.
   ops.def(
       "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
@@ -375,11 +360,9 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
   ops.def("mul_and_silu(Tensor! out, Tensor input) -> ()");
 
   // SwiGLU activation with input clamping.
-  // alpha scales the sigmoid (gate * sigmoid(alpha * gate)); beta is added to
-  // the up half (up + beta). Defaults alpha=1.0, beta=0.0 give silu(gate)*up.
   ops.def(
-      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
-      "float alpha=1.0, float beta=0.0) -> ()");
+      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
+      "-> ()");
 
   // Activation function used in GeGLU with `none` approximation.
   ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
@@ -491,33 +474,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor? initial_state_idx,"
       "Tensor? cu_chunk_seqlen,"
       "Tensor? last_chunk_indices) -> ()");
-
-  // Attention ops
-  // Compute the attention between an input query and the cached
-  // keys/values using PagedAttention.
-  ops.def(
-      "paged_attention_v1("
-      "    Tensor! out, Tensor query, Tensor key_cache,"
-      "    Tensor value_cache, int num_kv_heads, float scale,"
-      "    Tensor block_tables, Tensor seq_lens, int block_size,"
-      "    int max_seq_len, Tensor? alibi_slopes,"
-      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
-      "    int tp_rank, int blocksparse_local_blocks,"
-      "    int blocksparse_vert_stride, int blocksparse_block_size,"
-      "    int blocksparse_head_sliding_step) -> ()");
-
-  // PagedAttention V2.
-  ops.def(
-      "paged_attention_v2("
-      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
-      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
-      "    Tensor value_cache, int num_kv_heads, float scale,"
-      "    Tensor block_tables, Tensor seq_lens, int block_size,"
-      "    int max_seq_len, Tensor? alibi_slopes,"
-      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
-      "    int tp_rank, int blocksparse_local_blocks,"
-      "    int blocksparse_vert_stride, int blocksparse_block_size,"
-      "    int blocksparse_head_sliding_step) -> ()");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
@@ -584,8 +540,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   // Positional encoding kernels (shared CUDA/ROCm)
   ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding));
   ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope));
-  ops.impl("fused_minimax_m3_qknorm_rope_kv_insert",
-           TORCH_BOX(&fused_minimax_m3_qknorm_rope_kv_insert));
 
   // Sampler kernels (shared CUDA/ROCm)
   ops.impl("apply_repetition_penalties_",
@@ -627,9 +581,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8));
   ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec));
   ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd));
-
-  ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1));
-  ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2));
 }
 
 // These capability-check functions take only primitive args (no tensors), so
@@ -652,115 +603,4 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
   ops.impl("ggml_moe_get_block_size", TORCH_BOX(&ggml_moe_get_block_size));
 }
 
-// Cache ops
-STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
-  // Swap in (out) the cache blocks from src to dst.
-  ops.def(
-      "swap_blocks(Tensor src, Tensor! dst,"
-      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
-
-  // Batch swap: submit all block copies in a single driver call.
-  ops.def(
-      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
-      "                  Tensor sizes,"
-      "                  bool is_src_access_order_any=False) -> ()");
-
-  // Reshape the key and value tensors and cache them.
-  ops.def(
-      "reshape_and_cache(Tensor key, Tensor value,"
-      "                  Tensor! key_cache, Tensor! value_cache,"
-      "                  Tensor slot_mapping,"
-      "                  str kv_cache_dtype,"
-      "                  Tensor k_scale, Tensor v_scale) -> ()");
-
-  // Reshape the key and value tensors and cache them.
-  ops.def(
-      "reshape_and_cache_flash(Tensor key, Tensor value,"
-      "                        Tensor! key_cache,"
-      "                        Tensor! value_cache,"
-      "                        Tensor slot_mapping,"
-      "                        str kv_cache_dtype,"
-      "                        Tensor k_scale, Tensor v_scale) -> ()");
-
-  // Concat kv_c and k_pe and cache them.
-  ops.def(
-      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
-      "                     Tensor! kv_cache,"
-      "                     Tensor slot_mapping,"
-      "                     str kv_cache_dtype,"
-      "                     Tensor scale) -> ()");
-
-  // Rotate Q and K, then write to kv cache for MLA
-  ops.def(
-      "concat_and_cache_mla_rope_fused("
-      "                     Tensor positions,"
-      "                     Tensor! q_pe,"
-      "                     Tensor! k_pe,"
-      "                     Tensor kv_c,"
-      "                     Tensor cos_sin_cache,"
-      "                     bool is_neox,"
-      "                     Tensor slot_mapping,"
-      "                     Tensor! kv_cache,"
-      "                     str kv_cache_dtype,"
-      "                     Tensor kv_cache_scale) -> ()");
-
-  // Convert the key and value cache to fp8 data type.
-  ops.def(
-      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
-      "str kv_cache_dtype) -> ()");
-
-  // Gather cache blocks from src_cache to dst, dequantizing from
-  // src_cache's dtype to dst's dtype if necessary.
-  ops.def(
-      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
-      "                               Tensor block_table, Tensor cu_seq_lens, "
-      "                               Tensor token_to_seq, "
-      "                               int num_tokens, "
-      "                               str kv_cache_dtype, "
-      "                               Tensor scale, Tensor? seq_starts) -> ()");
-
-  ops.def(
-      "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
-      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
-
-  ops.def(
-      "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
-      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
-      "batch_size) -> ()");
-
-  ops.def(
-      "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
-      "slot_mapping, "
-      "int quant_block_size, str kv_cache_dtype) -> ()");
-
-  ops.def("concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
-
-  ops.def(
-      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
-      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CPU, ops) {
-  ops.impl("swap_blocks_batch", TORCH_BOX(&swap_blocks_batch));
-}
-
-STABLE_TORCH_LIBRARY_IMPL(_C_cache_ops, CUDA, ops) {
-  ops.impl("swap_blocks", TORCH_BOX(&swap_blocks));
-  ops.impl("reshape_and_cache", TORCH_BOX(&reshape_and_cache));
-  ops.impl("reshape_and_cache_flash", TORCH_BOX(&reshape_and_cache_flash));
-  ops.impl("concat_and_cache_mla", TORCH_BOX(&concat_and_cache_mla));
-  ops.impl("concat_and_cache_mla_rope_fused",
-           TORCH_BOX(&concat_and_cache_mla_rope_fused));
-  ops.impl("convert_fp8", TORCH_BOX(&convert_fp8));
-  ops.impl("gather_and_maybe_dequant_cache",
-           TORCH_BOX(&gather_and_maybe_dequant_cache));
-  ops.impl("cp_gather_cache", TORCH_BOX(&cp_gather_cache));
-  ops.impl("cp_gather_and_upconvert_fp8_kv_cache",
-           TORCH_BOX(&cp_gather_and_upconvert_fp8_kv_cache));
-  ops.impl("indexer_k_quant_and_cache", TORCH_BOX(&indexer_k_quant_and_cache));
-  ops.impl("concat_mla_q", TORCH_BOX(&concat_mla_q));
-  ops.impl("cp_gather_indexer_k_quant_cache",
-           TORCH_BOX(&cp_gather_indexer_k_quant_cache));
-}
-
 REGISTER_EXTENSION(_C_stable_libtorch)

csrc/minimax_m3_build_k2q_csr.cu

@@ -1,742 +0,0 @@
-// CUDA C++ q2k -> k2q CSR builder.
-//
-// Five-stage pipeline. q-ascending order within each CSR row is preserved
-// by partitioning q across (CTA, warp_in_CTA) units; each unit owns a
-// contiguous q-sub-range and reserves a contiguous slot range per row via
-// a precomputed exclusive prefix scan.
-//
-//   M:  build_row_map      -- round-robin packing of rows across batches
-//   H:  histogram + tile_counts
-//   PR: row prefix         -- single block per head, row_counts -> row_ptr
-//   PT: tile prefix        -- multi-block, scan tile_counts along (c, w) axis
-//   S:  scatter (sorted)   -- per-warp slot range, q-sequential within warp
-//
-// Per-warp partitioning: each CTA has kWarps warps; warp w of CTA c owns
-// q-range [c*q_per_cta + w*q_per_warp, c*q_per_cta + (w+1)*q_per_warp).
-// tile_counts is shaped [G * kWarps, H, total_rows]; the "row" dimension
-// of the prefix scan is the flattened (c * kWarps + w) index, scanned in
-// lexicographic order so that warp-local slot ranges concatenate to the
-// global q-sorted output.
-
-#include <torch/all.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-#include <algorithm>
-
-#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be CUDA")
-#define CHECK_CONTIGUOUS(x) \
-  TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
-#define CHECK_INT(x) \
-  TORCH_CHECK((x).scalar_type() == at::kInt, #x " must be int32")
-#define CHECK_INPUT(x) \
-  CHECK_CUDA(x);       \
-  CHECK_CONTIGUOUS(x); \
-  CHECK_INT(x)
-
-namespace {
-
-constexpr int kWarpSize = 32;
-
-__device__ __forceinline__ void advance_batch_only(int const* __restrict__ cu_q,
-                                                   int B, int q_abs, int& bi) {
-  while (bi < B && cu_q[bi + 1] <= q_abs) ++bi;
-}
-
-// Atomic increment of a 16-bit half within a 32-bit SMEM word; returns the
-// OLD 16-bit value (slot). Per-warp count must stay < 32768 so the low
-// half does not carry into the high half.
-//   base_int32 : int32 pointer; element i holds rows 2*i (low) and 2*i+1
-//   (high).
-__device__ __forceinline__ int atomic_inc_int16_packed(int* base_int32,
-                                                       int row) {
-  int idx = row >> 1;
-  int shift = (row & 1) << 4;  // 0 or 16
-  int delta = 1 << shift;
-  int old = atomicAdd(&base_int32[idx], delta);
-  return (old >> shift) & 0xFFFF;
-}
-
-// Read 16-bit half from packed int32 storage.
-__device__ __forceinline__ int read_int16_packed(int const* base_int32,
-                                                 int row) {
-  int v = base_int32[row >> 1];
-  int shift = (row & 1) << 4;
-  return (v >> shift) & 0xFFFF;
-}
-
-// ---------------------------------------------------------------------------
-// M: round-robin row map.
-// ---------------------------------------------------------------------------
-template <int kBlockK>
-__global__ void k2q_build_row_map_kernel(int const* __restrict__ cu_k,
-                                         int* __restrict__ row_map,
-                                         int* __restrict__ row_coords, int B,
-                                         int max_kv_blocks) {
-  int level = blockIdx.x;
-  if (level >= max_kv_blocks) return;
-  if (threadIdx.x != 0) return;
-  int rows_before = 0;
-  for (int b = 0; b < B; ++b) {
-    int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
-    rows_before += (rb < level ? rb : level);
-  }
-  int active_before = 0;
-  for (int b = 0; b < B; ++b) {
-    int rb = (cu_k[b + 1] - cu_k[b] + kBlockK - 1) / kBlockK;
-    if (rb > level) {
-      int row_linear = rows_before + active_before;
-      row_map[(size_t)b * max_kv_blocks + level] = row_linear;
-      if (row_coords != nullptr) {
-        row_coords[(size_t)row_linear * 2] = b;
-        row_coords[(size_t)row_linear * 2 + 1] = level;
-      }
-      ++active_before;
-    } else {
-      row_map[(size_t)b * max_kv_blocks + level] = -1;
-    }
-  }
-}
-
-// ---------------------------------------------------------------------------
-// H: per-warp histogram + tile_counts.
-// kWarps warps per CTA, each owns q-sub-range = q_per_cta / kWarps.
-// SMEM hist[kWarps, total_rows] int32 (stored as packed int16 cursor:
-// 2 entries per int32 word). Each warp counts to its own row.
-// At end-of-CTA, write tile_counts[c*kWarps + w, h, r] = smem_hist[w, r]
-// and atomicAdd(row_counts[h, r], sum over w of smem_hist[w, r]).
-// ---------------------------------------------------------------------------
-template <int kTopK, int kBlockK, int kWarps>
-__global__ void k2q_hist_kernel(int const* __restrict__ q2k,
-                                int const* __restrict__ cu_q,
-                                int const* __restrict__ row_map,
-                                int* __restrict__ row_counts,
-                                int* __restrict__ tile_counts, int H, int B,
-                                int S_Q, int total_rows, int max_kv_blocks,
-                                int q_per_cta, int q_per_warp) {
-  constexpr int kThreads = kWarps * kWarpSize;
-  extern __shared__ int smem_hist_int[];
-  int* smem_hist = smem_hist_int;
-  int tid = threadIdx.x;
-  int warp_id = tid >> 5;
-  int lane = tid & 31;
-  int c = blockIdx.x;
-  int q_start_cta = c * q_per_cta;
-  int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
-  int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
-  int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
-
-  constexpr int kInt4PerToken = kTopK / 4;
-  int packed_per_warp = (total_rows + 1) >> 1;
-  int* my_hist = smem_hist + warp_id * packed_per_warp;
-
-  for (int h = 0; h < H; ++h) {
-    for (int i = lane; i < packed_per_warp; i += kWarpSize) my_hist[i] = 0;
-    __syncthreads();
-
-    if (q_start_warp < q_end_warp) {
-      int bi = 0;
-      int qi = q_start_warp + lane;
-      advance_batch_only(cu_q, B, qi, bi);
-
-      int4 const* head_topk4 =
-          reinterpret_cast<int4 const*>(q2k + (size_t)h * S_Q * kTopK);
-
-      for (; qi < q_end_warp; qi += kWarpSize) {
-        advance_batch_only(cu_q, B, qi, bi);
-        int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
-
-        int4 buf[kInt4PerToken];
-#pragma unroll
-        for (int v = 0; v < kInt4PerToken; ++v) {
-          buf[v] = head_topk4[(size_t)qi * kInt4PerToken + v];
-        }
-#pragma unroll
-        for (int t = 0; t < kTopK; ++t) {
-          int kvb_local = reinterpret_cast<int const*>(buf)[t];
-          if (kvb_local >= 0 && kvb_local < max_kv_blocks) {
-            int row = my_row_map[kvb_local];
-            if (row >= 0 && row < total_rows) {
-              atomic_inc_int16_packed(my_hist, row);
-            }
-          }
-        }
-      }
-    }
-    __syncthreads();
-
-    int* head_row_counts = row_counts + (size_t)h * total_rows;
-    // Each warp writes its own slice of tile_counts (full int32) by
-    // unpacking int16 entries from SMEM.
-    int* my_tile =
-        tile_counts + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
-    for (int i = lane; i < total_rows; i += kWarpSize) {
-      my_tile[i] = read_int16_packed(my_hist, i);
-    }
-    __syncthreads();
-
-    // Sum across warps (int32 accumulator), atomicAdd to row_counts.
-    for (int i = tid; i < total_rows; i += kThreads) {
-      int sum = 0;
-#pragma unroll
-      for (int w = 0; w < kWarps; ++w) {
-        sum += read_int16_packed(smem_hist + w * packed_per_warp, i);
-      }
-      if (sum > 0) atomicAdd(&head_row_counts[i], sum);
-    }
-    if (h + 1 < H) __syncthreads();
-  }
-}
-
-// ---------------------------------------------------------------------------
-// PR: row prefix. One block per head.
-// ---------------------------------------------------------------------------
-template <int kThreads>
-__global__ void k2q_row_prefix_kernel(int const* __restrict__ row_counts,
-                                      int* __restrict__ row_ptr,
-                                      int const* __restrict__ row_coords,
-                                      int* __restrict__ scheduler_metadata,
-                                      int* __restrict__ work_count,
-                                      int total_rows, int target_q_per_cta,
-                                      int work_capacity) {
-  int h = blockIdx.x;
-  int tid = threadIdx.x;
-  __shared__ int scan_buf[kThreads];
-
-  int const* head_counts = row_counts + (size_t)h * total_rows;
-  int* head_rowptr = row_ptr + (size_t)h * (total_rows + 1);
-  int chunk = (total_rows + kThreads - 1) / kThreads;
-  int lo = tid * chunk;
-  int hi = min(lo + chunk, total_rows);
-
-  int local_sum = 0;
-  for (int i = lo; i < hi; ++i) local_sum += head_counts[i];
-  scan_buf[tid] = local_sum;
-  __syncthreads();
-
-  for (int off = 1; off < kThreads; off <<= 1) {
-    int add = (tid >= off) ? scan_buf[tid - off] : 0;
-    __syncthreads();
-    scan_buf[tid] += add;
-    __syncthreads();
-  }
-  int running = scan_buf[tid] - local_sum;
-  for (int i = lo; i < hi; ++i) {
-    int row_count = head_counts[i];
-    running += row_count;
-    head_rowptr[i + 1] = running;
-    if (scheduler_metadata != nullptr && work_count != nullptr &&
-        row_count > 0) {
-      int num_chunks = (row_count + target_q_per_cta - 1) / target_q_per_cta;
-      int base = atomicAdd(work_count, num_chunks);
-      int batch_idx = row_coords[(size_t)i * 2];
-      int kv_block_idx = row_coords[(size_t)i * 2 + 1];
-      for (int c = 0; c < num_chunks; ++c) {
-        int work_idx = base + c;
-        if (work_idx < work_capacity) {
-          int q_begin = c * target_q_per_cta;
-          int q_count = min(target_q_per_cta, row_count - q_begin);
-          int* meta = scheduler_metadata + (size_t)work_idx * 6;
-          meta[0] = h;
-          meta[1] = i;
-          meta[2] = q_begin;
-          meta[3] = q_count;
-          meta[4] = batch_idx;
-          meta[5] = kv_block_idx;
-        }
-      }
-    }
-  }
-}
-
-// ---------------------------------------------------------------------------
-// PT_smem: SMEM-staged tile prefix scan.
-// Each block handles kRowsPerBlock rows for one head h. Cooperative load
-// of tile_counts[*, h, base_r..base_r+M) into SMEM (better coalescing
-// than per-warp uncoalesced stride reads), then per-warp scan in SMEM,
-// then cooperative store back. Fuses row_ptr into the base.
-// ---------------------------------------------------------------------------
-template <int kThreads, int kRowsPerBlock>
-__global__ void k2q_tile_prefix_smem_kernel(int* __restrict__ tile_counts,
-                                            int const* __restrict__ row_ptr,
-                                            int H, int total_rows,
-                                            int G_total) {
-  static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
-  extern __shared__ int smem_tprefix[];
-  // smem layout: smem[r_off][g] for r_off in [0, M), g in [0, G_total).
-
-  int tid = threadIdx.x;
-  int lane = tid & 31;
-  int warp_id = tid >> 5;
-
-  // Grid: H * blocks_per_h. Each block stays within a single head h
-  // and processes kRowsPerBlock contiguous rows starting at b_in_h *
-  // kRowsPerBlock. (Earlier flat-grid mapping `h = block_job /
-  // total_rows; base_r = block_job - h*total_rows` skipped rows when
-  // total_rows was not a multiple of kRowsPerBlock and H > 1, because
-  // the last partial block of head h-1 left blocks of head h starting
-  // at a non-zero row offset.)
-  int blocks_per_h = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock;
-  int h = blockIdx.x / blocks_per_h;
-  int b_in_h = blockIdx.x - h * blocks_per_h;
-  if (h >= H) return;
-  int base_r = b_in_h * kRowsPerBlock;
-  if (base_r >= total_rows) return;
-  int actual_M = min(kRowsPerBlock, total_rows - base_r);
-
-  size_t stride_g = (size_t)H * total_rows;
-  int* base_ptr = tile_counts + (size_t)h * total_rows + base_r;
-  int total_elems = G_total * actual_M;
-
-  // Cooperative load. Pattern: thread tid -> (r_off=tid%M, g=tid/M),
-  // then strided. 32 lanes hit M r's × (32/M) g's, giving 32/M cache
-  // lines per warp (vs 32 in the naive stride-along-g pattern).
-  for (int i = tid; i < total_elems; i += kThreads) {
-    int r_off = i % actual_M;
-    int g = i / actual_M;
-    smem_tprefix[r_off * G_total + g] = base_ptr[g * stride_g + r_off];
-  }
-  __syncthreads();
-
-  // Per-warp scan: warp w scans row (base_r + w) if w < actual_M.
-  if (warp_id < actual_M) {
-    int abs_r = base_r + warp_id;
-    int rp = row_ptr[(size_t)h * (total_rows + 1) + abs_r];
-    int* my_smem = smem_tprefix + warp_id * G_total;
-    int running = rp;
-    for (int g0 = 0; g0 < G_total; g0 += kWarpSize) {
-      int g = g0 + lane;
-      int v = (g < G_total) ? my_smem[g] : 0;
-      int x = v;
-#pragma unroll
-      for (int off = 1; off < kWarpSize; off <<= 1) {
-        int nbr = __shfl_up_sync(0xFFFFFFFF, x, off);
-        if (lane >= off) x += nbr;
-      }
-      int excl = running + x - v;
-      if (g < G_total) my_smem[g] = excl;
-      int chunk_sum = __shfl_sync(0xFFFFFFFF, x, 31);
-      running += chunk_sum;
-    }
-  }
-  __syncthreads();
-
-  // Cooperative store back.
-  for (int i = tid; i < total_elems; i += kThreads) {
-    int r_off = i % actual_M;
-    int g = i / actual_M;
-    base_ptr[g * stride_g + r_off] = smem_tprefix[r_off * G_total + g];
-  }
-}
-
-// ---------------------------------------------------------------------------
-// S: scatter. kWarps warps per CTA, each owns q-sub-range. Per-warp SMEM
-// cursor and per-warp tile_offset slot range. Within a warp, q's are
-// processed sequentially; lanes 0..kTopK-1 handle the topK slots in
-// lockstep. Across distinct q's in the same warp, the lockstep ordering
-// guarantees q-monotonic atomicAdd on smem_cursor[r].
-// ---------------------------------------------------------------------------
-// kQPerIter * kTopK lanes are active per warp iter; remaining lanes idle.
-// For kTopK=16, kQPerIter=2 uses all 32 lanes; for kTopK=8, kQPerIter=4.
-// CORRECTNESS NOTE: relies on lane-ordered SMEM atomicAdd return values
-// within a single warp instruction (verified on B200; tests pass).
-//
-// SMEM cursor stored as packed int16 (two cursors per int32). Per-warp
-// row count must stay < 32768 (~q_per_warp * kTopK at max sink), which
-// holds for all task.md sizes up to 1024K.
-template <int kTopK, int kBlockK, int kWarps>
-__global__ void k2q_scatter_kernel(
-    int const* __restrict__ q2k, int const* __restrict__ cu_q,
-    int const* __restrict__ row_map, int const* __restrict__ abs_base,
-    int* __restrict__ q_idx, int* __restrict__ qsplit_idx,
-    int* __restrict__ split_counts, int H, int B, int S_Q, int total_rows,
-    int max_kv_blocks, int q_per_cta, int q_per_warp, int max_seqlen_q) {
-  constexpr int kQPerIter = kWarpSize / kTopK > 0 ? kWarpSize / kTopK : 1;
-  extern __shared__ int smem_cursor_int[];
-  int* smem_cursor = smem_cursor_int;
-  int tid = threadIdx.x;
-  int warp_id = tid >> 5;
-  int lane = tid & 31;
-  int c = blockIdx.x;
-  int q_start_cta = c * q_per_cta;
-  int q_end_cta = min(q_start_cta + q_per_cta, S_Q);
-  int q_start_warp = min(q_start_cta + warp_id * q_per_warp, q_end_cta);
-  int q_end_warp = min(q_start_warp + q_per_warp, q_end_cta);
-
-  int q_in_iter = lane / kTopK;
-  int slot_in_q = lane % kTopK;
-  bool lane_active = (lane < kQPerIter * kTopK);
-
-  // Per-warp packed cursor: total_rows int16 entries -> ceil(total_rows/2)
-  // int32.
-  int packed_per_warp = (total_rows + 1) >> 1;
-  int* my_cursor = smem_cursor + warp_id * packed_per_warp;
-
-  for (int h = 0; h < H; ++h) {
-    for (int i = lane; i < packed_per_warp; i += kWarpSize) my_cursor[i] = 0;
-    __syncwarp();
-
-    if (q_start_warp < q_end_warp) {
-      int bi = 0;
-      advance_batch_only(cu_q, B, q_start_warp, bi);
-
-      int const* head_q2k = q2k + (size_t)h * S_Q * kTopK;
-      int const* my_abs_base =
-          abs_base + ((size_t)(c * kWarps + warp_id) * H + h) * total_rows;
-      int* head_qidx = q_idx + (size_t)h * S_Q * kTopK;
-
-      // (Hot-row register cache experiment showed no measurable
-      // benefit; relying on L1 to keep row 0 / row total_rows-1
-      // hot since they're hit every iteration in sink workloads.)
-
-      constexpr int kUnroll = 16;
-      int qi_base = q_start_warp;
-      for (; qi_base + kUnroll * kQPerIter <= q_end_warp;
-           qi_base += kUnroll * kQPerIter) {
-        int kvb[kUnroll];
-        int qloc[kUnroll];
-        int batch[kUnroll];
-        int const* rmap[kUnroll];
-
-#pragma unroll
-        for (int u = 0; u < kUnroll; ++u) {
-          int qi_u = qi_base + u * kQPerIter + q_in_iter;
-          kvb[u] = -1;
-          qloc[u] = 0;
-          batch[u] = 0;
-          if (lane_active) {
-            advance_batch_only(cu_q, B, qi_u, bi);
-            qloc[u] = qi_u - cu_q[bi];
-            batch[u] = bi;
-            kvb[u] = head_q2k[(size_t)qi_u * kTopK + slot_in_q];
-          }
-          rmap[u] = row_map + (size_t)bi * max_kv_blocks;
-        }
-
-        int row[kUnroll];
-#pragma unroll
-        for (int u = 0; u < kUnroll; ++u) {
-          row[u] = -1;
-          if (lane_active && kvb[u] >= 0 && kvb[u] < max_kv_blocks)
-            row[u] = rmap[u][kvb[u]];
-        }
-
-        // Pre-issue all kUnroll abs_base loads in parallel before
-        // the atomic chain so memory pipeline runs concurrently
-        // with SMEM atomic-adds.
-        int abs_v[kUnroll];
-#pragma unroll
-        for (int u = 0; u < kUnroll; ++u) {
-          abs_v[u] =
-              (row[u] >= 0 && row[u] < total_rows) ? my_abs_base[row[u]] : 0;
-        }
-
-#pragma unroll
-        for (int u = 0; u < kUnroll; ++u) {
-          int r = row[u];
-          bool valid_edge = r >= 0 && r < total_rows;
-          unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
-          unsigned int group_mask =
-              (kTopK == 32) ? 0xFFFFFFFFu
-                            : (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
-          unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
-          int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
-          int valid_count = __popc(valid_mask & group_mask);
-          if (split_counts != nullptr && slot_in_q == 0) {
-            split_counts[((size_t)batch[u] * max_seqlen_q + qloc[u]) * H + h] =
-                valid_count;
-          }
-          if (valid_edge) {
-            int slot = atomic_inc_int16_packed(my_cursor, r);
-            int out_pos = abs_v[u] + slot;
-            head_qidx[out_pos] = qloc[u];
-            if (qsplit_idx != nullptr) {
-              qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
-                  qloc[u] | ((split_slot & 0xFF) << 24);
-            }
-          }
-        }
-      }
-      // Tail: 1-3 iters left.
-      for (; qi_base < q_end_warp; qi_base += kQPerIter) {
-        int my_qi = qi_base + q_in_iter;
-        bool valid_q = (my_qi < q_end_warp) && lane_active;
-        int kvb_local = -1;
-        int q_local = 0;
-        int batch_local = 0;
-        if (valid_q) {
-          advance_batch_only(cu_q, B, my_qi, bi);
-          batch_local = bi;
-          q_local = my_qi - cu_q[bi];
-          kvb_local = head_q2k[(size_t)my_qi * kTopK + slot_in_q];
-        }
-        int const* my_row_map = row_map + (size_t)bi * max_kv_blocks;
-        int row = -1;
-        if (valid_q && kvb_local >= 0 && kvb_local < max_kv_blocks) {
-          row = my_row_map[kvb_local];
-        }
-        bool valid_edge = row >= 0 && row < total_rows;
-        unsigned int valid_mask = __ballot_sync(0xFFFFFFFFu, valid_edge);
-        unsigned int group_mask =
-            (kTopK == 32) ? 0xFFFFFFFFu
-                          : (((1u << kTopK) - 1u) << (q_in_iter * kTopK));
-        unsigned int lower_lane_mask = lane == 0 ? 0u : ((1u << lane) - 1u);
-        int split_slot = __popc(valid_mask & group_mask & lower_lane_mask);
-        int valid_count = __popc(valid_mask & group_mask);
-        if (split_counts != nullptr && valid_q && slot_in_q == 0) {
-          split_counts[((size_t)batch_local * max_seqlen_q + q_local) * H + h] =
-              valid_count;
-        }
-        if (valid_edge) {
-          int slot = atomic_inc_int16_packed(my_cursor, row);
-          int out_pos = my_abs_base[row] + slot;
-          head_qidx[out_pos] = q_local;
-          if (qsplit_idx != nullptr) {
-            qsplit_idx[(size_t)h * S_Q * kTopK + out_pos] =
-                q_local | ((split_slot & 0xFF) << 24);
-          }
-        }
-      }
-    }
-    if (h + 1 < H) __syncthreads();
-  }
-}
-
-}  // anonymous namespace
-
-// ===========================================================================
-// Host orchestration
-// ===========================================================================
-
-template <int kTopK, int kBlockK>
-static void launch_pipeline(torch::Tensor q2k, torch::Tensor cu_q,
-                            torch::Tensor cu_k, torch::Tensor row_ptr,
-                            torch::Tensor q_idx, int total_rows,
-                            int max_kv_blocks,
-                            torch::Tensor scheduler_metadata = torch::Tensor(),
-                            torch::Tensor work_count = torch::Tensor(),
-                            torch::Tensor qsplit_idx = torch::Tensor(),
-                            torch::Tensor split_counts = torch::Tensor(),
-                            int target_q_per_cta = 1, int work_capacity = 0,
-                            int max_seqlen_q = 0) {
-  int H = (int)q2k.size(0);
-  int S_Q = (int)q2k.size(1);
-  int topK = (int)q2k.size(2);
-  TORCH_CHECK(topK == kTopK, "topK runtime != template kTopK");
-  int B = (int)cu_q.size(0) - 1;
-  auto device = q2k.device();
-  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-
-  AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
-                                (size_t)H * (total_rows + 1) * sizeof(int),
-                                stream));
-  AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
-                                (size_t)H * S_Q * kTopK * sizeof(int), stream));
-
-  auto opts = torch::TensorOptions().dtype(torch::kInt32).device(device);
-  auto row_counts = torch::zeros({H, total_rows}, opts);
-  auto row_map = torch::empty({B, max_kv_blocks}, opts);
-  bool emit_schedule = scheduler_metadata.defined();
-  auto row_coords =
-      emit_schedule ? torch::empty({total_rows, 2}, opts) : torch::Tensor();
-  int* scheduler_metadata_ptr =
-      emit_schedule ? scheduler_metadata.data_ptr<int>() : nullptr;
-  int* work_count_ptr = emit_schedule ? work_count.data_ptr<int>() : nullptr;
-  int* qsplit_idx_ptr = emit_schedule ? qsplit_idx.data_ptr<int>() : nullptr;
-  int* split_counts_ptr =
-      emit_schedule ? split_counts.data_ptr<int>() : nullptr;
-  int* row_coords_ptr = emit_schedule ? row_coords.data_ptr<int>() : nullptr;
-  if (emit_schedule) {
-    AT_CUDA_CHECK(cudaMemsetAsync(work_count_ptr, 0, sizeof(int), stream));
-    AT_CUDA_CHECK(cudaMemsetAsync(scheduler_metadata_ptr, 0,
-                                  (size_t)work_capacity * 6 * sizeof(int),
-                                  stream));
-  }
-
-  int dev = q2k.get_device();
-  int num_sms = 0;
-  AT_CUDA_CHECK(
-      cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev));
-
-  // -- Pick kWarps per CTA based on SMEM budget for cursor/hist ---------
-  // SMEM per CTA = kWarps * total_rows * sizeof(int) (for both H and S).
-  // Want at least 2 CTAs/SM for memory parallelism. SM100 SMEM = 228KB.
-  // Pick the largest kWarps that fits two CTAs/SM, capped at 4.
-  // SMEM cursor packed as int16 (2 entries per int32 word):
-  int per_warp_smem = ((total_rows + 1) >> 1) * (int)sizeof(int);
-  int kWarps_pick = 4;
-  while (kWarps_pick > 1 && (kWarps_pick * per_warp_smem) * 2 > 228 * 1024) {
-    kWarps_pick >>= 1;
-  }
-  if (kWarps_pick < 1) kWarps_pick = 1;
-
-  // -- Pick G (CTAs) ----------------------------------------------------
-  // For each (kWarps, per_warp_smem) pair, the SMEM-bound occupancy is
-  // 228KB / (kWarps*per_warp_smem) CTAs/SM. We size G as
-  // num_sms * occupancy so a single resident wave covers all CTAs and
-  // the memory pipeline runs at peak.
-  int per_cta_smem_bytes = kWarps_pick * per_warp_smem;
-  int max_ctas_per_sm =
-      std::max(1, (228 * 1024) / std::max(1, per_cta_smem_bytes));
-  if (max_ctas_per_sm > 8) max_ctas_per_sm = 8;
-  constexpr int kMinQPerCta = 256;
-  // Cap target_g at num_sms * 3 — empirically this balances
-  // per-CTA work-size against parallelism. Higher caps regress
-  // mid-size cases due to row_counts atomicAdd contention and
-  // smaller q_per_cta. SMEM-bound configurations naturally cap
-  // lower if max_ctas_per_sm < 3.
-  int target_g = num_sms * std::min(max_ctas_per_sm, 3);
-  int max_g_for_q = (S_Q + kMinQPerCta - 1) / kMinQPerCta;
-  int G = std::min({target_g, max_g_for_q, S_Q});
-  if (G < 1) G = 1;
-  int q_per_cta = (S_Q + G - 1) / G;
-  G = (S_Q + q_per_cta - 1) / q_per_cta;
-  int q_per_warp = (q_per_cta + kWarps_pick - 1) / kWarps_pick;
-  int G_total = G * kWarps_pick;
-
-  auto tile_counts = torch::empty({G_total, H, total_rows}, opts);
-
-  // -- Compile-time switch on kWarps for the templated kernels ---------
-  auto rmap_fn = k2q_build_row_map_kernel<kBlockK>;
-  auto rprefix_fn = k2q_row_prefix_kernel<1024>;
-  constexpr int kPtRowsPerBlock = 8;
-  constexpr int kPtThreads = 256;
-  auto tprefix_smem_fn =
-      k2q_tile_prefix_smem_kernel<kPtThreads, kPtRowsPerBlock>;
-
-  if (max_kv_blocks > 0) {
-    rmap_fn<<<max_kv_blocks, 32, 0, stream>>>(cu_k.data_ptr<int>(),
-                                              row_map.data_ptr<int>(),
-                                              row_coords_ptr, B, max_kv_blocks);
-  }
-
-  auto launch_hist_scatter = [&](auto kWarps_const) {
-    constexpr int W = decltype(kWarps_const)::value;
-    size_t smem_bytes = (size_t)W * per_warp_smem;
-    auto hist_fn = k2q_hist_kernel<kTopK, kBlockK, W>;
-    auto scat_fn = k2q_scatter_kernel<kTopK, kBlockK, W>;
-    AT_CUDA_CHECK(cudaFuncSetAttribute(
-        hist_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
-    AT_CUDA_CHECK(cudaFuncSetAttribute(
-        scat_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, (int)smem_bytes));
-
-    hist_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
-        q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
-        row_counts.data_ptr<int>(), tile_counts.data_ptr<int>(), H, B, S_Q,
-        total_rows, max_kv_blocks, q_per_cta, q_per_warp);
-
-    rprefix_fn<<<H, 1024, 0, stream>>>(
-        row_counts.data_ptr<int>(), row_ptr.data_ptr<int>(),
-        emit_schedule ? row_coords.data_ptr<int>() : nullptr,
-        scheduler_metadata_ptr, work_count_ptr, total_rows, target_q_per_cta,
-        work_capacity);
-
-    // Grid is H * blocks_per_h so each block stays within a single
-    // head; flat (H*total_rows) grid would skip rows when total_rows
-    // is not a multiple of kPtRowsPerBlock.
-    int blocks_per_h = (total_rows + kPtRowsPerBlock - 1) / kPtRowsPerBlock;
-    int pt_grid = H * blocks_per_h;
-    if (pt_grid < 1) pt_grid = 1;
-    size_t pt_smem = (size_t)kPtRowsPerBlock * G_total * sizeof(int);
-    AT_CUDA_CHECK(cudaFuncSetAttribute(
-        tprefix_smem_fn, cudaFuncAttributeMaxDynamicSharedMemorySize,
-        (int)pt_smem));
-    tprefix_smem_fn<<<pt_grid, kPtThreads, pt_smem, stream>>>(
-        tile_counts.data_ptr<int>(), row_ptr.data_ptr<int>(), H, total_rows,
-        G_total);
-
-    scat_fn<<<G, W * kWarpSize, smem_bytes, stream>>>(
-        q2k.data_ptr<int>(), cu_q.data_ptr<int>(), row_map.data_ptr<int>(),
-        tile_counts.data_ptr<int>(), q_idx.data_ptr<int>(), qsplit_idx_ptr,
-        split_counts_ptr, H, B, S_Q, total_rows, max_kv_blocks, q_per_cta,
-        q_per_warp, max_seqlen_q);
-  };
-
-  if (kWarps_pick == 4) {
-    launch_hist_scatter(std::integral_constant<int, 4>{});
-  } else if (kWarps_pick == 2) {
-    launch_hist_scatter(std::integral_constant<int, 2>{});
-  } else {
-    launch_hist_scatter(std::integral_constant<int, 1>{});
-  }
-}
-
-void run_minimax_m3_build_k2q_csr_with_schedule(
-    torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
-    torch::Tensor row_ptr, torch::Tensor q_idx,
-    torch::Tensor scheduler_metadata, torch::Tensor work_count,
-    torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
-    int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
-    int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q) {
-  CHECK_INPUT(q2k);
-  CHECK_INPUT(cu_q);
-  CHECK_INPUT(cu_k);
-  CHECK_INPUT(row_ptr);
-  CHECK_INPUT(q_idx);
-  CHECK_INPUT(scheduler_metadata);
-  CHECK_INPUT(work_count);
-  CHECK_INPUT(qsplit_idx);
-  CHECK_INPUT(split_counts);
-  TORCH_CHECK(blk_kv == 128, "build_k2q_csr only supports blk_kv == 128");
-  int H = (int)q2k.size(0);
-  int S_Q = (int)q2k.size(1);
-  int tr = (int)total_rows;
-  int mkv = (int)max_kv_blocks;
-  int target = (int)target_q_per_cta;
-  int capacity = (int)work_capacity;
-  int max_sq = (int)max_seqlen_q;
-  TORCH_CHECK(tr >= 0 && mkv >= 0 && target > 0 && capacity > 0 && max_sq >= 0,
-              "invalid schedule sizing arguments");
-  TORCH_CHECK(row_ptr.size(0) == H && row_ptr.size(1) == tr + 1,
-              "row_ptr shape mismatch");
-  TORCH_CHECK(q_idx.size(0) == H && q_idx.size(1) == (int64_t)S_Q * (int)topk,
-              "q_idx shape mismatch");
-  TORCH_CHECK(qsplit_idx.sizes() == q_idx.sizes(), "qsplit_idx shape mismatch");
-  TORCH_CHECK(
-      scheduler_metadata.size(0) == capacity && scheduler_metadata.size(1) == 6,
-      "scheduler_metadata shape mismatch");
-  TORCH_CHECK(work_count.numel() == 1,
-              "work_count must have one int32 element");
-  TORCH_CHECK(split_counts.dim() == 3 &&
-                  split_counts.size(0) == cu_q.size(0) - 1 &&
-                  split_counts.size(1) == max_sq && split_counts.size(2) == H,
-              "split_counts shape mismatch");
-  if (S_Q == 0 || tr == 0 || H == 0 || mkv == 0) {
-    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-    AT_CUDA_CHECK(cudaMemsetAsync(row_ptr.data_ptr<int>(), 0,
-                                  (size_t)H * (tr + 1) * sizeof(int), stream));
-    AT_CUDA_CHECK(cudaMemsetAsync(q_idx.data_ptr<int>(), 0xFF,
-                                  (size_t)H * S_Q * (int)topk * sizeof(int),
-                                  stream));
-    AT_CUDA_CHECK(
-        cudaMemsetAsync(work_count.data_ptr<int>(), 0, sizeof(int), stream));
-    if (split_counts.numel() > 0) {
-      AT_CUDA_CHECK(cudaMemsetAsync(split_counts.data_ptr<int>(), 0,
-                                    (size_t)split_counts.numel() * sizeof(int),
-                                    stream));
-    }
-    return;
-  }
-
-  if (topk == 16) {
-    launch_pipeline<16, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
-                             scheduler_metadata, work_count, qsplit_idx,
-                             split_counts, target, capacity, max_sq);
-  } else if (topk == 8) {
-    launch_pipeline<8, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
-                            scheduler_metadata, work_count, qsplit_idx,
-                            split_counts, target, capacity, max_sq);
-  } else if (topk == 32) {
-    launch_pipeline<32, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
-                             scheduler_metadata, work_count, qsplit_idx,
-                             split_counts, target, capacity, max_sq);
-  } else if (topk == 4) {
-    launch_pipeline<4, 128>(q2k, cu_q, cu_k, row_ptr, q_idx, tr, mkv,
-                            scheduler_metadata, work_count, qsplit_idx,
-                            split_counts, target, capacity, max_sq);
-  } else {
-    TORCH_CHECK(false, "unsupported topK ", topk,
-                " (expected 4, 8, 16, or 32)");
-  }
-}

csrc/nvfp4_kv_cache_kernels.cu

@@ -17,8 +17,11 @@
 #define NVFP4_ENABLE_ELTS16 1
 #include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
 
-#include "libtorch_stable/dispatch_utils.h"
-#include "libtorch_stable/torch_utils.h"
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include "dispatch_utils.h"
 
 namespace vllm {
 
@@ -181,13 +184,12 @@ __global__ void reshape_and_cache_nvfp4_kernel(
 // Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
 // Each KV side contains both data and scale:
 //   page = [K_data | K_scale | V_data | V_scale]
-void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
-                                      torch::stable::Tensor& value,
-                                      torch::stable::Tensor& key_cache,
-                                      torch::stable::Tensor& value_cache,
-                                      torch::stable::Tensor& slot_mapping,
-                                      torch::stable::Tensor& k_scale,
-                                      torch::stable::Tensor& v_scale) {
+void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
+                                      torch::Tensor& key_cache,
+                                      torch::Tensor& value_cache,
+                                      torch::Tensor& slot_mapping,
+                                      torch::Tensor& k_scale,
+                                      torch::Tensor& v_scale) {
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);
@@ -198,18 +200,17 @@ void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
   // key_cache is kv_cache[:, 0] with shape
   // [num_blocks, block_size, num_heads, full_dim] in logical order.
   // Strides encode the physical layout (HND or NHD).
-  STD_TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
-  STD_TORCH_CHECK(key_cache.size(3) == full_dim,
-                  "key_cache last dim must be data_dim + scale_dim, got ",
-                  key_cache.size(3), " expected ", full_dim);
+  TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
+  TORCH_CHECK(key_cache.size(3) == full_dim,
+              "key_cache last dim must be data_dim + scale_dim, got ",
+              key_cache.size(3), " expected ", full_dim);
 
   int block_size = key_cache.size(1);
 
-  STD_TORCH_CHECK(head_size % 16 == 0,
-                  "head_size must be divisible by 16 for NVFP4 KV cache");
-  STD_TORCH_CHECK(block_size % 4 == 0,
-                  "block_size must be divisible by 4 for NVFP4 KV cache "
-                  "swizzle");
+  TORCH_CHECK(head_size % 16 == 0,
+              "head_size must be divisible by 16 for NVFP4 KV cache");
+  TORCH_CHECK(block_size % 4 == 0,
+              "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
 
   // Detect physical layout from strides (based on full_dim).
   // HND: head stride > block_offset stride.
@@ -229,9 +230,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
   // Scale follows data within each KV side.
   int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
 
-  uint8_t* key_scale_ptr = key_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
-  uint8_t* value_scale_ptr =
-      value_cache.mutable_data_ptr<uint8_t>() + data_per_kv;
+  uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
+  uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
 
   // Scale strides: same page stride, inner strides from layout.
   int64_t scale_block_stride = data_block_stride;
@@ -244,8 +244,8 @@ void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
     scale_block_offset_stride = (int64_t)num_heads * scale_dim;
   }
 
-  const float* k_scale_ptr = k_scale.const_data_ptr<float>();
-  const float* v_scale_ptr = v_scale.const_data_ptr<float>();
+  const float* k_scale_ptr = k_scale.data_ptr<float>();
+  const float* v_scale_ptr = v_scale.data_ptr<float>();
 
   int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
   int total_groups = num_heads * groups_per_head;
@@ -256,22 +256,20 @@ void reshape_and_cache_nvfp4_dispatch(torch::stable::Tensor& key,
   dim3 grid(num_tokens);
   dim3 block(num_threads);
 
-  const torch::stable::accelerator::DeviceGuard device_guard(
-      key.get_device_index());
-  const cudaStream_t stream = get_current_cuda_stream();
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  VLLM_STABLE_DISPATCH_HALF_TYPES(
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(
       key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
         vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
             <<<grid, block, 0, stream>>>(
-                key.const_data_ptr<scalar_t>(),
-                value.const_data_ptr<scalar_t>(),
-                key_cache.mutable_data_ptr<uint8_t>(),
-                value_cache.mutable_data_ptr<uint8_t>(), key_scale_ptr,
-                value_scale_ptr, slot_mapping.const_data_ptr<int64_t>(),
-                k_scale_ptr, v_scale_ptr, key.stride(0), value.stride(0),
-                num_heads, head_size, block_size, data_block_stride,
-                data_head_stride, data_block_offset_stride, scale_block_stride,
-                scale_head_stride, scale_block_offset_stride);
+                key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
+                key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
+                key_scale_ptr, value_scale_ptr,
+                slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
+                key.stride(0), value.stride(0), num_heads, head_size,
+                block_size, data_block_stride, data_head_stride,
+                data_block_offset_stride, scale_block_stride, scale_head_stride,
+                scale_block_offset_stride);
       });
 }

csrc/ops.h

@@ -31,6 +31,29 @@ torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
   return new_tensor;
 }
 
+void paged_attention_v1(
+    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
+    int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+void paged_attention_v2(
+    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
+    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
+    int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
 // 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.
@@ -62,8 +85,7 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
 
 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_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
 
 void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
                         torch::Tensor& scale);
@@ -147,12 +169,4 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
     torch::Tensor const& norm_weight_k, torch::Tensor workspace,
     int64_t const q_size, int64_t const kv_size, int64_t const rank,
     int64_t const nranks, double const eps);
-
-void run_minimax_m3_build_k2q_csr_with_schedule(
-    torch::Tensor q2k, torch::Tensor cu_q, torch::Tensor cu_k,
-    torch::Tensor row_ptr, torch::Tensor q_idx,
-    torch::Tensor scheduler_metadata, torch::Tensor work_count,
-    torch::Tensor qsplit_idx, torch::Tensor split_counts, int64_t topk,
-    int64_t blk_kv, int64_t total_rows, int64_t max_kv_blocks,
-    int64_t target_q_per_cta, int64_t work_capacity, int64_t max_seqlen_q);
 #endif

csrc/quantization/w8a8/fp8/amd/quant_utils.cuh

@@ -6,7 +6,6 @@
 #include <hip/hip_bfloat16.h>
 
 #include "../../../../attention/attention_dtypes.h"
-#include <torch/headeronly/core/ScalarType.h>
 
 namespace vllm {
 #ifdef USE_ROCM
@@ -643,29 +642,27 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
     vllm::Fp8KVCacheDataType KV_CACHE_DTYPE =                                  \
         vllm::get_fp8_kv_cache_data_type(KV_DTYPE);                            \
     if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {                   \
-      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
+      if (SRC_DTYPE == at::ScalarType::Float) {                                \
         FN(float, float, vllm::Fp8KVCacheDataType::kAuto);                     \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
+      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
         FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);               \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
+      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
         FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);     \
       } else {                                                                 \
-        STD_TORCH_CHECK(false,                                                 \
-                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
+        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
       }                                                                        \
     } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) {         \
-      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
+      if (SRC_DTYPE == at::ScalarType::Float) {                                \
         FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);                \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
+      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
         FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);             \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
+      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
         FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);        \
       } else {                                                                 \
-        STD_TORCH_CHECK(false,                                                 \
-                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
+        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
       }                                                                        \
     } else {                                                                   \
-      STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
+      TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE);     \
     }
 
 }  // namespace fp8

csrc/quantization/w8a8/fp8/nvidia/quant_utils.cuh

@@ -1,7 +1,6 @@
 #pragma once
 
 #include "../../../../attention/attention_dtypes.h"
-#include <torch/headeronly/core/ScalarType.h>
 #include <assert.h>
 #include <float.h>
 #include <stdint.h>
@@ -547,40 +546,37 @@ __inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) {
     vllm::Fp8KVCacheDataType KV_CACHE_DTYPE =                                  \
         vllm::get_fp8_kv_cache_data_type(KV_DTYPE);                            \
     if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {                   \
-      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
+      if (SRC_DTYPE == at::ScalarType::Float) {                                \
         FN(float, float, vllm::Fp8KVCacheDataType::kAuto);                     \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
+      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
         FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto);               \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
+      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
         FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto);     \
       } else {                                                                 \
-        STD_TORCH_CHECK(false,                                                 \
-                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
+        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
       }                                                                        \
     } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E4M3) {         \
-      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
+      if (SRC_DTYPE == at::ScalarType::Float) {                                \
         FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);                \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
+      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
         FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);             \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
+      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
         FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3);        \
       } else {                                                                 \
-        STD_TORCH_CHECK(false,                                                 \
-                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
+        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
       }                                                                        \
     } else if (KV_CACHE_DTYPE == vllm::Fp8KVCacheDataType::kFp8E5M2) {         \
-      if (SRC_DTYPE == torch::headeronly::ScalarType::Float) {                 \
+      if (SRC_DTYPE == at::ScalarType::Float) {                                \
         FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);                \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::Half) {           \
+      } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
         FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);             \
-      } else if (SRC_DTYPE == torch::headeronly::ScalarType::BFloat16) {       \
+      } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
         FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2);        \
       } else {                                                                 \
-        STD_TORCH_CHECK(false,                                                 \
-                        "Unsupported input type of kv cache: ", SRC_DTYPE);    \
+        TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
       }                                                                        \
     } else {                                                                   \
-      STD_TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \
+      TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE);     \
     }
 
 }  // namespace fp8

csrc/rocm/ops.h

@@ -18,15 +18,6 @@ void wvSplitKQ(const at::Tensor& in_a, const at::Tensor& in_b,
                const at::Tensor& scale_a, const at::Tensor& scale_b,
                const int64_t CuCount);
 
-torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
-                              torch::Tensor b_qzeros, torch::Tensor b_scales,
-                              torch::Tensor b_g_idx, bool use_v2_format);
-
-torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
-                                   torch::Tensor b_qzeros,
-                                   torch::Tensor b_scales,
-                                   torch::Tensor b_g_idx, bool use_v2_format);
-
 void paged_attention(
     torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
     torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,

csrc/rocm/q_gemm_rdna3.cu

@@ -1,780 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// W4A16 GPTQ kernel for RDNA3 (gfx1100 / RX 7900 XTX class), templated on the
-// activation dtype (half or __hip_bfloat16). Adapted from exllamav2's 4-bit
-// kernel (csrc/quantization/gptq/q_gemm.cu) with the following changes:
-//
-//   1. Direct write to the T-typed output via packed CAS-loop on a 64-bit
-//      word (atomic_add_pk4_{f16,bf16}). gfx11 has no native
-//      v_global_atomic_pk_add_{f16,bf16}, so the kernel emulates one with
-//      global_atomic_cmpswap_b64. This avoids the M*N*4-byte FP32 scratch
-//      buffer + memset + cast-pass that an fp32-accumulator design would
-//      need; the caller passes a zero-initialised T-typed output tensor
-//      and every block atomically adds its partial sum into it.
-//
-//   2. The bf16 path uses a dedicated bit-trick that avoids the fp16-only
-//      "upper nibble * 16" trick, which would overflow the 7-bit bf16
-//      mantissa. See qdq_4_rdna3.cuh for details.
-//
-//   3. Wave32 geometry sized for high CU saturation: THREADS_X=256
-//      (8 waves per block) and BLOCK_KN_SIZE=256, with each thread
-//      computing 4 N output columns. gridDim.z = K / BLOCK_KN_SIZE
-//      splits K and the output is atomically accumulated. fp16 uses
-//      v_dot2_f32_f16 (__builtin_amdgcn_fdot2) for the inner dot;
-//      bf16 widens to fp32 (no v_pk_fma_bf16 on gfx11) and accumulates
-//      with v_fma_f32. M_COUNT ∈ {1,2,4,8} is selected at launch
-//      based on size_m.
-//
-//   4. The bf16 dispatch with M >= 16 forwards to the WMMA kernel in
-//      q_gemm_rdna3_wmma.cu (separate translation unit) where
-//      v_wmma_f32_16x16x16_bf16_w32 wins. The fp16 path always stays
-//      scalar (the bit-trick dequant beats WMMA below M=64).
-
-#include <cstdint>
-#include <cstdio>
-
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <ATen/cuda/CUDAContext.h>
-
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp16.h>
-
-#include "qdq_4_rdna3.cuh"
-
-#if defined(__HIPCC__) && defined(__gfx1100__)
-  #define __HIP__RDNA3__
-#endif
-
-namespace vllm {
-namespace gptq_rdna3 {
-
-// BLOCK_KN_SIZE = 256 (was 128 in exllama). Each block covers 256 K
-// elements and THREADS_X*4 = 1024 N columns. For Qwen-class K=4096 this
-// halves gridDim.z (32 → 16) and therefore halves the atomic count per
-// output position vs the exllama default. THREADS_X=256 = 8 waves on RDNA3
-// wave32; with ~32 wave slots per CU we still fit 4 blocks per CU at peak.
-//
-// We tried BLOCK_KN_SIZE=512 (microbench on Qwen3.6-27B): bf16 improved
-// 5-10% at large M (atomic CAS halved), but fp16 decode regressed up to
-// +40% on qkv-square (32 → 45 μs at M=1). Cause: 16 waves/block × 16
-// total blocks for [M=1, K=N=4096] only saturates ~8 of the 96 CUs,
-// breaking memory-latency hiding for the fp16 path which is already
-// memory-bound. Reverted to 256; bf16 keeps most of its gains from the
-// fp32 dequant rewrite alone.
-#define BLOCK_KN_SIZE 256
-#define THREADS_X 256
-
-// Device code below is RDNA3-only; non-RDNA3 device passes fall through to
-// the empty __global__ stub at the #else below for symbol parity.
-#if defined(__HIP__RDNA3__) || !defined(__HIP_DEVICE_COMPILE__)
-
-// ---------------------------------------------------------------------------
-// Per-dtype helpers. We avoid heavy template metaprogramming and just provide
-// overloaded inline functions; the kernel below selects via `if constexpr`.
-// ---------------------------------------------------------------------------
-
-// Type-generic zero — both half and bf16_t in HIP/ROCm have a converting
-// constructor from float, but going through __float2half_rn / __float2bfloat16
-// is the unambiguously correct path on every ROCm version.
-template <typename T>
-__forceinline__ __device__ T tzero();
-
-template <>
-__forceinline__ __device__ half tzero<half>() {
-  return __float2half_rn(0.0f);
-}
-
-template <>
-__forceinline__ __device__ bf16_t tzero<bf16_t>() {
-  return __float2bfloat16(0.0f);
-}
-
-__forceinline__ __device__ float dot22_8_f(half2 (&dq)[4], const half* a_ptr) {
-  // RDNA3 has v_dot2_f32_f16 (`__builtin_amdgcn_fdot2`) which computes
-  // fp32 += a.x*b.x + a.y*b.y in a single instruction with the accumulator
-  // staying in fp32 throughout. hipcc 7.2 does NOT peephole the obvious
-  // `__hfma2 + cast + add` pattern into v_dot2 (verified by ISA
-  // disassembly: 0 v_dot2_f32_f16 vs 256 v_cvt_f32_f16 + 218 v_add_f32 in
-  // the M_COUNT=8 kernel before this change), so we issue the builtin
-  // explicitly. Saves the trailing 2× v_cvt_f32_f16 + v_add_f32 (3 ops)
-  // per dot22_8_f call vs the half2-accumulator form. With 128 calls per
-  // K=32 step that's ~384 ops/K-step less issue pressure on the VALU.
-  //
-  // Numerical bonus: accumulator stays fp32 throughout the dot. The old
-  // form accumulated 8 muladds in fp16 (10-bit mantissa) before casting,
-  // which could lose ~3 bits of precision on borderline magnitudes.
-  float result = 0.0f;
-  const half2* a2_ptr = (const half2*)a_ptr;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    result = __builtin_amdgcn_fdot2(dq[i], *a2_ptr++, result, /*clamp=*/false);
-  }
-  return result;
-}
-
-__forceinline__ __device__ float dot22_8_f(bf162_t (&dq)[4],
-                                           const bf16_t* a_ptr) {
-  // RDNA3 (gfx1100) lacks a packed bf16 FMA: there is no v_pk_fma_bf16 in
-  // the gfx11 ISA (it only landed on CDNA3+ / gfx94x and later). hipcc
-  // therefore lowers __hfma2(bf162_t, bf162_t, bf162_t) to a serialised
-  // fallback (single-element FMAs or fp32 round-trips), which empirically
-  // runs ~2× the cycle count of v_pk_fma_f16 on the same VALU. The bf16
-  // decode path was paying that tax in full, scaling linearly with M (the
-  // fp16 path scales sub-linearly because its v_pk_fma_f16 is full rate
-  // and the kernel becomes memory-bound).
-  //
-  // Fix: widen bf16 → fp32 explicitly (a left-shift by 16, free in VGPRs)
-  // and accumulate with v_fma_f32, which IS full rate on RDNA3. Same FMA
-  // count, but each FMA is fast. Bonus: the accumulator is now fp32
-  // throughout instead of bf16, which is also numerically more accurate
-  // (no compounding bf16-rounding inside the dot loop).
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw, dw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    __builtin_memcpy(&dw, &dq[i], sizeof(uint32_t));
-    // bf16 in low 16 bits  → fp32 by left-shifting into the upper half.
-    // bf16 in high 16 bits → already aligned with fp32's upper half.
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    float d_x = __uint_as_float((dw & 0xFFFFu) << 16);
-    float d_y = __uint_as_float(dw & 0xFFFF0000u);
-    result = __fmaf_rn(d_x, a_x, result);
-    result = __fmaf_rn(d_y, a_y, result);
-  }
-  return result;
-}
-
-// fp32-input dot product: paired with dequant_4bit_8_bf16_f32 which already
-// produces fp32 dq[8]. Saves the bf16→fp32 widening that the bf162_t
-// overload above does for dq (still need to widen A from bf16). Wins more
-// at high N: the bf162_t version's per-call widening cost scales with the
-// number of dequants × M_COUNT × 4 dot calls; the fp32 version pays only
-// for A widening (M_COUNT × 4 × 4 widens, half as many).
-__forceinline__ __device__ float dot22_8_f(float (&dq)[8],
-                                           const bf16_t* a_ptr) {
-  float result = 0.0f;
-  #pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t aw;
-    __builtin_memcpy(&aw, a_ptr + 2 * i, sizeof(uint32_t));
-    float a_x = __uint_as_float((aw & 0xFFFFu) << 16);
-    float a_y = __uint_as_float(aw & 0xFFFF0000u);
-    result = __fmaf_rn(dq[2 * i + 0], a_x, result);
-    result = __fmaf_rn(dq[2 * i + 1], a_y, result);
-  }
-  return result;
-}
-
-// ---------------------------------------------------------------------------
-// Packed atomic-add via CAS-loop on a 64-bit word (4 fp16/bf16 lanes per CAS).
-// RDNA3 (gfx11) does NOT have native v_global_atomic_pk_add_f16 / _bf16 (those
-// landed on gfx940 / gfx1250 respectively), so this lowers to
-// global_atomic_cmpswap_b64 plus retry. We use this in the kernel epilogue to
-// write 4 output columns per row in a single atomic operation — half the
-// atomic instruction count and half the contention vs two 32-bit CAS calls.
-//
-// Writing directly to fp16/bf16 (instead of through an FP32 scratch buffer +
-// cast pass) saves M*N*4 bytes of allocation, the memset, and the epilogue
-// cast pass that an fp32-accumulator design would need.
-//
-// 64-bit alignment: the kernel writes at `out + n` where n = offset_n + t*4
-// (always multiple of 4), and partition_weight_shape[1] is required to be a
-// multiple of 8 by can_implement(), so every (m, n) write target is 8-byte
-// aligned. Required by global_atomic_cmpswap_b64.
-// ---------------------------------------------------------------------------
-
-__forceinline__ __device__ void atomic_add_pk4_f16(half* addr, half2 v01,
-                                                   half2 v23) {
-  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
-  unsigned long long old = *addr_u;
-  while (true) {
-    union {
-      unsigned long long u;
-      half2 h2[2];
-    } cur, sum;
-    cur.u = old;
-    sum.h2[0] = __hadd2(cur.h2[0], v01);
-    sum.h2[1] = __hadd2(cur.h2[1], v23);
-    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
-    if (prev == old) break;
-    old = prev;
-  }
-}
-
-__forceinline__ __device__ void atomic_add_pk4_bf16(bf16_t* addr, bf162_t v01,
-                                                    bf162_t v23) {
-  unsigned long long* addr_u = reinterpret_cast<unsigned long long*>(addr);
-  unsigned long long old = *addr_u;
-  while (true) {
-    union {
-      unsigned long long u;
-      bf162_t b2[2];
-    } cur, sum;
-    cur.u = old;
-    sum.b2[0] = __hadd2(cur.b2[0], v01);
-    sum.b2[1] = __hadd2(cur.b2[1], v23);
-    unsigned long long prev = atomicCAS(addr_u, old, sum.u);
-    if (prev == old) break;
-    old = prev;
-  }
-}
-
-// Load one row's worth of 4 packed zeros (column n..n+3) from a [groups, N/8]
-// uint32 tensor. n is a multiple of 4 by construction (n = offset_n + t*4 with
-// offset_n = blockIdx.x * 512), so the 4 nibbles always live within one or two
-// uint32 words; in practice within one because n & 7 is 0 or 4.
-__forceinline__ __device__ void load4_zeros(const uint32_t* qzeros_row, int n,
-                                            int (&zeros)[4]) {
-  int qcol = n / 8;
-  int shift = (n & 0x07) * 4;
-  uint32_t d = qzeros_row[qcol] >> shift;
-  zeros[0] = (int)(d & 0xF);
-  zeros[1] = (int)((d >> 4) & 0xF);
-  zeros[2] = (int)((d >> 8) & 0xF);
-  zeros[3] = (int)((d >> 12) & 0xF);
-}
-
-template <typename T>
-__forceinline__ __device__ void load4_scales(const T* scales_row, int n,
-                                             T (&scales)[4]) {
-  scales[0] = scales_row[n + 0];
-  scales[1] = scales_row[n + 1];
-  scales[2] = scales_row[n + 2];
-  scales[3] = scales_row[n + 3];
-}
-
-// ---------------------------------------------------------------------------
-// Main kernel.
-// ---------------------------------------------------------------------------
-
-template <typename T, int M_COUNT>
-__global__ void gemm_q4_kernel_rdna3(
-    const T* __restrict__ a, const uint32_t* __restrict__ b_q_weight,
-    const uint32_t* __restrict__ b_qzeros, const T* __restrict__ b_scales,
-    T* __restrict__ c, const int size_m, const int size_n, const int size_k,
-    const int groups, const int zero_offset, const int* __restrict__ b_q_perm) {
-  const int t = threadIdx.x;
-  const int offset_n = blockIdx.x * BLOCK_KN_SIZE * 4;
-  const int offset_m = blockIdx.y * M_COUNT;
-  const int offset_k = blockIdx.z * BLOCK_KN_SIZE;
-  const int end_k = min(offset_k + BLOCK_KN_SIZE, size_k);
-  const int n = offset_n + t * 4;
-
-  // LDS layout: [M_COUNT][BLOCK_KN_SIZE + LDS_PAD]. The PAD=8 elements per M
-  // row break the natural 256-element/512-byte alignment that would otherwise
-  // collide on the same LDS bank when a thread reads block_a[0..M_COUNT-1][k]
-  // (same k, different m). Row stride becomes 264 elements * 2B = 528B = 132
-  // 4-byte banks, so m-stride hits banks (m*132)%32 = (m*4)%32 — distinct for
-  // all M_COUNT ≤ 8. Cost: 16B LDS per block, irrelevant.
-  constexpr int LDS_PAD = 8;
-  __shared__ T block_a[M_COUNT][BLOCK_KN_SIZE + LDS_PAD];
-
-  // Stage A: each thread loads 1 K element per M row into LDS (with optional
-  // act-order permutation). THREADS_X == BLOCK_KN_SIZE so this is a 1:1 map.
-  // For M_COUNT > 1 with size_m not a multiple of M_COUNT, slots past size_m
-  // are zero-padded so the dot product contribution is 0 (we then skip the
-  // atomic write for those rows below).
-  //
-  // M=1 fast path: skip LDS staging + __syncthreads entirely. All 256 threads
-  // read the SAME 8-element A window per inner step (a_off is uniform across
-  // the block), so the cache-line broadcast through L1 makes global reads as
-  // cheap as LDS reads. Measured: ~1% on 4B b=1, ~6% on 27B b=1 in=128.
-  static_assert(BLOCK_KN_SIZE == THREADS_X,
-                "BLOCK_KN_SIZE must equal THREADS_X (1 K element per thread)");
-  // The M=1 fast path (skip LDS) only has a global-read code path for bf16
-  // (the v_dot2_f32_bf16 branch).  The fp16 inner loop still indexes
-  // block_a[m][a_off] unconditionally, so for fp16 we MUST stage A through
-  // LDS even at M=1 to avoid reading uninitialized shared memory.
-  constexpr bool USE_LDS_A = (M_COUNT > 1) || std::is_same<T, half>::value;
-  if constexpr (USE_LDS_A) {
-    if (offset_k + t < end_k) {
-  #pragma unroll
-      for (int m = 0; m < M_COUNT; ++m) {
-        T av;
-        if (offset_m + m < size_m) {
-          const T* a_row = a + (offset_m + m) * size_k;
-          if (b_q_perm)
-            av = a_row[b_q_perm[offset_k + t]];
-          else
-            av = a_row[offset_k + t];
-        } else {
-          av = tzero<T>();  // zero-pad invalid M rows
-        }
-        block_a[m][t] = av;
-      }
-    }
-
-    // Threads beyond the right edge of N have nothing to do. Note: we must NOT
-    // return before __syncthreads() if any thread in the block participates in
-    // the LDS load above — but here all THREADS_X (=256) threads always do,
-    // regardless of whether their `n` is in bounds.
-    __syncthreads();
-  } else if (b_q_perm) {
-    // bf16 M=1 fast path skips LDS, but its global read below is sequential
-    // and cannot apply act-order. When a permutation is present, stage the
-    // single A row through LDS (as fp16 / M>1 do) so the read picks it up.
-    // b_q_perm is block-uniform, so the __syncthreads is non-divergent.
-    if (offset_k + t < end_k)
-      block_a[0][t] = a[offset_m * size_k + b_q_perm[offset_k + t]];
-    __syncthreads();
-  }
-  if (n >= size_n) return;
-
-  // Group bookkeeping. We require size_k % groups == 0 (groupsize divides K).
-  const int groupsize = size_k / groups;
-  int group = offset_k / groupsize;
-  int nextgroup = (group + 1) * groupsize;
-
-  // qweight stride: weights are [K/8, N] uint32 with K packed at dim 0.
-  int qk = offset_k / 8;
-  const uint32_t* b_ptr = b_q_weight + qk * size_n + n;
-
-  // Per-column dequant constants. We hold one set of (z, y) pairs per column.
-  // fp16 uses the exllama (z1z16, y1y16) double-pair to enable the upper-
-  // nibble-*16 trick. bf16 uses fp32 scalars (z, y) because the dequant
-  // produces fp32 directly — see prep_zero_scale_bf16_f32 / the FMA
-  // bypass for the missing v_pk_fma_bf16 on gfx11.
-  half2 z1z16_h[4][2], y1y16_h[4][2];
-  float z_b_f[4], y_b_f[4];
-
-  auto refresh_group = [&](int g) {
-    const uint32_t* qz_row = b_qzeros + g * (size_n / 8);
-    const T* sc_row = b_scales + g * size_n;
-    int zeros[4];
-    T scales[4];
-    load4_zeros(qz_row, n, zeros);
-    load4_scales<T>(sc_row, n, scales);
-    if constexpr (std::is_same<T, half>::value) {
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        prep_zero_scale_fp16((uint32_t)(zeros[i] + zero_offset), scales[i],
-                             z1z16_h[i], y1y16_h[i]);
-      }
-    } else {
-  #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        prep_zero_scale_bf16_f32((uint32_t)(zeros[i] + zero_offset), scales[i],
-                                 z_b_f[i], y_b_f[i]);
-      }
-    }
-  };
-
-  refresh_group(group);
-
-  float block_c[M_COUNT][4];
-  #pragma unroll
-  for (int m = 0; m < M_COUNT; ++m) {
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) block_c[m][j] = 0.0f;
-  }
-
-  // Note on group-transition granularity: we check `k == nextgroup` at the
-  // start of each outer iteration (which advances K by 32). This is correct
-  // when group_size >= 32 OR group_size divides 32 evenly (groupsize is one
-  // of {1,2,4,8,16,32,64,128,...}). For group_size in {16, 8, 4, ...} the
-  // inner loop would cross a group boundary between j-iterations; we require
-  // group_size >= 32 here, mirroring exllama's assumption.
-  //
-  // Software pipelining: we issue all 4 vectorized weight loads up front
-  // before any dequant/FMA depends on them. This gives the AMDGPU backend
-  // freedom to schedule the global_loads early and overlap their latency
-  // with dequant + v_pk_fma_f16 of earlier iterations. Cost: 4×int4 = 16
-  // VGPRs in flight per thread, plenty of headroom on RDNA3.
-  int k = offset_k;
-  while (k < end_k) {
-    if (k == nextgroup) {
-      group++;
-      nextgroup += groupsize;
-      refresh_group(group);
-    }
-
-    // Prefetch all four j-iterations' weight words. The compiler emits 4
-    // global_load_b128 instructions back-to-back; the dependent dequant +
-    // FMA work below hides their latency.
-    int4 b_w[4];
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) {
-      b_w[j] = *(const int4*)(b_ptr + j * size_n);
-    }
-    b_ptr += 4 * size_n;
-
-  #pragma unroll
-    for (int j = 0; j < 4; ++j) {
-      const int a_off = (k - offset_k) + 8 * j;
-
-      if constexpr (std::is_same<T, half>::value) {
-        half2 dq[4][4];
-        dequant_4bit_8_fp16((uint32_t)b_w[j].x, dq[0], z1z16_h[0], y1y16_h[0]);
-        dequant_4bit_8_fp16((uint32_t)b_w[j].y, dq[1], z1z16_h[1], y1y16_h[1]);
-        dequant_4bit_8_fp16((uint32_t)b_w[j].z, dq[2], z1z16_h[2], y1y16_h[2]);
-        dequant_4bit_8_fp16((uint32_t)b_w[j].w, dq[3], z1z16_h[3], y1y16_h[3]);
-
-  #pragma unroll
-        for (int m = 0; m < M_COUNT; ++m) {
-          const half* a_ptr = reinterpret_cast<const half*>(&block_a[m][a_off]);
-          block_c[m][0] += dot22_8_f(dq[0], a_ptr);
-          block_c[m][1] += dot22_8_f(dq[1], a_ptr);
-          block_c[m][2] += dot22_8_f(dq[2], a_ptr);
-          block_c[m][3] += dot22_8_f(dq[3], a_ptr);
-        }
-      } else if constexpr (M_COUNT == 1) {
-        // bf16 decode (M=1), v_dot2_f32_bf16 path. Mirrors the data-flow of
-        // Hybrid PR #40977's wvSplitK_int4 kernel exactly so clang's
-        // InstCombine cannot fold the bf16→fp32 widening (LLVM #76000):
-        //   * activations and magic-value weights share a fp32-aliased
-        //     union (bytes written as uint32, read as bf16x2_t for the
-        //     dot — pointer-cast opacity defeats the fold)
-        //   * sum_a computed via a *second* v_dot2 with bf162(1,1) as the
-        //     second operand, avoiding any explicit bf16→fp32 widen of A
-        //   * bias correction y_b_f * partial + z_b_f * sum_a, identical
-        //     to the previous fp32-FMA-chain path
-        //
-        // Net: 20 v_dot2_f32_bf16 + 8 fp32 FMA per int32 weight vs the
-        // previous 40 fp32 FMA. v_dot2 runs at full rate on gfx1100, so
-        // the substitution is ~2× cheaper for the inner accumulator.
-        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
-        constexpr uint32_t BF16_MAGIC = 0x43004300u;  // bf162(128, 128)
-        constexpr uint32_t BF16_ONES = 0x3F803F80u;   // bf162(1.0, 1.0)
-        union pack4 {
-          float f[4];
-          uint32_t u[4];
-        };
-
-        uint32_t w[4];
-        __builtin_memcpy(w, &b_w[j], sizeof(int4));
-
-        // Load 8 bf16 activations as 4 uint32s (= 4 bf16x2 pairs) into a
-        // fp32-aliased union. Storing as uint32 keeps the IR-level type
-        // opaque so the inner v_dot2 cannot be folded to fp32 widening.
-        //
-        // A is read direct from global (no LDS staging — see USE_LDS_A above),
-        // except under act-order, where it comes from the permuted LDS copy.
-        pack4 a_pack;
-        {
-          const uint32_t* a_words =
-              b_q_perm
-                  ? reinterpret_cast<const uint32_t*>(&block_a[0][a_off])
-                  : reinterpret_cast<const uint32_t*>(a + offset_k + a_off);
-          a_pack.u[0] = a_words[0];
-          a_pack.u[1] = a_words[1];
-          a_pack.u[2] = a_words[2];
-          a_pack.u[3] = a_words[3];
-        }
-
-        // sum_a = Σ a[i]. Computed via 4× v_dot2_f32_bf16 with bf162(1,1) as
-        // the second operand — every bf16 pair contributes 1·a_lo + 1·a_hi.
-        // No fp32 widening of activations: the bytes go straight from LDS
-        // through v_dot2 into the fp32 accumulator.
-        float sum_a = 0.0f;
-  #pragma unroll
-        for (int b = 0; b < 4; ++b) {
-          sum_a = __builtin_amdgcn_fdot2_f32_bf16(
-              *((bf16x2_t*)(&a_pack.f[b])), *((const bf16x2_t*)&BF16_ONES),
-              sum_a, /*clamp=*/false);
-        }
-
-        // unroll 1 keeps q_pack alive only one col at a time (8 fp32 VGPRs
-        // recycled across cols), avoiding straight-line expansion that
-        // would inflate live-range to 32 VGPRs.
-  #pragma unroll 1
-        for (int col = 0; col < 4; ++col) {
-          // Build dequant magic values bf16(128 + nibble) directly into a
-          // fp32-aliased union via uint32 stores. No fp32 in the data flow
-          // until v_dot2 consumes the bytes.
-          pack4 q_pack;
-          const uint32_t qa = w[col];
-          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
-
-          // partial = Σ (128 + nibble[i]) · a[i], via 4× v_dot2_f32_bf16.
-          float partial = 0.0f;
-  #pragma unroll
-          for (int b = 0; b < 4; ++b) {
-            partial = __builtin_amdgcn_fdot2_f32_bf16(
-                *((bf16x2_t*)(&a_pack.f[b])), *((bf16x2_t*)(&q_pack.f[b])),
-                partial, /*clamp=*/false);
-          }
-
-          // block_c += y_b_f * partial + z_b_f * sum_a
-          //   y_b_f = scale, z_b_f = -(128+zero)*scale
-          // partial holds (128 + nibble) · a; subtracting (128+zero)·sum_a
-          // and scaling yields scale · (nibble - zero) · a as required.
-          block_c[0][col] =
-              __fmaf_rn(y_b_f[col], partial,
-                        __fmaf_rn(z_b_f[col], sum_a, block_c[0][col]));
-        }
-      } else {
-        // bf16 M_COUNT > 1 path with v_dot2_f32_bf16. Same opacity trick as
-        // the M=1 branch: activations + magic-value weights stored in
-        // fp32-aliased unions, dot via __builtin_amdgcn_fdot2_f32_bf16 with
-        // pointer-cast to bf16x2_t. sum_a[m] computed via second v_dot2
-        // with BF16_ONES; bias correction (y_b_f * partial + z_b_f * sum_a)
-        // applied after the dot. Magic values built once per col and reused
-        // across all M rows — amortizes dequant cost across M_COUNT.
-        typedef short __attribute__((ext_vector_type(2))) bf16x2_t;
-        constexpr uint32_t BF16_MAGIC = 0x43004300u;  // bf162(128, 128)
-        constexpr uint32_t BF16_ONES = 0x3F803F80u;   // bf162(1.0, 1.0)
-        union pack4 {
-          float f[4];
-          uint32_t u[4];
-        };
-
-        uint32_t w[4];
-        __builtin_memcpy(w, &b_w[j], sizeof(int4));
-
-        // Load M_COUNT × 8 bf16 activations as 4 uint32s each into pack4
-        // unions. Stored as uint32 to keep IR-level types opaque (defeats
-        // InstCombine fold). At M_COUNT=8 this is 32 fp32 VGPRs — within RDNA3
-        // budget.
-        pack4 a_pack[M_COUNT];
-  #pragma unroll
-        for (int m = 0; m < M_COUNT; ++m) {
-          const uint32_t* a_words =
-              reinterpret_cast<const uint32_t*>(&block_a[m][a_off]);
-          a_pack[m].u[0] = a_words[0];
-          a_pack[m].u[1] = a_words[1];
-          a_pack[m].u[2] = a_words[2];
-          a_pack[m].u[3] = a_words[3];
-        }
-
-        // sum_a[m] = Σ a[m][i] via 4× v_dot2 with bf162(1,1) — no fp32 widen.
-        float sum_a[M_COUNT];
-  #pragma unroll
-        for (int m = 0; m < M_COUNT; ++m) {
-          float s = 0.0f;
-  #pragma unroll
-          for (int b = 0; b < 4; ++b) {
-            s = __builtin_amdgcn_fdot2_f32_bf16(*((bf16x2_t*)(&a_pack[m].f[b])),
-                                                *((const bf16x2_t*)&BF16_ONES),
-                                                s, /*clamp=*/false);
-          }
-          sum_a[m] = s;
-        }
-
-        // Per col: build magic-value pack, dot against all M activations.
-        // unroll 1 keeps q_pack live one col at a time (8 fp32 VGPRs recycled)
-        // — same register-pressure trick as the previous fp32 path.
-  #pragma unroll 1
-        for (int col = 0; col < 4; ++col) {
-          pack4 q_pack;
-          const uint32_t qa = w[col];
-          q_pack.u[0] = ((qa >> 0) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[1] = ((qa >> 4) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[2] = ((qa >> 8) & 0x000F000Fu) | BF16_MAGIC;
-          q_pack.u[3] = ((qa >> 12) & 0x000F000Fu) | BF16_MAGIC;
-
-  #pragma unroll
-          for (int m = 0; m < M_COUNT; ++m) {
-            float partial = 0.0f;
-  #pragma unroll
-            for (int b = 0; b < 4; ++b) {
-              partial = __builtin_amdgcn_fdot2_f32_bf16(
-                  *((bf16x2_t*)(&a_pack[m].f[b])), *((bf16x2_t*)(&q_pack.f[b])),
-                  partial, /*clamp=*/false);
-            }
-            // block_c += y_b_f * partial + z_b_f * sum_a (same correction as
-            // M=1)
-            block_c[m][col] =
-                __fmaf_rn(y_b_f[col], partial,
-                          __fmaf_rn(z_b_f[col], sum_a[m], block_c[m][col]));
-          }
-        }
-      }
-    }
-    k += 32;  // 4 weight words * 8 nibbles = 32 K elements
-  }
-
-  // Pack the 4 FP32 partial sums into 2 packed pairs and atomically add all
-  // four lanes in a single 64-bit CAS write directly to the T-typed output
-  // (caller pre-zeros it). On gfx11 the packed atomic is a CAS-loop, but with
-  // a single b64 op we halve the atomic instruction count vs two b32 CAS
-  // calls, AND save the FP32 buffer + memset + cast pass entirely.
-  #pragma unroll
-  for (int m = 0; m < M_COUNT; ++m) {
-    if (offset_m + m >= size_m) continue;  // skip padding rows past size_m
-    T* out = c + (offset_m + m) * size_n + n;
-    if constexpr (std::is_same<T, half>::value) {
-      half2 r01 = __halves2half2(__float2half_rn(block_c[m][0]),
-                                 __float2half_rn(block_c[m][1]));
-      half2 r23 = __halves2half2(__float2half_rn(block_c[m][2]),
-                                 __float2half_rn(block_c[m][3]));
-      atomic_add_pk4_f16(out, r01, r23);
-    } else {
-      bf162_t r01;
-      r01.x = __float2bfloat16(block_c[m][0]);
-      r01.y = __float2bfloat16(block_c[m][1]);
-      bf162_t r23;
-      r23.x = __float2bfloat16(block_c[m][2]);
-      r23.y = __float2bfloat16(block_c[m][3]);
-      atomic_add_pk4_bf16(out, r01, r23);
-    }
-  }
-}
-
-#else  // non-RDNA3 device pass: empty __global__ for symbol parity.
-
-template <typename T, int M_COUNT>
-__global__ void gemm_q4_kernel_rdna3(const T*, const uint32_t*, const uint32_t*,
-                                     const T*, T*, const int, const int,
-                                     const int, const int, const int,
-                                     const int*) {}
-
-#endif  // __HIP__RDNA3__ || !__HIP_DEVICE_COMPILE__
-
-// ---------------------------------------------------------------------------
-// Launcher.
-// ---------------------------------------------------------------------------
-
-template <typename T, int M_COUNT>
-void launch_gemm_q4_for_mcount(const T* a, const uint32_t* b_q_weight,
-                               const uint32_t* b_qzeros, const T* b_scales,
-                               const int* b_q_perm, T* c, int size_m,
-                               int size_n, int size_k, int groups,
-                               int zero_offset, cudaStream_t stream) {
-  dim3 block(THREADS_X);
-  dim3 grid((size_n + BLOCK_KN_SIZE * 4 - 1) / (BLOCK_KN_SIZE * 4),
-            (size_m + M_COUNT - 1) / M_COUNT,
-            (size_k + BLOCK_KN_SIZE - 1) / BLOCK_KN_SIZE);
-
-  gemm_q4_kernel_rdna3<T, M_COUNT><<<grid, block, 0, stream>>>(
-      a, b_q_weight, b_qzeros, b_scales, c, size_m, size_n, size_k, groups,
-      zero_offset, b_q_perm);
-}
-
-// Dispatch to the largest M_COUNT template that doesn't waste more than
-// half a tile. Caps at 8: above that, the WMMA-prefill kernel (M >= 16) is
-// the right tool, not bigger M_COUNT in the scalar dot-product path.
-//
-// Tile-waste table:
-//   M=1   -> M_COUNT=1   (no waste)
-//   M=2,3 -> M_COUNT=2   (M=3 wastes 1/2 of last tile)
-//   M=4-7 -> M_COUNT=4   (worst case M=5: wastes 3/4 of last tile)
-//   M=8-15-> M_COUNT=8   (worst case M=9: wastes 7/8 of last tile)
-// "Wasted" rows are zero-padded in LDS and skip the atomic write, so they
-// only burn instructions on the last block, never affect correctness.
-template <typename T>
-void launch_gemm_q4(const T* a, const uint32_t* b_q_weight,
-                    const uint32_t* b_qzeros, const T* b_scales,
-                    const int* b_q_perm, T* c, int size_m, int size_n,
-                    int size_k, int groups, bool use_v2_format,
-                    cudaStream_t stream) {
-  const int zero_offset = use_v2_format ? 0 : 1;
-
-  if (size_m == 1) {
-    launch_gemm_q4_for_mcount<T, 1>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
-                                    c, size_m, size_n, size_k, groups,
-                                    zero_offset, stream);
-  } else if (size_m <= 3) {
-    launch_gemm_q4_for_mcount<T, 2>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
-                                    c, size_m, size_n, size_k, groups,
-                                    zero_offset, stream);
-  } else if (size_m <= 7) {
-    launch_gemm_q4_for_mcount<T, 4>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
-                                    c, size_m, size_n, size_k, groups,
-                                    zero_offset, stream);
-  } else {
-    // M_COUNT=8 covers M up to 15 here; M >= 16 should ideally take the
-    // WMMA path, but if it falls through we still produce correct output —
-    // just leaving 3-5× of throughput on the table for prefill workloads.
-    launch_gemm_q4_for_mcount<T, 8>(a, b_q_weight, b_qzeros, b_scales, b_q_perm,
-                                    c, size_m, size_n, size_k, groups,
-                                    zero_offset, stream);
-  }
-}
-
-}  // namespace gptq_rdna3
-}  // namespace vllm
-
-// ---------------------------------------------------------------------------
-// Public entry point.
-// ---------------------------------------------------------------------------
-//
-// Inputs:
-//   a         [M, K]            half or bfloat16
-//   b_q_weight[K/8, N]          uint32 (already shuffled via gptq_shuffle)
-//   b_qzeros  [groups, N/8]     uint32 (packed 4-bit zeros)
-//   b_scales  [groups, N]       half or bfloat16
-//   b_g_idx   [K] or empty      int32 (act-order permutation; empty=identity)
-//   use_v2_format                bool   (true = GPTQv2, no +1 zero offset)
-//
-// Output:
-//   c         [M, N]            same dtype as a
-
-torch::Tensor gptq_gemm_rdna3_wmma(torch::Tensor a, torch::Tensor b_q_weight,
-                                   torch::Tensor b_qzeros,
-                                   torch::Tensor b_scales,
-                                   torch::Tensor b_g_idx, bool use_v2_format);
-
-torch::Tensor gptq_gemm_rdna3(torch::Tensor a, torch::Tensor b_q_weight,
-                              torch::Tensor b_qzeros, torch::Tensor b_scales,
-                              torch::Tensor b_g_idx, bool use_v2_format) {
-  if (a.dim() == 2 && b_q_weight.dim() == 2 && a.size(1) % 16 == 0 &&
-      b_q_weight.size(1) % 16 == 0 &&
-      ((a.scalar_type() == torch::kBFloat16 && a.size(0) >= 16) ||
-       (a.scalar_type() == torch::kHalf && a.size(0) >= 64))) {
-    return gptq_gemm_rdna3_wmma(a, b_q_weight, b_qzeros, b_scales, b_g_idx,
-                                use_v2_format);
-  }
-
-  TORCH_CHECK(a.is_cuda(), "a must be a CUDA/HIP tensor");
-  TORCH_CHECK(b_q_weight.is_cuda(), "b_q_weight must be a CUDA/HIP tensor");
-  TORCH_CHECK(b_qzeros.is_cuda(), "b_qzeros must be a CUDA/HIP tensor");
-  TORCH_CHECK(b_scales.is_cuda(), "b_scales must be a CUDA/HIP tensor");
-  TORCH_CHECK(a.dim() == 2, "a must be 2D [M, K]");
-  TORCH_CHECK(b_q_weight.dim() == 2, "b_q_weight must be 2D [K/8, N]");
-  TORCH_CHECK(
-      a.scalar_type() == torch::kHalf || a.scalar_type() == torch::kBFloat16,
-      "a must be half or bfloat16");
-  TORCH_CHECK(a.scalar_type() == b_scales.scalar_type(),
-              "b_scales dtype must match a");
-
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
-  auto stream = at::cuda::getCurrentCUDAStream();
-
-  int size_m = (int)a.size(0);
-  int size_k = (int)a.size(1);
-  int size_n = (int)b_q_weight.size(1);
-  int groups = (int)b_qzeros.size(0);
-
-  TORCH_CHECK(b_q_weight.size(0) * 8 == size_k,
-              "b_q_weight first dim must be K/8");
-  TORCH_CHECK(b_scales.size(0) == groups,
-              "b_scales must have same group count as qzeros");
-  TORCH_CHECK(b_scales.size(1) == size_n, "b_scales last dim must be N");
-  TORCH_CHECK(size_n % 8 == 0, "N must be a multiple of 8 (64-bit atomic CAS)");
-
-  auto opts = torch::TensorOptions().dtype(a.dtype()).device(a.device());
-  at::Tensor c = torch::zeros({size_m, size_n}, opts);
-
-  const int* g_idx_ptr = nullptr;
-  if (!b_g_idx.device().is_meta() && b_g_idx.numel() > 0) {
-    TORCH_CHECK(b_g_idx.scalar_type() == torch::kInt32,
-                "b_g_idx must be int32");
-    g_idx_ptr = (const int*)b_g_idx.data_ptr();
-  }
-
-  if (a.scalar_type() == torch::kHalf) {
-    vllm::gptq_rdna3::launch_gemm_q4<half>(
-        (const half*)a.data_ptr(), (const uint32_t*)b_q_weight.data_ptr(),
-        (const uint32_t*)b_qzeros.data_ptr(), (const half*)b_scales.data_ptr(),
-        g_idx_ptr, (half*)c.data_ptr(), size_m, size_n, size_k, groups,
-        use_v2_format, stream);
-  } else {
-    vllm::gptq_rdna3::launch_gemm_q4<vllm::gptq_rdna3::bf16_t>(
-        (const vllm::gptq_rdna3::bf16_t*)a.data_ptr(),
-        (const uint32_t*)b_q_weight.data_ptr(),
-        (const uint32_t*)b_qzeros.data_ptr(),
-        (const vllm::gptq_rdna3::bf16_t*)b_scales.data_ptr(), g_idx_ptr,
-        (vllm::gptq_rdna3::bf16_t*)c.data_ptr(), size_m, size_n, size_k, groups,
-        use_v2_format, stream);
-  }
-
-  return c;
-}

csrc/rocm/qdq_4_rdna3.cuh

@@ -1,239 +0,0 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-//
-// W4A16 dequant primitives for RDNA3 (gfx1100/gfx1101/gfx1102), templated on
-// the activation/scale dtype (half or __hip_bfloat16). The fp16 path reuses
-// the classic exllamav2 bit-trick:
-//
-//   (qa & 0x000F000F) | 0x64006400  ->  half2(1024+q_lo, 1024+q_hi)
-//   (qa & 0x00F000F0) | 0x64006400  ->  half2(1024+q_lo*16, 1024+q_hi*16)
-//
-// The "*16 then divide by 16 in the FMA" trick for the upper-nibble pairs
-// works in fp16 because the mantissa (10 bits) is wide enough to hold a value
-// shifted by 4 bits. In bf16 the mantissa is only 7 bits, so shifting an upper
-// nibble into bits [7:4] would spill into the exponent. To avoid that, the
-// bf16 path shifts each pair of nibbles down to bits [3:0]/[19:16] with a
-// single right-shift before the OR with 0x43004300 (= bf162(128, 128)).
-
-#ifndef _qdq_4_rdna3_cuh
-#define _qdq_4_rdna3_cuh
-
-#include <cstdint>
-
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp16.h>
-
-namespace vllm {
-namespace gptq_rdna3 {
-
-using bf16_t = __hip_bfloat16;
-using bf162_t = __hip_bfloat162;
-
-// Bit-shuffle for an int32 holding 8 sequential 4-bit weights q[0..7]:
-//   in:  q[7] q[6] q[5] q[4] q[3] q[2] q[1] q[0]   (LSB first)
-//   out: q[7] q[5] q[3] q[1] q[6] q[4] q[2] q[0]   (even/odd interleaved)
-//
-// After shuffle, q[2k]   sits at bits [4k   : 4k+3]   (lower 16)
-//                q[2k+1] sits at bits [16+4k: 16+4k+3] (upper 16)
-// so a single mask 0x000F000F selects the matching even/odd pair, ready to
-// bitcast to half2 / bfloat162 after OR-ing with the magic constant.
-__forceinline__ __device__ void shuffle_4bit_8(uint32_t* q) {
-  uint32_t qa = q[0];
-  uint32_t qb = 0;
-#pragma unroll
-  for (int i = 0; i < 4; i++) {
-    uint32_t qa0 = qa & 0x0F;
-    uint32_t qa1 = (qa & 0xF0) >> 4;
-    qa >>= 8;
-    qb |= (qa1 << (i * 4 + 16));
-    qb |= (qa0 << (i * 4));
-  }
-  q[0] = qb;
-}
-
-// ---------------------------------------------------------------------------
-// fp16 path
-// ---------------------------------------------------------------------------
-
-// Precompute scale-baked constants for a single zero/scale pair.
-//   z1z16[0] = scale * (-1024 - zero)            (used for "low" pairs)
-//   z1z16[1] = scale * (-64   - zero)            (used for "high" pairs)
-//   y1y16[0] = scale * 1                          (low pairs are q + 1024)
-//   y1y16[1] = scale * (1/16)                     (high pairs are q*16 + 1024)
-__forceinline__ __device__ void prep_zero_scale_fp16(uint32_t zero, half scale,
-                                                     half2 (&z1z16)[2],
-                                                     half2 (&y1y16)[2]) {
-  // half(-1024 - zero) via the exllamav2 bit-trick:
-  //   half bits 0xE400 == -1024.0 ; ORing the zero into mantissa subtracts it.
-  union {
-    uint16_t u;
-    half h;
-  } z1u;
-  z1u.u = (uint16_t)(0xE400 | zero);
-  half z1 = z1u.h;
-  half z16 = __hsub(__int2half_rn(-64), __int2half_rn((int)zero));
-
-  half2 scale2 = __half2half2(scale);
-  z1z16[0] = __hmul2(scale2, __half2half2(z1));
-  z1z16[1] = __hmul2(scale2, __half2half2(z16));
-
-  half y1 = __float2half_rn(1.0f);
-  half y16 = __float2half_rn(1.0f / 16.0f);
-  y1y16[0] = __hmul2(scale2, __half2half2(y1));
-  y1y16[1] = __hmul2(scale2, __half2half2(y16));
-}
-
-// Dequantize one int32 (8 shuffled 4-bit weights) into 4 half2 pairs:
-//   dq[0] = (q[0], q[1]) * scale - zero*scale
-//   dq[1] = (q[2], q[3]) * scale - zero*scale
-//   dq[2] = (q[4], q[5]) * scale - zero*scale
-//   dq[3] = (q[6], q[7]) * scale - zero*scale
-__forceinline__ __device__ void dequant_4bit_8_fp16(uint32_t qa, half2 (&dq)[4],
-                                                    half2 (&z1z16)[2],
-                                                    half2 (&y1y16)[2]) {
-  const uint32_t c0 = 0x64006400;
-
-  union {
-    uint32_t u;
-    half2 h2;
-  } q0, q1, q2, q3;
-  q0.u = (qa & 0x000F000F) | c0;  // half2(q[0]+1024, q[1]+1024)
-  q1.u = (qa & 0x00F000F0) | c0;  // half2(q[2]*16+1024, q[3]*16+1024)
-  uint32_t qa_hi = qa >> 8;
-  q2.u = (qa_hi & 0x000F000F) | c0;  // half2(q[4]+1024, q[5]+1024)
-  q3.u = (qa_hi & 0x00F000F0) | c0;  // half2(q[6]*16+1024, q[7]*16+1024)
-
-  dq[0] = __hfma2(q0.h2, y1y16[0], z1z16[0]);
-  dq[1] = __hfma2(q1.h2, y1y16[1], z1z16[1]);
-  dq[2] = __hfma2(q2.h2, y1y16[0], z1z16[0]);
-  dq[3] = __hfma2(q3.h2, y1y16[1], z1z16[1]);
-}
-
-// ---------------------------------------------------------------------------
-// bf16 path
-// ---------------------------------------------------------------------------
-
-// Bit-trick magic for bf16:
-//   bf16(128) == 0x4300 (sign 0, exp 134, mantissa 0).
-//   For nibble n in [0..15], bits [3:0] of mantissa hold n exactly because
-//   bf16's ULP at 128 is 1 (mantissa step = 2^(7-7) = 1). So
-//   ((qa & 0x000F000F) | 0x43004300) bitcasts to bfloat162(128+n_lo, 128+n_hi).
-//
-// Because bf16's mantissa is only 7 bits, we cannot use the fp16 "upper nibble
-// * 16" trick. Instead each pair of nibbles is shifted down to [3:0]/[19:16]
-// via a single 4/8/12-bit right-shift before the OR. That costs one extra
-// shift per pair vs fp16, but keeps the FMA structure identical.
-__forceinline__ __device__ void prep_zero_scale_bf16(uint32_t zero,
-                                                     bf16_t scale,
-                                                     bf162_t& z_prep,
-                                                     bf162_t& y_prep) {
-  // z = scale * -(128 + zero); y = scale.
-  float scale_f = __bfloat162float(scale);
-  float zf = -(128.0f + (float)zero) * scale_f;
-  bf16_t zb = __float2bfloat16(zf);
-  z_prep = __bfloat162bfloat162(zb);
-  y_prep = __bfloat162bfloat162(scale);
-}
-
-__forceinline__ __device__ void dequant_4bit_8_bf16(uint32_t qa,
-                                                    bf162_t (&dq)[4],
-                                                    bf162_t z_prep,
-                                                    bf162_t y_prep) {
-  const uint32_t c0 = 0x43004300;
-
-  union {
-    uint32_t u;
-    bf162_t b2;
-  } q0, q1, q2, q3;
-  q0.u = ((qa >> 0) & 0x000F000F) | c0;   // bf162(128+q[0], 128+q[1])
-  q1.u = ((qa >> 4) & 0x000F000F) | c0;   // bf162(128+q[2], 128+q[3])
-  q2.u = ((qa >> 8) & 0x000F000F) | c0;   // bf162(128+q[4], 128+q[5])
-  q3.u = ((qa >> 12) & 0x000F000F) | c0;  // bf162(128+q[6], 128+q[7])
-
-  // dq = q_b * scale + (-(128+zero)*scale) = (q - zero) * scale
-  dq[0] = __hfma2(q0.b2, y_prep, z_prep);
-  dq[1] = __hfma2(q1.b2, y_prep, z_prep);
-  dq[2] = __hfma2(q2.b2, y_prep, z_prep);
-  dq[3] = __hfma2(q3.b2, y_prep, z_prep);
-}
-
-// ---------------------------------------------------------------------------
-// bf16-input → fp32-output dequant (RDNA3 scalar path).
-//
-// RDNA3 (gfx1100) has no v_pk_fma_bf16; packed bf16 FMA lowers to a slow
-// fallback. Rather than computing dq in bf16 and widening at FMA time in
-// the dot product, we widen to fp32 here once (a free left-shift by 16) and
-// emit the (q - zero) * scale FMA directly in fp32. This:
-//   * Replaces 4× slow bf16 packed FMA with 8× fast fp32 FMA per int32.
-//   * Eliminates 4× bf16→fp32 widens that the dot product would do.
-//   * Keeps the dot product accumulator in fp32 without a roundtrip.
-//
-// Output: fp32 dq[8], one element per K position (consumed by the
-// fp32-overload of dot22_8_f in q_gemm_rdna3.cu).
-__forceinline__ __device__ void prep_zero_scale_bf16_f32(uint32_t zero,
-                                                         bf16_t scale,
-                                                         float& z_prep,
-                                                         float& y_prep) {
-  float scale_f = __bfloat162float(scale);
-  z_prep = -(128.0f + (float)zero) * scale_f;
-  y_prep = scale_f;
-}
-
-// Pure-q dequant for the M_COUNT=1 factored path: outputs the unscaled fp32
-// values 128+nibble, without folding scale/zero. The caller folds scale/zb
-// into the accumulator outside the inner loop using a precomputed sum_a,
-// which saves ~27% of the FMA count vs the per-col-dequant approach above
-// (only beneficial at M_COUNT=1; break-even at M_COUNT=2).
-//
-// Cost: 0 FMAs (pure bit-trick + as_float reinterprets).
-__forceinline__ __device__ void dequant_4bit_8_bf16_q_only(uint32_t qa,
-                                                           float (&q_f32)[8]) {
-  const uint32_t c0 = 0x43004300;
-  const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
-  const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
-  const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
-  const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
-  q_f32[0] = __uint_as_float((q0 & 0xFFFFu) << 16);
-  q_f32[1] = __uint_as_float(q0 & 0xFFFF0000u);
-  q_f32[2] = __uint_as_float((q1 & 0xFFFFu) << 16);
-  q_f32[3] = __uint_as_float(q1 & 0xFFFF0000u);
-  q_f32[4] = __uint_as_float((q2 & 0xFFFFu) << 16);
-  q_f32[5] = __uint_as_float(q2 & 0xFFFF0000u);
-  q_f32[6] = __uint_as_float((q3 & 0xFFFFu) << 16);
-  q_f32[7] = __uint_as_float(q3 & 0xFFFF0000u);
-}
-
-__forceinline__ __device__ void dequant_4bit_8_bf16_f32(uint32_t qa,
-                                                        float (&dq)[8],
-                                                        float z_prep,
-                                                        float y_prep) {
-  const uint32_t c0 = 0x43004300;
-  const uint32_t q0 = ((qa >> 0) & 0x000F000F) | c0;
-  const uint32_t q1 = ((qa >> 4) & 0x000F000F) | c0;
-  const uint32_t q2 = ((qa >> 8) & 0x000F000F) | c0;
-  const uint32_t q3 = ((qa >> 12) & 0x000F000F) | c0;
-  // bf16(128+nibble) bits → fp32(128+nibble) bits via left-shift by 16
-  // (just zero-extends the mantissa from 7 to 23 bits; exponent preserved).
-  const float q0x = __uint_as_float((q0 & 0xFFFFu) << 16);
-  const float q0y = __uint_as_float(q0 & 0xFFFF0000u);
-  const float q1x = __uint_as_float((q1 & 0xFFFFu) << 16);
-  const float q1y = __uint_as_float(q1 & 0xFFFF0000u);
-  const float q2x = __uint_as_float((q2 & 0xFFFFu) << 16);
-  const float q2y = __uint_as_float(q2 & 0xFFFF0000u);
-  const float q3x = __uint_as_float((q3 & 0xFFFFu) << 16);
-  const float q3y = __uint_as_float(q3 & 0xFFFF0000u);
-  // dq[i] = q_f32 * scale + (-(128+zero)*scale) = (nibble - zero) * scale
-  dq[0] = __fmaf_rn(q0x, y_prep, z_prep);
-  dq[1] = __fmaf_rn(q0y, y_prep, z_prep);
-  dq[2] = __fmaf_rn(q1x, y_prep, z_prep);
-  dq[3] = __fmaf_rn(q1y, y_prep, z_prep);
-  dq[4] = __fmaf_rn(q2x, y_prep, z_prep);
-  dq[5] = __fmaf_rn(q2y, y_prep, z_prep);
-  dq[6] = __fmaf_rn(q3x, y_prep, z_prep);
-  dq[7] = __fmaf_rn(q3y, y_prep, z_prep);
-}
-
-}  // namespace gptq_rdna3
-}  // namespace vllm
-
-#endif  // _qdq_4_rdna3_cuh

csrc/rocm/torch_bindings.cpp

@@ -39,19 +39,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
       "          Tensor scale_b, int CuCount) -> ()");
   rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
 
-#ifdef VLLM_ROCM_GFX1100
-  // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
-  rocm_ops.def(
-      "gptq_gemm_rdna3(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
-      "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
-  rocm_ops.impl("gptq_gemm_rdna3", torch::kCUDA, &gptq_gemm_rdna3);
-
-  rocm_ops.def(
-      "gptq_gemm_rdna3_wmma(Tensor a, Tensor b_q_weight, Tensor b_qzeros, "
-      "Tensor b_scales, Tensor b_g_idx, bool use_v2_format) -> Tensor");
-  rocm_ops.impl("gptq_gemm_rdna3_wmma", torch::kCUDA, &gptq_gemm_rdna3_wmma);
-#endif
-
   // Custom attention op
   // Compute the attention between an input query and the cached
   // keys/values using PagedAttention.

csrc/torch_bindings.cpp

@@ -1,7 +1,4 @@
-// Provides torch::Tensor for ops.h (previously included transitively via
-// cache.h, which is no longer included here after cache ops moved to
-// _C_stable_libtorch).
-#include <torch/all.h>
+#include "cache.h"
 #include "cuda_utils.h"
 #include "ops.h"
 #include "core/registration.h"
@@ -36,6 +33,35 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
            &get_cuda_view_from_cpu_tensor);
 
+  // Attention ops
+  // Compute the attention between an input query and the cached
+  // keys/values using PagedAttention.
+  ops.def(
+      "paged_attention_v1("
+      "    Tensor! out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
+  ops.impl("paged_attention_v1", torch::kCUDA, &paged_attention_v1);
+
+  // PagedAttention V2.
+  ops.def(
+      "paged_attention_v2("
+      "    Tensor! out, Tensor! exp_sums, Tensor! max_logits,"
+      "    Tensor! tmp_out, Tensor query, Tensor key_cache,"
+      "    Tensor value_cache, int num_kv_heads, float scale,"
+      "    Tensor block_tables, Tensor seq_lens, int block_size,"
+      "    int max_seq_len, Tensor? alibi_slopes,"
+      "    str kv_cache_dtype, Tensor k_scale, Tensor v_scale,"
+      "    int tp_rank, int blocksparse_local_blocks,"
+      "    int blocksparse_vert_stride, int blocksparse_block_size,"
+      "    int blocksparse_head_sliding_step) -> ()");
+  ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2);
+
   // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
   ops.def(
       "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
@@ -187,31 +213,118 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "float eps) -> (Tensor, Tensor)");
   ops.impl("minimax_allreduce_rms_qk", torch::kCUDA, &minimax_allreduce_rms_qk);
 
-  ops.def(
-      "minimax_m3_build_k2q_csr_with_schedule("
-      "Tensor q2k,"
-      "Tensor cu_q,"
-      "Tensor cu_k,"
-      "Tensor! row_ptr,"
-      "Tensor! q_idx,"
-      "Tensor! scheduler_metadata,"
-      "Tensor! work_count,"
-      "Tensor! qsplit_idx,"
-      "Tensor! split_counts,"
-      "int topk,"
-      "int blk_kv,"
-      "int total_rows,"
-      "int max_kv_blocks,"
-      "int target_q_per_cta,"
-      "int work_capacity,"
-      "int max_seqlen_q) -> ()");
-  ops.impl("minimax_m3_build_k2q_csr_with_schedule", torch::kCUDA,
-           &run_minimax_m3_build_k2q_csr_with_schedule);
-
   //  conditionally compiled so impl in source file
 #endif
 }
 
+TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
+  // Cache ops
+  // Swap in (out) the cache blocks from src to dst.
+  cache_ops.def(
+      "swap_blocks(Tensor src, Tensor! dst,"
+      "            int block_size_in_bytes, Tensor block_mapping) -> ()");
+  cache_ops.impl("swap_blocks", torch::kCUDA, &swap_blocks);
+
+  // Batch swap: submit all block copies in a single driver call.
+  cache_ops.def(
+      "swap_blocks_batch(Tensor src_ptrs, Tensor dst_ptrs,"
+      "                  Tensor sizes,"
+      "                  bool is_src_access_order_any=False) -> ()");
+  cache_ops.impl("swap_blocks_batch", torch::kCPU, &swap_blocks_batch);
+
+  // Reshape the key and value tensors and cache them.
+  cache_ops.def(
+      "reshape_and_cache(Tensor key, Tensor value,"
+      "                  Tensor! key_cache, Tensor! value_cache,"
+      "                  Tensor slot_mapping,"
+      "                  str kv_cache_dtype,"
+      "                  Tensor k_scale, Tensor v_scale) -> ()");
+  cache_ops.impl("reshape_and_cache", torch::kCUDA, &reshape_and_cache);
+
+  // Reshape the key and value tensors and cache them.
+  cache_ops.def(
+      "reshape_and_cache_flash(Tensor key, Tensor value,"
+      "                        Tensor! key_cache,"
+      "                        Tensor! value_cache,"
+      "                        Tensor slot_mapping,"
+      "                        str kv_cache_dtype,"
+      "                        Tensor k_scale, Tensor v_scale) -> ()");
+  cache_ops.impl("reshape_and_cache_flash", torch::kCUDA,
+                 &reshape_and_cache_flash);
+
+  // Concat kv_c and k_pe and cache them.
+  cache_ops.def(
+      "concat_and_cache_mla(Tensor kv_c, Tensor k_pe,"
+      "                     Tensor! kv_cache,"
+      "                     Tensor slot_mapping,"
+      "                     str kv_cache_dtype,"
+      "                     Tensor scale) -> ()");
+  cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
+
+  // Rotate Q and K, then write to kv cache for MLA
+  cache_ops.def(
+      "concat_and_cache_mla_rope_fused("
+      "                     Tensor positions,"
+      "                     Tensor! q_pe,"
+      "                     Tensor! k_pe,"
+      "                     Tensor kv_c,"
+      "                     Tensor cos_sin_cache,"
+      "                     bool is_neox,"
+      "                     Tensor slot_mapping,"
+      "                     Tensor! kv_cache,"
+      "                     str kv_cache_dtype,"
+      "                     Tensor kv_cache_scale) -> ()");
+  cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
+                 &concat_and_cache_mla_rope_fused);
+
+  // Convert the key and value cache to fp8 data type.
+  cache_ops.def(
+      "convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
+      "str kv_cache_dtype) -> ()");
+  cache_ops.impl("convert_fp8", torch::kCUDA, &convert_fp8);
+
+  // Gather cache blocks from src_cache to dst, dequantizing from
+  // src_cache's dtype to dst's dtype if necessary.
+  cache_ops.def(
+      "gather_and_maybe_dequant_cache(Tensor src_cache, Tensor! dst, "
+      "                               Tensor block_table, Tensor cu_seq_lens, "
+      "                               Tensor token_to_seq, "
+      "                               int num_tokens, "
+      "                               str kv_cache_dtype, "
+      "                               Tensor scale, Tensor? seq_starts) -> ()");
+  cache_ops.impl("gather_and_maybe_dequant_cache", torch::kCUDA,
+                 &gather_and_maybe_dequant_cache);
+
+  cache_ops.def(
+      "cp_gather_cache(Tensor src_cache, Tensor! dst, Tensor block_table, "
+      "Tensor cu_seq_lens, int batch_size, Tensor? seq_starts) -> ()");
+  cache_ops.impl("cp_gather_cache", torch::kCUDA, &cp_gather_cache);
+
+  cache_ops.def(
+      "cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
+      "Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
+      "batch_size) -> ()");
+  cache_ops.impl("cp_gather_and_upconvert_fp8_kv_cache", torch::kCUDA,
+                 &cp_gather_and_upconvert_fp8_kv_cache);
+
+  cache_ops.def(
+      "indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
+      "slot_mapping, "
+      "int quant_block_size, str kv_cache_dtype) -> ()");
+  cache_ops.impl("indexer_k_quant_and_cache", torch::kCUDA,
+                 &indexer_k_quant_and_cache);
+
+  cache_ops.def(
+      "concat_mla_q(Tensor ql_nope, Tensor q_pe, Tensor! q_out) -> ()");
+  cache_ops.impl("concat_mla_q", torch::kCUDA, &concat_mla_q);
+
+  cache_ops.def(
+      "cp_gather_indexer_k_quant_cache(Tensor kv_cache, Tensor! dst_k, Tensor! "
+      "dst_scale, Tensor block_table, Tensor cu_seq_lens) -> ()");
+  cache_ops.impl("cp_gather_indexer_k_quant_cache", torch::kCUDA,
+                 &cp_gather_indexer_k_quant_cache);
+}
+
 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) {
   // Cuda utils
 

docker/Dockerfile

@@ -757,7 +757,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.12
+ARG FLASHINFER_VERSION=0.6.11.post2
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

docker/Dockerfile.cpu

@@ -165,23 +165,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
     VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
 
-######################### TRITON-CPU BUILD IMAGE #########################
-FROM base AS vllm-triton-cpu-build
-
-WORKDIR /vllm-workspace
-
-RUN mkdir dist
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=cache,target=/root/.cache/ccache \
-    --mount=type=cache,target=/vllm-workspace/.deps,sharing=locked \
-    if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
-        git clone --recurse-submodules "https://github.com/triton-lang/triton-cpu.git"; \
-        cd triton-cpu; \
-        git checkout "270e696d"; \
-        uv build --wheel --out-dir=../dist; \
-    fi
-
 ######################### TEST DEPS #########################
 FROM base AS vllm-test-deps
 
@@ -274,14 +257,7 @@ WORKDIR /vllm-workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=bind,from=vllm-build,src=/vllm-workspace/dist,target=dist \
-    uv pip install "$(realpath dist/*.whl)[audio]"
-
-RUN --mount=type=cache,target=/root/.cache/uv \
-    --mount=type=cache,target=/root/.cache/ccache \
-    --mount=type=bind,from=vllm-triton-cpu-build,src=/vllm-workspace/dist,target=dist \
-    if [ "$TARGETARCH" = "amd64" ] || [ "$VLLM_CPU_X86" != "0" ]; then \
-        uv pip install "$(realpath dist/*.whl)"; \
-    fi
+    uv pip install "$(realpath dist/*.whl)[audio,triton-cpu]"
 
 # Add labels to document build configuration
 LABEL org.opencontainers.image.title="vLLM CPU"

docker/Dockerfile.nightly_torch

@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.12
+# release version: v0.6.11.post2
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/uv \
     echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.11.post2 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.12"
+      "default": "0.6.11.post2"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

docs/configuration/optimization.md

@@ -46,14 +46,14 @@ In V1, **chunked prefill is enabled by default whenever possible**. With chunked
 
 This policy has two benefits:
 
-- It improves inter-token latency (ITL) and generation decode because decode requests are prioritized.
+- It improves ITL and generation decode because decode requests are prioritized.
 - It helps achieve better GPU utilization by locating compute-bound (prefill) and memory-bound (decode) requests to the same batch.
 
 ### Performance Tuning with Chunked Prefill
 
 You can tune the performance by adjusting `max_num_batched_tokens`:
 
-- Smaller values (e.g., 2048) achieve better ITL because there are fewer prefills slowing down decodes.
+- Smaller values (e.g., 2048) achieve better inter-token latency (ITL) because there are fewer prefills slowing down decodes.
 - Higher values achieve better time to first token (TTFT) as you can process more prefill tokens in a batch.
 - For optimal throughput, we recommend setting `max_num_batched_tokens > 8192` especially for smaller models on large GPUs.
 - If `max_num_batched_tokens` is the same as `max_model_len`, that's almost the equivalent to the V0 default scheduling policy (except that it still prioritizes decodes).

docs/design/attention_backends.md

@@ -170,14 +170,14 @@ 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 | 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 | 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 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
-| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
+| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
 | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
@@ -187,18 +187,6 @@ Priority is **1 = highest** (tried first).
 >
 > **\*** Specify the FlashAttention version via `--attention-config.flash_attn_version=2`, `3`, or `4`. Default is FA4 on SM100+ (Blackwell), FA3 on SM90 (Hopper), FA2 otherwise.
 
-## MiniMax M3 Sparse Attention Backends
-
-Block-sparse GQA backend used by MiniMax M3 sparse ("lightning indexer")
-layers. It is wired in directly by the model and is not part of the
-automatic priority lists above. A lightning indexer scores KV blocks, the
-top-k blocks (plus fixed init/local blocks) are selected, and attention
-attends only to those blocks; index keys live in a separate side cache.
-
-| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
-| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
-| `MINIMAX_M3_SPARSE` | bf16, fp16 | `bfloat16` | 128 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
-
 ## MLA (Multi-head Latent Attention) Backends
 
 MLA uses separate backends for prefill and decode phases.

docs/getting_started/installation/gpu.xpu.inc.md

@@ -41,7 +41,7 @@ pip install -v -r requirements/xpu.txt
 
     ```bash
     pip uninstall -y triton triton-xpu
-    pip install triton-xpu==3.7.0 --extra-index-url https://download.pytorch.org/whl/xpu
+    pip install triton-xpu==3.6.0 --extra-index-url https://download.pytorch.org/whl/xpu
     ```
 
     !!! note

docs/governance/committers.md

@@ -17,7 +17,6 @@ Sorted alphabetically by GitHub handle:
 - [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
 - [@benchislett](https://github.com/benchislett): Engine core and spec decode
 - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
-- [@BugenZhao](https://github.com/BugenZhao): Rust frontend
 - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
 - [@DarkLight1337](https://github.com/DarkLight1337): Multimodality, API server
 - [@esmeetu](https://github.com/esmeetu): developer marketing, community
@@ -131,8 +130,6 @@ If you have PRs touching the area, please feel free to ping the area owner for r
     - @DarkLight1337
 - API Server: The OpenAI-compatible API server
     - @DarkLight1337, @njhill, @aarnphm, @simon-mo, @heheda12345 (Responses API)
-- Rust Frontend: The experimental API server in Rust
-    - @BugenZhao, @njhill
 - Batch Runner: The OpenAI-compatible batch runner
     - @simon-mo
 

docs/models/supported_models.md

@@ -633,7 +633,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | |
 | `Step3VLForConditionalGeneration` | Step3-VL | T + I<sup>+</sup> | `stepfun-ai/step3` | | ✅︎ |
 | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I<sup>+</sup> | `stepfun-ai/Step3-VL-10B` | | ✅︎ |
-| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I<sup>+</sup> | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |
 | `TarsierForConditionalGeneration` | Tarsier | T + I<sup>E+</sup> | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |
 | `Tarsier2ForConditionalGeneration`<sup>^</sup> | Tarsier2 | T + I<sup>E+</sup> + V<sup>E+</sup> | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ |
 | `UltravoxModel` | Ultravox | T + A<sup>E+</sup> | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |

docs/pre_run_check.sh

@@ -1,60 +1,41 @@
-if [ "$READTHEDOCS_VERSION_TYPE" != "external" ]; then
-  echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
-  exit 0
-fi
-
-echo "Checking for changes to docs-affecting files vs origin/main..."
-DOCS_PATHS=(
-  docs/                       # Actual docs content
-  examples/                   # Examples are rendered in docs
-  vllm/                       # API & CLI reference
-  requirements/test/cuda.txt  # CLI reference (see docs/mkdocs/hooks/generate_argparse.py)
-  mkdocs.yaml                 # Affects build process
-  .readthedocs.yaml           # Affects build process
-  requirements/docs.txt       # Affects build process
-  requirements/docs.in        # Affects build process
-)
-if git diff --quiet origin/main -- "${DOCS_PATHS[@]}"; then
-  echo "No docs-affecting files changed vs origin/main; cancelling build."
-  # See https://docs.readthedocs.com/platform/latest/guides/build/skip-build.html for info on exit code
-  exit 183
-fi
-echo "Docs-affecting files changed; continuing pre-run-check."
-echo "Checking pre-commit/pre-run-check status..."
-MAX_WAIT=300
-INTERVAL=60
-ELAPSED=0
-while :; do
-  RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
-  HTTP_CODE=$(printf %s "$RAW" | tail -n1)
-  BODY=$(printf %s "$RAW" | sed '$d')
-  if [ "$HTTP_CODE" != "200" ]; then
-    echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-commit/pre-run-check gate."
-    break
-  fi
-  STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
-  CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
-  CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")")
-  if [ "$STATUS" = "none" ]; then
-    echo "no pre-commit/pre-run-check found for this commit; skipping gate."
-    break
-  fi
-  if [ -n "$CONCLUSION" ]; then
-    echo "pre-commit/pre-run-check conclusion: $CONCLUSION"
-    if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
-      echo "pre-commit/pre-run-check did not pass; skipping docs build."
-      if [ -n "$CHECK_URL" ]; then
-        echo "pre-commit/pre-run-check failure reason: $CHECK_URL"
-      fi
-      exit 1
+if [ "$READTHEDOCS_VERSION_TYPE" = "external" ]; then
+  MAX_WAIT=300
+  INTERVAL=60
+  ELAPSED=0
+  while :; do
+    RAW=$(curl -sS -w "\n%{http_code}" "https://api.github.com/repos/vllm-project/vllm/commits/${READTHEDOCS_GIT_COMMIT_HASH}/check-runs?check_name=pre-run-check&filter=latest")
+    HTTP_CODE=$(printf %s "$RAW" | tail -n1)
+    BODY=$(printf %s "$RAW" | sed '$d')
+    if [ "$HTTP_CODE" != "200" ]; then
+      echo "GitHub API returned HTTP $HTTP_CODE (likely rate-limited); skipping pre-run-check gate."
+      break
     fi
-    break
-  fi
-  if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
-    echo "pre-commit/pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
-    break
-  fi
-  echo "pre-commit/pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
-  sleep "$INTERVAL"
-  ELAPSED=$((ELAPSED + INTERVAL))
-done
+    STATUS=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"status\") or \"\") if r else \"none\")")
+    CONCLUSION=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"conclusion\") or \"\") if r else \"\")")
+    CHECK_URL=$(printf %s "$BODY" | python3 -c "import sys, json; r=json.load(sys.stdin).get(\"check_runs\",[]); print((r[0].get(\"html_url\") or \"\") if r else \"\")")
+    if [ "$STATUS" = "none" ]; then
+      echo "no pre-run-check found for this commit; skipping gate."
+      break
+    fi
+    if [ -n "$CONCLUSION" ]; then
+      echo "pre-run-check conclusion: $CONCLUSION"
+      if [ "$CONCLUSION" = "failure" ] || [ "$CONCLUSION" = "cancelled" ] || [ "$CONCLUSION" = "timed_out" ]; then
+        echo "pre-run-check did not pass; skipping docs build."
+        if [ -n "$CHECK_URL" ]; then
+          echo "pre-run-check failure reason: $CHECK_URL"
+        fi
+        exit 1
+      fi
+      break
+    fi
+    if [ "$ELAPSED" -ge "$MAX_WAIT" ]; then
+      echo "pre-run-check status=$STATUS after ${MAX_WAIT}s; skipping gate."
+      break
+    fi
+    echo "pre-run-check status=$STATUS; waiting ${INTERVAL}s..."
+    sleep "$INTERVAL"
+    ELAPSED=$((ELAPSED + INTERVAL))
+  done
+else
+  echo "Not a PR build (version type=$READTHEDOCS_VERSION_TYPE); skipping pre-run-check gate."
+fi

docs/serving/expert_parallel_deployment.md

@@ -151,7 +151,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
 | `step_interval` | Frequency of rebalancing (every N engine steps) | 3000 |
 | `log_balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` |
 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
-| `use_async` | Use non-blocking EPLB for reduced latency overhead | `true` |
+| `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
 | `policy` | The policy type for expert parallel load balancing | `"default"` |
 | `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`,  or `null` (auto) | `null` |
 

examples/tool_chat_template_gemma4.jinja

@@ -295,15 +295,6 @@
                                 {%- endif -%}
                             {%- endfor -%}
                             {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
-                            {%- for part in tool_body -%}
-                                {%- if part.get('type') == 'image' -%}
-                                    {{- '<|image|>' -}}
-                                {%- elif part.get('type') == 'audio' -%}
-                                    {{- '<|audio|>' -}}
-                                {%- elif part.get('type') == 'video' -%}
-                                    {{- '<|video|>' -}}
-                                {%- endif -%}
-                            {%- endfor -%}
                         {%- else -%}
                             {{- format_tool_response_block(ns_tname.name, tool_body) -}}
                         {%- endif -%}

pyproject.toml

@@ -162,12 +162,7 @@ dout = "dout"
 Pn = "Pn"
 arange = "arange"
 thw = "thw"
-# temporal position ids (parallels hpos/wpos in vision RoPE)
-tpos = "tpos"
 subtile = "subtile"
-subtiles = "subtiles"
-reord = "reord"
-Ot = "Ot"
 HSA = "HSA"
 setp = "setp"
 CPY = "CPY"

requirements/common.txt

@@ -29,7 +29,6 @@ xgrammar >= 0.2.0, < 1.0.0; platform_machine == "x86_64" or platform_machine ==
 typing_extensions >= 4.10
 filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
 partial-json-parser # used for parsing partial JSON outputs
-jsonschema >= 4.23.0 # required for MiniMax M3 tool schema validation
 pyzmq >= 25.0.0
 msgspec
 gguf >= 0.17.0

requirements/cuda.txt

@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.12
-flashinfer-cubin==0.6.12
+flashinfer-python==0.6.11.post2
+flashinfer-cubin==0.6.11.post2
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to

requirements/test/cuda.txt

@@ -360,7 +360,6 @@ jsonpointer==3.0.0
     # via jsonschema
 jsonschema==4.23.0
     # via
-    #   -c requirements/common.txt
     #   hypothesis-jsonschema
     #   mistral-common
     #   ray

requirements/test/rocm.txt

@@ -440,8 +440,6 @@ jsonpointer==3.1.0
     # via jsonschema
 jsonschema==4.26.0
     # via
-    #   -c requirements/common.txt
-    #   -r requirements/test/../common.txt
     #   hypothesis-jsonschema
     #   mcp
     #   mistral-common

requirements/test/xpu.txt

@@ -229,7 +229,6 @@ jsonlines==4.0.0
     # via lm-eval
 jsonschema==4.26.0
     # via
-    #   -c requirements/common.txt
     #   hypothesis-jsonschema
     #   mistral-common
     #   schemathesis

requirements/xpu.txt

@@ -17,4 +17,4 @@ torchaudio
 torchvision
 
 auto_round_lib>=0.13.0
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.9/vllm_xpu_kernels-0.1.9-cp38-abi3-manylinux_2_28_x86_64.whl
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.8/vllm_xpu_kernels-0.1.8-cp38-abi3-manylinux_2_28_x86_64.whl

rust/Cargo.lock

@@ -3458,75 +3458,6 @@ version = "0.1.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e0c5ccf5294c6ccd63a74f1565028353830a9c2f5eb0c682c355c471726a6e3f"
 
-[[package]]
-name = "pyo3"
-version = "0.28.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "91fd8e38a3b50ed1167fb981cd6fd60147e091784c427b8f7183a7ee32c31c12"
-dependencies = [
- "libc",
- "once_cell",
- "portable-atomic",
- "pyo3-build-config",
- "pyo3-ffi",
- "pyo3-macros",
-]
-
-[[package]]
-name = "pyo3-build-config"
-version = "0.28.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e368e7ddfdeb98c9bca7f8383be1648fd84ab466bf2bc015e94008db6d35611e"
-dependencies = [
- "target-lexicon",
-]
-
-[[package]]
-name = "pyo3-ffi"
-version = "0.28.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7f29e10af80b1f7ccaf7f69eace800a03ecd13e883acfacc1e5d0988605f651e"
-dependencies = [
- "libc",
- "pyo3-build-config",
-]
-
-[[package]]
-name = "pyo3-macros"
-version = "0.28.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "df6e520eff47c45997d2fc7dd8214b25dd1310918bbb2642156ef66a67f29813"
-dependencies = [
- "proc-macro2",
- "pyo3-macros-backend",
- "quote",
- "syn 2.0.117",
-]
-
-[[package]]
-name = "pyo3-macros-backend"
-version = "0.28.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c4cdc218d835738f81c2338f822078af45b4afdf8b2e33cbb5916f108b813acb"
-dependencies = [
- "heck",
- "proc-macro2",
- "pyo3-build-config",
- "quote",
- "syn 2.0.117",
-]
-
-[[package]]
-name = "pythonize"
-version = "0.28.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0b79f670c9626c8b651c0581011b57b6ba6970bb69faf01a7c4c0cfc81c43f95"
-dependencies = [
- "pyo3",
- "serde",
- "serde_json",
-]
-
 [[package]]
 name = "qoi"
 version = "0.4.1"
@@ -4738,12 +4669,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "target-lexicon"
-version = "0.13.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
-
 [[package]]
 name = "task-local"
 version = "0.1.1"
@@ -5697,7 +5622,6 @@ dependencies = [
  "expect-test",
  "futures",
  "half",
- "indexmap 2.13.0",
  "itertools 0.14.0",
  "llm-multimodal",
  "minijinja",
@@ -5976,17 +5900,6 @@ dependencies = [
  "winnow",
 ]
 
-[[package]]
-name = "vllm-tool-parser-py"
-version = "0.1.0"
-dependencies = [
- "pyo3",
- "pythonize",
- "serde_json",
- "thiserror-ext",
- "vllm-tool-parser",
-]
-
 [[package]]
 name = "walkdir"
 version = "2.5.0"

rust/Cargo.toml

@@ -12,7 +12,6 @@ members = [
     "src/text",
     "src/tokenizer",
     "src/tool-parser",
-    "src/tool-parser/python",
 ]
 resolver = "3"
 
@@ -44,7 +43,6 @@ half = { version = "2.7.1", features = ["bytemuck"] }
 hex = "0.4.3"
 hf-hub = { version = "0.5.0", features = ["tokio"] }
 http-body = "1.0.1"
-indexmap = "2.13.0"
 itertools = "0.14.0"
 libc = "0.2.177"
 llm-multimodal = { git = "https://github.com/vllm-project/llm-multimodal", rev = "5b558989844d1c7af3e43d0f604069ffd9c06320" }
@@ -61,8 +59,6 @@ prometheus-client = "0.24.0"
 prometheus-client-derive-encode = "0.5.0"
 prost = "0.14.3"
 prost-types = "0.14.3"
-pyo3 = "0.28.3"
-pythonize = "0.28.0"
 rand = "0.9.2"
 reasoning-parser = "1.2.2"
 reqwest = { version = "0.12.8", default-features = false, features = ["rustls-tls"] }
@@ -73,7 +69,7 @@ rustc-hash = "1.1.0"
 serde = { version = "1.0.228", features = ["derive"] }
 serde-json-fmt = "0.1.0"
 serde_default = "0.2.0"
-serde_json = { version = "1.0.145", features = ["preserve_order"] }
+serde_json = { version = "1.0.145", features = ["arbitrary_precision", "preserve_order"] }
 serde_repr = "0.1.20"
 serde_tuple = "1.1.3"
 serde_with = "3.18.0"

rust/src/chat/Cargo.toml

@@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true
 easy-ext.workspace = true
 futures.workspace = true
 half.workspace = true
-indexmap.workspace = true
 itertools.workspace = true
 llm-multimodal.workspace = true
 minijinja.workspace = true

rust/src/chat/src/lib.rs

@@ -93,7 +93,7 @@ pub struct ChatLlm {
     text: TextLlm,
     backend: DynChatBackend,
     /// Effective model dtype reported by the engine.
-    model_dtype: ModelDtype,
+    model_dtype: Option<ModelDtype>,
     /// Tool-call parser selection.
     tool_call_parser: ParserSelection,
     /// Reasoning parser selection.
@@ -135,7 +135,7 @@ impl ChatLlm {
     }
 
     /// Override the effective model dtype used for multimodal tensor encoding.
-    pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
+    pub fn with_model_dtype(mut self, model_dtype: Option<ModelDtype>) -> Self {
         self.model_dtype = model_dtype;
         self
     }
@@ -233,7 +233,7 @@ mod tests {
         )
         .unwrap_err();
 
-        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, minimax_m3, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
+        expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, gemma4, glm45, glm47, hermes, hy_v3, kimi_k2, llama3_json, llama4_json, minimax_m2, mistral, qwen3_coder, qwen3_xml)"].assert_eq(&error.to_report_string());
     }
 
     #[test]

rust/src/chat/src/multimodal.rs

@@ -12,7 +12,7 @@
 use std::collections::{HashMap, HashSet};
 use std::fs;
 use std::path::Path;
-use std::sync::{Arc, LazyLock};
+use std::sync::{Arc, LazyLock, Once};
 
 use itertools::izip;
 use llm_multimodal::{
@@ -239,7 +239,7 @@ pub(crate) async fn finalize_rendered_prompt(
     request: &ChatRequest,
     rendered: RenderedPrompt,
     info: Option<&MultimodalModelInfo>,
-    model_dtype: ModelDtype,
+    model_dtype: Option<ModelDtype>,
 ) -> Result<(Prompt, Option<MmFeatures>)> {
     if !request.has_multimodal() {
         return Ok((rendered.prompt, None));
@@ -249,6 +249,16 @@ pub(crate) async fn finalize_rendered_prompt(
         bail_multimodal!("multimodal chat renderer must return a text prompt before expansion");
     };
     let media_parts = extract_media_parts(request)?;
+    let model_dtype = model_dtype.unwrap_or_else(|| {
+        static WARN_ONCE: Once = Once::new();
+        WARN_ONCE.call_once(|| {
+            warn!(
+                "engine handshake did not report model dtype; \
+                 falling back to float32 for multimodal tensor encoding"
+            );
+        });
+        ModelDtype::Float32
+    });
 
     let mut prompt_token_ids = info
         .context

rust/src/chat/src/parser/reasoning/mod.rs

@@ -5,9 +5,8 @@ use std::sync::LazyLock;
 pub use vllm_reasoning_parser::{
     CohereCmdReasoningParser, DeepSeekR1ReasoningParser, DeepSeekV3ReasoningParser,
     DeepSeekV4ReasoningParser, Gemma4ReasoningParser, Glm45ReasoningParser, KimiK2ReasoningParser,
-    KimiReasoningParser, MiniMaxM2ReasoningParser, MiniMaxM3ReasoningParser,
-    NemotronV3ReasoningParser, Qwen3ReasoningParser, ReasoningDelta, ReasoningError,
-    ReasoningParser, Step3ReasoningParser,
+    KimiReasoningParser, MiniMaxM2ReasoningParser, NemotronV3ReasoningParser, Qwen3ReasoningParser,
+    ReasoningDelta, ReasoningError, ReasoningParser, Step3ReasoningParser,
 };
 use vllm_tokenizer::DynTokenizer;
 
@@ -24,7 +23,6 @@ pub mod names {
     pub const KIMI: &str = "kimi";
     pub const KIMI_K2: &str = "kimi_k2";
     pub const MINIMAX_M2: &str = "minimax_m2";
-    pub const MINIMAX_M3: &str = "minimax_m3";
     pub const NEMOTRON_V3: &str = "nemotron_v3";
     pub const QWEN3: &str = "qwen3";
     pub const STEP3: &str = "step3";
@@ -61,7 +59,6 @@ impl ReasoningParserFactory {
             .register_parser::<KimiReasoningParser>(names::KIMI)
             .register_parser::<KimiK2ReasoningParser>(names::KIMI_K2)
             .register_parser::<MiniMaxM2ReasoningParser>(names::MINIMAX_M2)
-            .register_parser::<MiniMaxM3ReasoningParser>(names::MINIMAX_M3)
             .register_parser::<NemotronV3ReasoningParser>(names::NEMOTRON_V3)
             .register_parser::<Qwen3ReasoningParser>(names::QWEN3)
             .register_parser::<Step3ReasoningParser>(names::STEP3);
@@ -81,8 +78,6 @@ impl ReasoningParserFactory {
             .register_pattern("kimi-k2", names::KIMI_K2)
             .register_pattern("kimi", names::KIMI)
             .register_pattern("step3", names::STEP3)
-            .register_pattern("minimax-m3", names::MINIMAX_M3)
-            .register_pattern("mm-m3", names::MINIMAX_M3)
             .register_pattern("minimax", names::MINIMAX_M2)
             .register_pattern("mm-m2", names::MINIMAX_M2)
             .register_pattern("cohere", names::COHERE_CMD)

rust/src/chat/src/parser/reasoning/tests.rs

@@ -32,10 +32,8 @@ fn factory_contains_and_lists_registered_parsers() {
     let factory = ReasoningParserFactory::new();
     assert!(factory.contains(names::QWEN3));
     assert!(factory.contains(names::DEEPSEEK_V4));
-    assert!(factory.contains(names::MINIMAX_M3));
     assert!(factory.list().contains(&names::QWEN3.to_string()));
     assert!(factory.list().contains(&names::DEEPSEEK_V4.to_string()));
-    assert!(factory.list().contains(&names::MINIMAX_M3.to_string()));
 }
 
 #[test]
@@ -51,19 +49,6 @@ fn factory_resolves_deepseek_v4_to_qwen3_alias() {
     );
 }
 
-#[test]
-fn factory_resolves_minimax_m3_before_generic_minimax() {
-    let factory = ReasoningParserFactory::new();
-    assert_eq!(
-        factory.resolve_name_for_model("MiniMaxAI/Minimax-M3-preview"),
-        Some(names::MINIMAX_M3)
-    );
-    assert_eq!(
-        factory.resolve_name_for_model("mm-m3"),
-        Some(names::MINIMAX_M3)
-    );
-}
-
 #[test]
 fn factory_rejects_unknown_parser_names() {
     let tokenizer = Arc::new(FakeTokenizer);

rust/src/chat/src/parser/tool/mod.rs

@@ -5,9 +5,9 @@ use std::sync::LazyLock;
 pub use vllm_tool_parser::{
     DeepSeekV3ToolParser, DeepSeekV4ToolParser, DeepSeekV31ToolParser, DeepSeekV32ToolParser,
     Gemma4ToolParser, Glm45MoeToolParser, Glm47MoeToolParser, HermesToolParser, HyV3ToolParser,
-    KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MinimaxM3ToolParser,
-    MistralToolParser, Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser,
-    ToolParserError, ToolParserOutput,
+    KimiK2ToolParser, Llama3JsonToolParser, MinimaxM2ToolParser, MistralToolParser,
+    Qwen3CoderToolParser, Qwen3XmlToolParser, ToolCallDelta, ToolParser, ToolParserError,
+    ToolParserOutput,
 };
 
 use crate::parser::ParserFactory;
@@ -28,7 +28,6 @@ pub mod names {
     pub const LLAMA3_JSON: &str = "llama3_json";
     pub const LLAMA4_JSON: &str = "llama4_json";
     pub const MINIMAX_M2: &str = "minimax_m2";
-    pub const MINIMAX_M3: &str = "minimax_m3";
     pub const MISTRAL: &str = "mistral";
     pub const QWEN3_CODER: &str = "qwen3_coder";
     pub const QWEN3_XML: &str = "qwen3_xml";
@@ -67,7 +66,6 @@ impl ToolParserFactory {
             .register_parser::<Llama3JsonToolParser>(names::LLAMA3_JSON)
             .register_parser::<Llama3JsonToolParser>(names::LLAMA4_JSON)
             .register_parser::<MinimaxM2ToolParser>(names::MINIMAX_M2)
-            .register_parser::<MinimaxM3ToolParser>(names::MINIMAX_M3)
             .register_parser::<MistralToolParser>(names::MISTRAL)
             .register_parser::<Qwen3XmlToolParser>(names::QWEN3_XML)
             .register_parser::<Qwen3CoderToolParser>(names::QWEN3_CODER);
@@ -98,8 +96,6 @@ impl ToolParserFactory {
             .register_pattern("gemma4", names::GEMMA4)
             .register_pattern("gemma-4", names::GEMMA4)
             .register_pattern("kimi-k2", names::KIMI_K2)
-            .register_pattern("minimax-m3", names::MINIMAX_M3)
-            .register_pattern("mm-m3", names::MINIMAX_M3)
             .register_pattern("minimax", names::MINIMAX_M2)
             .register_pattern("mm-m2", names::MINIMAX_M2);
 

rust/src/chat/src/parser/tool/tests.rs

@@ -153,14 +153,6 @@ fn factory_new_resolves_default_patterns() {
         factory.resolve_name_for_model("tencent/Hy3-preview"),
         Some(names::HY_V3)
     );
-    assert_eq!(
-        factory.resolve_name_for_model("MiniMax/MiniMax-M3-Text"),
-        Some(names::MINIMAX_M3)
-    );
-    assert_eq!(
-        factory.resolve_name_for_model("org/mm-m3-base"),
-        Some(names::MINIMAX_M3)
-    );
     assert_eq!(
         factory.resolve_name_for_model("MiniMax/MiniMax-M2-01"),
         Some(names::MINIMAX_M2)

rust/src/chat/src/renderer/hf/mod.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 
 use serde::Serialize;
-use serde_json::Value as JsonValue;
+use serde_json::Value;
 use thiserror_ext::AsReport as _;
 use tracing::{info, trace, warn};
 use vllm_text::Prompt;
@@ -13,7 +13,6 @@ use self::format::{
     ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption,
 };
 use self::template::{CompiledChatTemplate, TemplateContext};
-use self::value::{TemplateValue, to_template_value};
 use super::{ChatRenderer, RenderedPrompt};
 use crate::error::Result;
 use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
@@ -25,7 +24,6 @@ mod error;
 mod format;
 mod template;
 mod tojson;
-mod value;
 
 pub use template::{load_chat_template, resolve_chat_template};
 
@@ -40,7 +38,7 @@ pub struct MultimodalRenderInfo {
 /// state.
 pub struct HfChatRenderer {
     default_template: Option<CompiledChatTemplate>,
-    default_template_kwargs: HashMap<String, JsonValue>,
+    default_template_kwargs: HashMap<String, Value>,
     content_format: ContentFormatOption,
     special_tokens: Option<HfSpecialTokens>,
     multimodal: Option<MultimodalRenderInfo>,
@@ -50,7 +48,7 @@ impl HfChatRenderer {
     /// Create a renderer from the given template string.
     pub fn new(
         template: Option<String>,
-        default_template_kwargs: HashMap<String, JsonValue>,
+        default_template_kwargs: HashMap<String, Value>,
         content_format: ContentFormatOption,
     ) -> Result<Self> {
         Ok(Self {
@@ -247,7 +245,7 @@ struct TemplateToolCall {
 #[derive(Debug, Serialize)]
 struct TemplateToolFunction {
     name: String,
-    arguments: TemplateValue,
+    arguments: Value,
 }
 
 #[derive(Debug, Serialize)]
@@ -261,7 +259,7 @@ pub(super) struct TemplateTool {
 struct TemplateToolDefinition {
     name: String,
     description: Option<String>,
-    parameters: TemplateValue,
+    parameters: Value,
     strict: Option<bool>,
 }
 
@@ -347,14 +345,13 @@ fn to_template_tool_calls(
     let mut tool_calls = Vec::new();
 
     for tool_call in content.tool_calls() {
-        let arguments = serde_json::from_str(&tool_call.arguments).map_err(|error| {
+        let arguments = serde_json::from_str::<Value>(&tool_call.arguments).map_err(|error| {
             Error::ChatTemplate(format!(
                 "assistant tool call `{}` has invalid JSON arguments: {}",
                 tool_call.id,
                 error.as_report()
             ))
         })?;
-        let arguments = to_template_value(arguments);
 
         tool_calls.push(TemplateToolCall {
             id: tool_call.id.clone(),
@@ -437,7 +434,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec<TemplateTool> {
             function: TemplateToolDefinition {
                 name: tool.name.clone(),
                 description: tool.description.clone(),
-                parameters: to_template_value(tool.parameters.clone()),
+                parameters: tool.parameters.clone(),
                 strict: tool.strict,
             },
         })
@@ -912,29 +909,6 @@ mod tests {
         assert_eq!(rendered, "get_weather|Paris|call_1|Sunny");
     }
 
-    #[test]
-    fn chat_template_tool_call_argument_items_method_is_not_shadowed_by_field() {
-        let request = sample_request(vec![ChatMessage::assistant_blocks(vec![
-            AssistantContentBlock::ToolCall(crate::AssistantToolCall {
-                id: "call_1".to_string(),
-                name: "add".to_string(),
-                arguments: r#"{"items":"operands","x":2,"y":1.0}"#.to_string(),
-            }),
-        ])]);
-
-        let rendered = render(
-            Some(
-                "{%- set arguments = messages[0].tool_calls[0].function.arguments -%}
-{%- for key, value in arguments.items() -%}{{ key }}={{ value }};{%- endfor -%}
-|{{ arguments['items'] }}",
-            ),
-            &request,
-        )
-        .unwrap();
-
-        assert_eq!(rendered, "items=operands;x=2;y=1.0;|operands");
-    }
-
     #[test]
     fn qwen35_template_renders_prefilled_reasoning_start_when_thinking_enabled() {
         let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);

rust/src/chat/src/renderer/hf/tojson.rs

@@ -208,27 +208,11 @@ mod tests {
     }
 
     #[test]
-    fn tojson_uses_standard_serde_json_number_spelling() {
+    fn tojson_preserves_arbitrary_precision_number_spelling() {
         let payload = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
         let rendered = render("{{ payload|tojson }}", payload);
 
-        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
-        // `arbitrary_precision` feature, otherwise the following test
-        // `serialized_json_numbers_do_not_leak_serde_private_representation` will fail.
-        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
-        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.0}");
-    }
-
-    #[test]
-    fn serialized_json_numbers_do_not_leak_serde_private_representation() {
-        let payload: serde_json::Value = serde_json::from_str(r#"{"x":2,"y":1.00}"#).unwrap();
-        let rendered = render("{{ payload }}", payload);
-
-        // TODO: we cannot preserve the original number precision by enabling `serde_json`'s
-        // `arbitrary_precision` feature, otherwise this will fail.
-        // See issue: https://github.com/mitsuhiko/minijinja/issues/641
-        assert!(!rendered.contains("$serde_json::private::Number"));
-        assert_eq!(rendered, r#"{"x": 2, "y": 1.0}"#);
+        assert_eq!(rendered, "{\"x\": 2, \"y\": 1.00}");
     }
 
     #[test]

rust/src/chat/src/renderer/hf/value.rs

@@ -1,77 +0,0 @@
-use std::sync::Arc;
-
-use indexmap::IndexMap;
-use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr};
-use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State};
-use serde::Serialize;
-use serde_json::Value as JsonValue;
-
-/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used
-/// as a value in the chat template.
-#[derive(Debug, Serialize)]
-#[serde(transparent)]
-pub(super) struct TemplateValue(minijinja::Value);
-
-pub(super) fn to_template_value(value: JsonValue) -> TemplateValue {
-    TemplateValue(match value {
-        JsonValue::Array(values) => values
-            .into_iter()
-            .map(to_template_value)
-            .map(|value| value.0)
-            .collect::<minijinja::Value>(),
-        JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap(
-            values
-                .into_iter()
-                .map(|(key, value)| (key, to_template_value(value).0))
-                .collect(),
-        )),
-        // For primitive values, directly convert them to `minijinja::Value` using `from_serialize`.
-        value => minijinja::Value::from_serialize(value),
-    })
-}
-
-/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can
-/// always handle dict methods through the unknown-method callback.
-///
-/// Use `IndexMap` to preserve the original key order when iterating.
-///
-/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates
-/// commonly call `dict.items()`, which would fail if the map had an `items` field.
-/// See issue: https://github.com/mitsuhiko/minijinja/issues/903
-#[derive(Debug)]
-struct TemplateMap(IndexMap<String, minijinja::Value>);
-
-impl Object for TemplateMap {
-    fn repr(self: &Arc<Self>) -> ObjectRepr {
-        ObjectRepr::Map
-    }
-
-    fn get_value(self: &Arc<Self>, key: &minijinja::Value) -> Option<minijinja::Value> {
-        self.0.get(key.as_str()?).cloned()
-    }
-
-    fn get_value_by_str(self: &Arc<Self>, key: &str) -> Option<minijinja::Value> {
-        self.0.get(key).cloned()
-    }
-
-    fn enumerate(self: &Arc<Self>) -> Enumerator {
-        self.mapped_rev_enumerator(|this| {
-            Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str())))
-        })
-    }
-
-    fn enumerator_len(self: &Arc<Self>) -> Option<usize> {
-        Some(self.0.len())
-    }
-
-    fn call_method(
-        self: &Arc<Self>,
-        _state: &State<'_, '_>,
-        _method: &str,
-        _args: &[minijinja::Value],
-    ) -> std::result::Result<minijinja::Value, TemplateError> {
-        // Always return `UnknownMethod` for method calls,
-        // so that pycompat can handle dict methods through the unknown-method callback.
-        Err(TemplateError::from(TemplateErrorKind::UnknownMethod))
-    }
-}

rust/src/chat/tests/roundtrip.rs

@@ -183,7 +183,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
         "roundtrip-reasoning-tools",
         vec![ChatMessage::text(
             ChatRole::User,
-            "Check Shanghai weather and add 1.0 plus 2.",
+            "Check Shanghai weather and add 1.00 plus 2.",
         )],
         test_tools(),
     );
@@ -210,10 +210,9 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
                 AssistantContentBlock::ToolCall(AssistantToolCall {
                     id: "functions.add:1".to_string(),
                     name: "add".to_string(),
-                    // Intentionally use a non-lexical order of keys to verify text-level
-                    // fidelity of the roundtrip where JSON formatting remains stable. The
-                    // `items` key also exercises templates that call `arguments.items()`.
-                    arguments: r#"{"y":1.0,"x":2,"items":["left","right"]}"#.to_string(),
+                    // Intentionally use a non-lexical order of keys and a different number
+                    // formatting style to verify text-level fidelity of the roundtrip.
+                    arguments: r#"{"y":1.00,"x":2}"#.to_string(),
                 }),
             ],
         },
@@ -241,7 +240,7 @@ async fn run_roundtrip_tool_call_mix(case: RoundtripCase) -> Result<()> {
     assert_eq!(tool_calls[1].name, "add");
     assert_eq!(
         tool_calls[1].arguments,
-        expected_arguments(&case, r#"{"y": 1.0, "x": 2, "items": ["left", "right"]}"#)?,
+        expected_arguments(&case, r#"{"y": 1.00, "x": 2}"#)?,
     );
 
     assert_eq!(
@@ -532,13 +531,9 @@ fn test_tools() -> Vec<ChatTool> {
                 "type": "object",
                 "properties": {
                     "y": { "type": "number" },
-                    "x": { "type": "number" },
-                    "items": {
-                        "type": "array",
-                        "items": { "type": "string" }
-                    }
+                    "x": { "type": "number" }
                 },
-                "required": ["y", "x", "items"]
+                "required": ["y", "x"]
             }),
             strict: None,
         },

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

@@ -290,8 +290,10 @@ impl EngineCoreClient {
 
         // If any engine reported a dp_stats_address in its ready response, use it
         // as the external coordinator address.
-        let dp_stats_address: Option<String> =
-            engines.iter().find_map(|engine| engine.ready_response.dp_stats_address.clone());
+        let dp_stats_address: Option<String> = engines
+            .iter()
+            .filter_map(|e| e.ready_response.as_ref())
+            .find_map(|r| r.dp_stats_address.clone());
 
         let (coordinator, coordinator_output_task, coordinator_task) =
             if let Some(coordinator_transport) = connected.coordinator {
@@ -366,44 +368,40 @@ impl EngineCoreClient {
     /// Return the ready responses received from all engines on the input
     /// socket.
     pub fn ready_responses(&self) -> Vec<&EngineCoreReadyResponse> {
-        self.engines.iter().map(|engine| &engine.ready_response).collect()
+        self.engines
+            .iter()
+            .filter_map(|engine| engine.ready_response.as_ref())
+            .collect()
     }
 
-    /// Return the engine-reported effective model dtype.
-    pub fn model_dtype(&self) -> ModelDtype {
+    /// Return the engine-reported effective model dtype, when available.
+    pub fn model_dtype(&self) -> Option<ModelDtype> {
         self.engines
-            .first()
-            .expect("engine core client requires at least one engine")
-            .ready_response
-            .dtype
-    }
-
-    /// Return the engine-reported Python vLLM version.
-    pub fn vllm_version(&self) -> &str {
-        self.engines
-            .first()
-            .expect("engine core client requires at least one engine")
-            .ready_response
-            .vllm_version
-            .as_str()
+            .iter()
+            .filter_map(|engine| engine.ready_response.as_ref())
+            .find_map(|response| response.dtype)
     }
 
     /// Return the total number of GPU blocks summed across all connected
     /// engines.
     pub fn total_num_gpu_blocks(&self) -> u64 {
-        self.engines.iter().map(|engine| engine.ready_response.num_gpu_blocks).sum()
+        self.engines
+            .iter()
+            .filter_map(|engine| engine.ready_response.as_ref())
+            .map(|r| r.num_gpu_blocks)
+            .sum()
     }
 
     /// Return the minimum engine-reported `max_model_len` across all engines.
     ///
     /// This is the auto-fitted value after KV cache profiling and may differ
     /// from the originally configured value.
-    pub fn max_model_len(&self) -> u32 {
+    pub fn max_model_len(&self) -> Option<u32> {
         self.engines
             .iter()
-            .map(|engine| engine.ready_response.max_model_len as u32)
+            .filter_map(|e| e.ready_response.as_ref())
+            .map(|r| r.max_model_len as u32)
             .min()
-            .expect("engine core client requires at least one engine")
     }
 
     /// Get the model name associated with this client used for metrics

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

@@ -382,7 +382,6 @@ mod tests {
     use zeromq::{RouterSocket, Socket};
 
     use super::*;
-    use crate::mock_engine::default_ready_response;
 
     async fn test_inner() -> ClientInner {
         let mut socket = RouterSocket::new();
@@ -393,7 +392,7 @@ mod tests {
             "test-model".to_string(),
             &[ConnectedEngine {
                 engine_id: EngineId::from(b"engine-0"),
-                ready_response: default_ready_response(),
+                ready_response: None,
             }],
         )
     }

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

@@ -339,20 +339,15 @@ mod tests {
     use super::{EngineRoutingState, RequestRegistry, UtilityRegistry};
     use crate::EngineId;
     use crate::client::state::EngineLoadSnapshot;
-    use crate::mock_engine::default_ready_response;
     use crate::protocol::{EngineCoreFinishReason, EngineCoreOutput};
     use crate::transport::ConnectedEngine;
 
-    fn connected_engine(engine_id: EngineId) -> ConnectedEngine {
-        ConnectedEngine {
-            engine_id,
-            ready_response: default_ready_response(),
-        }
-    }
-
     #[test]
     fn registry_rejects_duplicate_request_ids() {
-        let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
+        let mut registry = RequestRegistry::new(&[ConnectedEngine {
+            engine_id: EngineId::from(b"engine-0"),
+            ready_response: None,
+        }]);
         registry.register("req-1".to_string(), None).unwrap();
         let error = registry.register("req-1".to_string(), None).unwrap_err();
         assert!(matches!(
@@ -363,7 +358,10 @@ mod tests {
 
     #[test]
     fn registry_removes_finished_request_on_output() {
-        let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
+        let mut registry = RequestRegistry::new(&[ConnectedEngine {
+            engine_id: EngineId::from(b"engine-0"),
+            ready_response: None,
+        }]);
         registry.register("req-1".to_string(), None).unwrap();
 
         let sender = registry.sender_for_output(&EngineCoreOutput {
@@ -378,7 +376,10 @@ mod tests {
 
     #[test]
     fn registry_closes_all_requests_on_failure() {
-        let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
+        let mut registry = RequestRegistry::new(&[ConnectedEngine {
+            engine_id: EngineId::from(b"engine-0"),
+            ready_response: None,
+        }]);
         registry.register("req-1".to_string(), None).unwrap();
         registry.register("req-2".to_string(), None).unwrap();
 
@@ -393,8 +394,14 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            connected_engine(engine_0.clone()),
-            connected_engine(engine_1.clone()),
+            ConnectedEngine {
+                engine_id: engine_0.clone(),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: engine_1.clone(),
+                ready_response: None,
+            },
         ]);
         let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap();
         let (chosen_1, _) = registry.register("req-2".to_string(), None).unwrap();
@@ -421,8 +428,14 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            connected_engine(engine_0.clone()),
-            connected_engine(engine_1.clone()),
+            ConnectedEngine {
+                engine_id: engine_0.clone(),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: engine_1.clone(),
+                ready_response: None,
+            },
         ]);
 
         let (chosen_0, _) = registry.register("req-1".to_string(), None).unwrap();
@@ -475,8 +488,14 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            connected_engine(engine_0.clone()),
-            connected_engine(engine_1.clone()),
+            ConnectedEngine {
+                engine_id: engine_0.clone(),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: engine_1.clone(),
+                ready_response: None,
+            },
         ]);
 
         assert!(registry.apply_scheduler_counts(
@@ -504,9 +523,18 @@ mod tests {
         let engine_1 = EngineId::from_engine_index(1);
         let engine_2 = EngineId::from_engine_index(2);
         let mut registry = RequestRegistry::new(&[
-            connected_engine(engine_0.clone()),
-            connected_engine(engine_1.clone()),
-            connected_engine(engine_2.clone()),
+            ConnectedEngine {
+                engine_id: engine_0.clone(),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: engine_1.clone(),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: engine_2.clone(),
+                ready_response: None,
+            },
         ]);
 
         // Explicitly target rank 2 (third engine).
@@ -527,8 +555,14 @@ mod tests {
         let engine_0 = EngineId::from_engine_index(0);
         let engine_1 = EngineId::from_engine_index(1);
         let mut registry = RequestRegistry::new(&[
-            connected_engine(engine_0.clone()),
-            connected_engine(engine_1.clone()),
+            ConnectedEngine {
+                engine_id: engine_0.clone(),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: engine_1.clone(),
+                ready_response: None,
+            },
         ]);
 
         // Load-balance: first two go to engine_0 and engine_1.
@@ -543,8 +577,14 @@ mod tests {
     #[test]
     fn register_with_out_of_range_rank_returns_error() {
         let mut registry = RequestRegistry::new(&[
-            connected_engine(EngineId::from_engine_index(0)),
-            connected_engine(EngineId::from_engine_index(1)),
+            ConnectedEngine {
+                engine_id: EngineId::from_engine_index(0),
+                ready_response: None,
+            },
+            ConnectedEngine {
+                engine_id: EngineId::from_engine_index(1),
+                ready_response: None,
+            },
         ]);
 
         let error = registry.register("req-1".to_string(), Some(2)).unwrap_err();
@@ -560,7 +600,10 @@ mod tests {
     #[test]
     fn register_with_rank_on_single_engine_only_accepts_zero() {
         let engine_0 = EngineId::from_engine_index(0);
-        let mut registry = RequestRegistry::new(&[connected_engine(engine_0.clone())]);
+        let mut registry = RequestRegistry::new(&[ConnectedEngine {
+            engine_id: engine_0.clone(),
+            ready_response: None,
+        }]);
 
         let (chosen, _) = registry.register("req-ok".to_string(), Some(0)).unwrap();
         assert_eq!(chosen, engine_0);

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

@@ -47,8 +47,7 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
         max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
         num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
         dp_stats_address: None,
-        dtype: ModelDtype::Float32,
-        vllm_version: "test-vllm-version".to_string(),
+        dtype: Some(ModelDtype::Float32),
     }
 }
 

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

@@ -39,9 +39,10 @@ pub struct EngineCoreReadyResponse {
     /// DP coordinator stats publish address, if applicable.
     pub dp_stats_address: Option<String>,
     /// Effective model dtype after Python vLLM resolves `--dtype`.
-    pub dtype: ModelDtype,
-    /// Python vLLM version reported by the engine process.
-    pub vllm_version: String,
+    // TODO: This is currently not wired up on the engine side. After it's added, remove `Option`
+    // and `serde(default)`.
+    #[serde(default)]
+    pub dtype: Option<ModelDtype>,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup
@@ -68,3 +69,22 @@ pub struct HandshakeInitMessage {
     pub addresses: HandshakeAddresses,
     pub parallel_config: BTreeMap<String, OpaqueValue>,
 }
+
+#[cfg(test)]
+mod tests {
+    use super::EngineCoreReadyResponse;
+    use crate::protocol::ModelDtype;
+
+    #[test]
+    fn ready_response_accepts_effective_dtype() {
+        let response: EngineCoreReadyResponse = serde_json::from_value(serde_json::json!({
+            "max_model_len": 4096,
+            "num_gpu_blocks": 2,
+            "dp_stats_address": null,
+            "dtype": "bfloat16"
+        }))
+        .unwrap();
+
+        assert_eq!(response.dtype, Some(ModelDtype::BFloat16));
+    }
+}

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

@@ -10,9 +10,9 @@ use crate::EngineId;
 pub use crate::mock_engine::{MockCoordinatorSockets, MockEngineSockets};
 use crate::mock_engine::{
     MockEngineConfig, MockEngineDataSockets, connect_to_bootstrapped_frontend, connect_to_frontend,
-    default_ready_response,
 };
-use crate::protocol::handshake::HandshakeInitMessage;
+use crate::protocol::ModelDtype;
+use crate::protocol::handshake::{EngineCoreReadyResponse, HandshakeInitMessage};
 
 /// Per-test IPC endpoint namespace backed by a unique temporary directory.
 ///
@@ -57,7 +57,12 @@ fn test_mock_engine_config() -> MockEngineConfig {
     MockEngineConfig {
         local: true,
         headless: true,
-        ready_response: default_ready_response(),
+        ready_response: EngineCoreReadyResponse {
+            max_model_len: 4096,
+            num_gpu_blocks: 0,
+            dp_stats_address: None,
+            dtype: Some(ModelDtype::Float32),
+        },
         ..Default::default()
     }
 }

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

@@ -925,7 +925,6 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() {
     .await;
     assert_eq!(client.engine_identities()[0], b"engine-0");
     assert!(client.ready_responses()[0].max_model_len > 0);
-    assert_eq!(client.vllm_version(), "test-vllm-version");
 
     let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap();
     let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap();

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

@@ -104,8 +104,8 @@ pub struct ConnectedEngine {
     /// The identity of the connected engine.
     pub engine_id: EngineId,
     /// Post-initialization configuration received from the engine on the input
-    /// socket registration message.
-    pub ready_response: EngineCoreReadyResponse,
+    /// socket registration message. `None` until the registration is received.
+    pub ready_response: Option<EngineCoreReadyResponse>,
 }
 
 /// Represents the connected shared transport plus all registered engines after
@@ -295,9 +295,18 @@ pub async fn connect_handshake(
         }
     }
 
-    // 6. Wait for every engine to connect to the shared input socket and register itself.
-    let engines =
-        wait_for_input_registrations(&mut input_socket, engines.into_keys(), ready_timeout).await?;
+    // 4. Wait for every engine to connect to the shared input socket and register itself. The
+    //    `ready_response` is a placeholder; it is populated for each engine by
+    //    `wait_for_input_registrations` below.
+    let mut engines: Vec<_> = engines
+        .into_keys()
+        .map(|engine_id| ConnectedEngine {
+            engine_id,
+            ready_response: None,
+        })
+        .collect();
+
+    wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?;
     debug!(
         engine_count = engines.len(),
         "all engines registered on shared input socket"
@@ -340,13 +349,15 @@ pub async fn connect_bootstrapped(
     let mut output_socket = PullSocket::new();
     let output_address = output_socket.bind(output_address).await?.to_string();
 
-    let engines = wait_for_input_registrations(
-        &mut input_socket,
-        // TODO: follow start rank
-        (0..engine_count).map(|index| EngineId::from((index as u16).to_le_bytes().to_vec())),
-        ready_timeout,
-    )
-    .await?;
+    // TODO: follow start rank
+    let mut engines = (0..engine_count)
+        .map(|index| ConnectedEngine {
+            engine_id: EngineId::from((index as u16).to_le_bytes().to_vec()),
+            ready_response: None,
+        })
+        .collect::<Vec<_>>();
+
+    wait_for_input_registrations(&mut input_socket, &mut engines, ready_timeout).await?;
     info!(
         engine_count = engines.len(),
         "bootstrapped engines connected"
@@ -444,14 +455,17 @@ async fn send_init_message(
 /// Simplify API server handshake"), the payload is a msgpack-encoded
 /// [`EngineCoreReadyResponse`] carrying post-initialization values such as
 /// `max_model_len`.
+///
+/// Older engines sent an empty second frame here just to establish the
+/// ROUTER/DEALER backchannel, with no structured payload on the input socket.
+/// We continue to tolerate that legacy shape so the frontend can still connect
+/// to slightly older local engine checkouts.
 async fn wait_for_input_registrations(
     input_socket: &mut RouterSocket,
-    expected_engines: impl IntoIterator<Item = EngineId>,
+    engines: &mut [ConnectedEngine],
     ready_timeout: Duration,
-) -> Result<Vec<ConnectedEngine>> {
-    let expected_engines = expected_engines.into_iter().collect::<Vec<_>>();
-    let mut pending = expected_engines.iter().cloned().collect::<BTreeSet<_>>();
-    let mut ready_responses = BTreeMap::new();
+) -> Result<()> {
+    let mut pending = engines.iter().map(|e| e.engine_id.clone()).collect::<BTreeSet<_>>();
 
     while !pending.is_empty() {
         let registration = timeout(ready_timeout, input_socket.recv()).await.map_err(|_| {
@@ -475,33 +489,29 @@ async fn wait_for_input_registrations(
             );
         }
 
-        if frames[1].is_empty() {
-            bail_unexpected_handshake_message!(
-                "expected msgpack EngineCoreReadyResponse for engine input registration, got empty payload from engine id {actual_id:?}"
+        let ready_response = if frames[1].is_empty() {
+            debug!(
+                ?actual_id,
+                "received legacy empty input registration from engine"
             );
-        }
+            None
+        } else {
+            let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?;
+            debug!(
+                ?actual_id,
+                ?ready_response,
+                "received input registration from engine"
+            );
+            Some(ready_response)
+        };
 
-        let ready_response: EngineCoreReadyResponse = decode_msgpack(&frames[1])?;
-        debug!(
-            ?actual_id,
-            ?ready_response,
-            "received input registration from engine"
-        );
-        ready_responses.insert(actual_id, ready_response);
+        // Store the ready response in the corresponding engine entry.
+        if let Some(engine) = engines.iter_mut().find(|e| e.engine_id == actual_id) {
+            engine.ready_response = ready_response;
+        }
     }
 
-    Ok(expected_engines
-        .into_iter()
-        .map(|engine_id| {
-            let ready_response = ready_responses
-                .remove(&engine_id)
-                .expect("every expected engine id has a decoded ready response");
-            ConnectedEngine {
-                engine_id,
-                ready_response,
-            }
-        })
-        .collect())
+    Ok(())
 }
 
 /// Send an encoded message to the engine through the input socket.

rust/src/mock-engine/src/tests.rs

@@ -98,8 +98,7 @@ async fn mock_engine_connects_over_tcp() {
     let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await;
     assert_eq!(client.engine_count(), 1);
     assert_eq!(client.engine_identities()[0], &[0, 0]);
-    assert_eq!(client.max_model_len(), 1024 * 1024);
-    assert_eq!(client.vllm_version(), "test-vllm-version");
+    assert_eq!(client.max_model_len(), Some(1024 * 1024));
     shutdown_mock(client, shutdown, task).await;
 }
 

rust/src/reasoning-parser/src/lib.rs

@@ -19,7 +19,6 @@ mod deepseek_r1;
 mod delimited;
 mod gemma4;
 mod kimi;
-mod minimax_m3;
 mod qwen3;
 
 use thiserror::Error;
@@ -30,7 +29,6 @@ pub use self::deepseek_r1::DeepSeekR1ReasoningParser;
 pub(crate) use self::delimited::DelimitedReasoningParser;
 pub use self::gemma4::Gemma4ReasoningParser;
 pub use self::kimi::KimiReasoningParser;
-pub use self::minimax_m3::MiniMaxM3ReasoningParser;
 pub use self::qwen3::Qwen3ReasoningParser;
 
 /// DeepSeek V3 currently shares the standard `<think>...</think>` parser.

rust/src/reasoning-parser/src/minimax_m3.rs

@@ -1,43 +0,0 @@
-use vllm_tokenizer::DynTokenizer;
-
-use super::{DelimitedReasoningParser, ReasoningDelta, ReasoningParser, Result};
-
-/// Reasoning parser for MiniMax M3 style outputs.
-///
-/// MiniMax M3 uses `<mm:think>...</mm:think>` delimiters. Its chat template may
-/// prefill either delimiter depending on the requested thinking mode, so the
-/// shared delimited parser derives the starting state from the rendered prompt.
-pub struct MiniMaxM3ReasoningParser {
-    inner: DelimitedReasoningParser,
-}
-
-impl MiniMaxM3ReasoningParser {
-    /// Create a MiniMax M3 parser backed by the shared delimited state machine.
-    pub fn new(tokenizer: DynTokenizer) -> Result<Self> {
-        Ok(Self {
-            inner: DelimitedReasoningParser::new(tokenizer, "<mm:think>", "</mm:think>", false)?,
-        })
-    }
-}
-
-impl ReasoningParser for MiniMaxM3ReasoningParser {
-    fn create(tokenizer: DynTokenizer) -> Result<Box<dyn ReasoningParser>>
-    where
-        Self: Sized + 'static,
-    {
-        Ok(Box::new(Self::new(tokenizer)?))
-    }
-
-    fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> {
-        self.inner.initialize(prompt_token_ids);
-        Ok(())
-    }
-
-    fn push(&mut self, delta: &str) -> Result<ReasoningDelta> {
-        Ok(self.inner.push(delta))
-    }
-
-    fn finish(&mut self) -> Result<ReasoningDelta> {
-        Ok(self.inner.finish())
-    }
-}

rust/src/reasoning-parser/src/tests.rs

@@ -3,8 +3,7 @@ use std::sync::Arc;
 use vllm_tokenizer::Tokenizer;
 
 use super::{
-    DeepSeekR1ReasoningParser, DelimitedReasoningParser, MiniMaxM3ReasoningParser,
-    Qwen3ReasoningParser, ReasoningParser,
+    DeepSeekR1ReasoningParser, DelimitedReasoningParser, Qwen3ReasoningParser, ReasoningParser,
 };
 
 struct FakeTokenizer;
@@ -33,8 +32,6 @@ impl Tokenizer for FakeTokenizer {
             "<|END_THINKING|>" => Some(4),
             "◁think▷" => Some(5),
             "◁/think▷" => Some(6),
-            "<mm:think>" => Some(8),
-            "</mm:think>" => Some(9),
             _ => None,
         }
     }
@@ -162,35 +159,3 @@ fn deepseek_r1_stops_scanning_at_last_special_token() {
     assert_eq!(delta.reasoning.as_deref(), Some("reason"));
     assert_eq!(delta.content.as_deref(), Some("answer"));
 }
-
-#[test]
-fn minimax_m3_handles_explicit_think_delimiters() {
-    let tokenizer = Arc::new(FakeTokenizer);
-    let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
-
-    let delta = parser.push("<mm:think>reason</mm:think>answer").unwrap();
-    assert_eq!(delta.reasoning.as_deref(), Some("reason"));
-    assert_eq!(delta.content.as_deref(), Some("answer"));
-}
-
-#[test]
-fn minimax_m3_uses_prompt_prefilled_start_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
-    let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
-    parser.initialize(&[8]).unwrap();
-
-    let delta = parser.push("reason</mm:think>answer").unwrap();
-    assert_eq!(delta.reasoning.as_deref(), Some("reason"));
-    assert_eq!(delta.content.as_deref(), Some("answer"));
-}
-
-#[test]
-fn minimax_m3_uses_prompt_prefilled_end_marker() {
-    let tokenizer = Arc::new(FakeTokenizer);
-    let mut parser = MiniMaxM3ReasoningParser::new(tokenizer).unwrap();
-    parser.initialize(&[9]).unwrap();
-
-    let delta = parser.push("answer").unwrap();
-    assert_eq!(delta.reasoning, None);
-    assert_eq!(delta.content.as_deref(), Some("answer"));
-}

rust/src/server/src/routes.rs

@@ -6,7 +6,6 @@ mod load;
 mod metrics;
 pub(crate) mod openai;
 mod sleep;
-mod version;
 
 use std::sync::Arc;
 
@@ -36,7 +35,6 @@ fn build_router_with_dev_mode(state: Arc<AppState>, dev_mode_enabled: bool) -> R
         .route("/health", get(health::health))
         .route("/metrics", get(metrics::scrape))
         .route("/load", get(load::load))
-        .route("/version", get(version::version))
         // OpenAI-compatible endpoints
         .route("/v1/models", get(openai::list_models))
         .route("/v1/completions", post(openai::completions))

rust/src/server/src/routes/tests.rs

@@ -994,27 +994,6 @@ async fn list_models_returns_configured_model() {
     assert_eq!(json["data"][0]["id"], "Qwen/Qwen1.5-0.5B-Chat");
 }
 
-#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-#[serial]
-async fn version_returns_engine_vllm_version() {
-    let mut app = test_app().await;
-    let response = app
-        .call(Request::builder().uri("/version").body(Body::empty()).expect("build request"))
-        .await
-        .expect("call app");
-
-    assert_eq!(response.status(), StatusCode::OK);
-    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
-    let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
-    assert_eq!(
-        json,
-        json!({
-            "version": "test-vllm-version",
-            "rust_frontend_version": env!("CARGO_PKG_VERSION"),
-        })
-    );
-}
-
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn http_metrics_record_list_models_requests() {

rust/src/server/src/routes/version.rs

@@ -1,23 +0,0 @@
-use std::sync::Arc;
-
-use axum::Json;
-use axum::extract::State;
-use serde::Serialize;
-
-use crate::state::AppState;
-
-#[derive(Serialize)]
-pub(crate) struct VersionResponse {
-    version: String,
-    rust_frontend_version: &'static str,
-}
-
-/// Get engine and Rust frontend version metadata.
-pub async fn version(State(state): State<Arc<AppState>>) -> Json<VersionResponse> {
-    let version = state.engine_core_client().vllm_version().to_string();
-
-    Json(VersionResponse {
-        version,
-        rust_frontend_version: env!("CARGO_PKG_VERSION"),
-    })
-}

rust/src/text/src/lib.rs

@@ -45,9 +45,9 @@ pub struct TextLlm {
     /// Tokenizer/model metadata backend responsible for prompt encode/decode
     /// and sampling hints.
     backend: DynTextBackend,
-    /// Context window size reported by the engine startup handshake, with
-    /// optional override from config.
-    max_model_len: u32,
+    /// Context window size derived by the backend or from engine startup
+    /// handshake, with optional override from config.
+    max_model_len: Option<u32>,
 }
 
 impl TextLlm {
@@ -71,7 +71,7 @@ impl TextLlm {
     /// This takes priority over both the engine-reported default and any
     /// tokenizer/model metadata exposed by the backend.
     pub fn with_max_model_len(mut self, max_model_len: u32) -> Self {
-        self.max_model_len = max_model_len;
+        self.max_model_len = Some(max_model_len);
         self
     }
 
@@ -129,7 +129,9 @@ impl TextLlm {
         };
 
         let mut sampling_hints = self.backend.sampling_hints()?;
-        sampling_hints.max_model_len = Some(self.max_model_len);
+        if let Some(max_model_len) = self.max_model_len {
+            sampling_hints.max_model_len = Some(max_model_len);
+        }
         let PreparedTextRequest {
             text_request,
             generate_request,

rust/src/tool-parser/python/Cargo.toml

@@ -1,23 +0,0 @@
-[package]
-name = "vllm-tool-parser-py"
-version.workspace = true
-edition.workspace = true
-license.workspace = true
-
-[lib]
-name = "_rust_tool_parser"
-crate-type = ["cdylib", "rlib"]
-
-[features]
-default = []
-extension-module = ["pyo3/extension-module"]
-
-[dependencies]
-pyo3.workspace = true
-pythonize = { workspace = true, features = ["serde_json"] }
-serde_json.workspace = true
-thiserror-ext.workspace = true
-vllm-tool-parser.workspace = true
-
-[lints]
-workspace = true

rust/src/tool-parser/python/src/lib.rs

@@ -1,367 +0,0 @@
-//! Thin PyO3 bindings for `vllm_tool_parser`.
-//!
-//! This crate exposes the Rust tool parser trait and data shapes to Python
-//! while keeping parser state, grammar, and schema-aware argument conversion in
-//! Rust. Python callers should use this module as a typed bridge and keep any
-//! vLLM protocol adaptation outside the binding.
-
-use pyo3::exceptions::PyValueError;
-use pyo3::prelude::*;
-use pyo3::types::{PyAny, PyModule};
-use pythonize::{depythonize, pythonize};
-use serde_json::Value;
-use thiserror_ext::AsReport as _;
-use vllm_tool_parser::{Tool, ToolCallDelta, ToolParser, ToolParserOutput};
-
-macro_rules! tool_parser_factory {
-    ($($parser:ident),+ $(,)?) => {
-        fn create_tool_parser(
-            name: &str,
-            tools: &[Tool],
-        ) -> PyResult<Box<dyn ToolParser>> {
-            match name {
-                $(
-                    stringify!($parser) => {
-                        <vllm_tool_parser::$parser as ToolParser>::create(tools)
-                    }
-                )+
-                _ => {
-                    return Err(PyValueError::new_err(format!(
-                        "unsupported tool parser `{name}`"
-                    )));
-                }
-            }
-            .map_err(|error| PyValueError::new_err(error.to_report_string()))
-        }
-    };
-}
-
-// Export a tool parser to Python by registering it here.
-tool_parser_factory! {
-    DeepSeekV4ToolParser,
-    MinimaxM3ToolParser,
-}
-
-#[pyclass(name = "Tool", module = "vllm._rust_tool_parser", skip_from_py_object)]
-#[derive(Clone)]
-struct PyTool(Tool);
-
-#[pymethods]
-impl PyTool {
-    #[new]
-    #[pyo3(signature = (name, description, parameters, strict=None))]
-    fn new(
-        name: String,
-        description: Option<String>,
-        parameters: &Bound<'_, PyAny>,
-        strict: Option<bool>,
-    ) -> PyResult<Self> {
-        let parameters = depythonize::<Value>(parameters).map_err(|error| {
-            PyValueError::new_err(format!(
-                "failed to convert tool parameters from Python to JSON: {error}"
-            ))
-        })?;
-        Ok(Self(Tool {
-            name,
-            description,
-            parameters,
-            strict,
-        }))
-    }
-
-    #[getter]
-    fn name(&self) -> &str {
-        &self.0.name
-    }
-
-    #[getter]
-    fn description(&self) -> Option<&str> {
-        self.0.description.as_deref()
-    }
-
-    #[getter]
-    fn parameters(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
-        pythonize(py, &self.0.parameters).map(Bound::unbind).map_err(|error| {
-            PyValueError::new_err(format!(
-                "failed to convert tool parameters from JSON to Python: {error}"
-            ))
-        })
-    }
-
-    #[getter]
-    fn strict(&self) -> Option<bool> {
-        self.0.strict
-    }
-}
-
-#[pyclass(
-    name = "ToolCallDelta",
-    module = "vllm._rust_tool_parser",
-    skip_from_py_object
-)]
-#[derive(Clone)]
-struct PyToolCallDelta(ToolCallDelta);
-
-#[pymethods]
-impl PyToolCallDelta {
-    #[new]
-    #[pyo3(signature = (tool_index, name, arguments))]
-    fn new(tool_index: usize, name: Option<String>, arguments: String) -> Self {
-        Self(ToolCallDelta {
-            tool_index,
-            name,
-            arguments,
-        })
-    }
-
-    #[getter]
-    fn tool_index(&self) -> usize {
-        self.0.tool_index
-    }
-
-    #[getter]
-    fn name(&self) -> Option<&str> {
-        self.0.name.as_deref()
-    }
-
-    #[getter]
-    fn arguments(&self) -> &str {
-        &self.0.arguments
-    }
-}
-
-#[pyclass(
-    name = "ToolParserOutput",
-    module = "vllm._rust_tool_parser",
-    skip_from_py_object
-)]
-#[derive(Clone)]
-struct PyToolParserOutput(ToolParserOutput);
-
-#[pymethods]
-impl PyToolParserOutput {
-    #[new]
-    #[pyo3(signature = (normal_text="", calls=None))]
-    fn new(py: Python<'_>, normal_text: &str, calls: Option<Vec<Py<PyToolCallDelta>>>) -> Self {
-        let calls =
-            calls.unwrap_or_default().iter().map(|call| call.borrow(py).0.clone()).collect();
-        Self(ToolParserOutput {
-            normal_text: normal_text.to_owned(),
-            calls,
-        })
-    }
-
-    #[getter]
-    fn normal_text(&self) -> &str {
-        &self.0.normal_text
-    }
-
-    #[getter]
-    fn calls(&self) -> Vec<PyToolCallDelta> {
-        self.0.calls.iter().cloned().map(PyToolCallDelta).collect()
-    }
-
-    fn append(&mut self, other: PyRef<'_, PyToolParserOutput>) {
-        self.0.append(other.0.clone());
-    }
-
-    fn coalesce_calls(&self) -> Self {
-        Self(self.0.clone().coalesce_calls())
-    }
-}
-
-#[pyclass(name = "ToolParser", module = "vllm._rust_tool_parser", unsendable)]
-struct PyToolParser(Box<dyn ToolParser>);
-
-impl PyToolParser {
-    fn parse_into_output(&mut self, chunk: &str, output: &mut PyToolParserOutput) -> PyResult<()> {
-        self.0
-            .parse_into(chunk, &mut output.0)
-            .map_err(|error| PyValueError::new_err(error.to_report_string()))
-    }
-}
-
-#[pymethods]
-impl PyToolParser {
-    #[new]
-    fn new(py: Python<'_>, parser_name: &str, tools: Vec<Py<PyTool>>) -> PyResult<Self> {
-        let tools = tools.iter().map(|tool| tool.borrow(py).0.clone()).collect::<Vec<_>>();
-        create_tool_parser(parser_name, &tools).map(Self)
-    }
-
-    fn parse_into(
-        &mut self,
-        chunk: &str,
-        mut output: PyRefMut<'_, PyToolParserOutput>,
-    ) -> PyResult<()> {
-        self.parse_into_output(chunk, &mut output)
-    }
-
-    fn finish(&mut self) -> PyResult<PyToolParserOutput> {
-        self.0
-            .finish()
-            .map(PyToolParserOutput)
-            .map_err(|error| PyValueError::new_err(error.to_report_string()))
-    }
-
-    fn reset(&mut self) -> String {
-        self.0.reset()
-    }
-
-    fn preserve_special_tokens(&self) -> bool {
-        self.0.preserve_special_tokens()
-    }
-}
-
-#[pymodule]
-fn _rust_tool_parser(m: &Bound<'_, PyModule>) -> PyResult<()> {
-    m.add_class::<PyTool>()?;
-    m.add_class::<PyToolCallDelta>()?;
-    m.add_class::<PyToolParserOutput>()?;
-    m.add_class::<PyToolParser>()?;
-    Ok(())
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use serde_json::json;
-
-    fn with_python<R>(f: impl for<'py> FnOnce(Python<'py>) -> R) -> R {
-        Python::initialize();
-        Python::attach(f)
-    }
-
-    fn tool_schema() -> Value {
-        json!({
-            "type": "object",
-            "properties": {
-                "user_id": {"type": "integer"},
-                "shipping": {
-                    "type": "object",
-                    "properties": {
-                        "city": {"type": "string"},
-                        "zip": {"type": "integer"}
-                    }
-                }
-            }
-        })
-    }
-
-    fn build_call() -> String {
-        r#"<｜DSML｜tool_calls>
-<｜DSML｜invoke name="create_order">
-<｜DSML｜parameter name="user_id" string="false">42</｜DSML｜parameter>
-<｜DSML｜parameter name="shipping" string="false">{"city":"Singapore","zip":18956}</｜DSML｜parameter>
-</｜DSML｜invoke>
-</｜DSML｜tool_calls>"#
-            .to_owned()
-    }
-
-    fn make_py_tool(py: Python<'_>) -> PyResult<Py<PyTool>> {
-        let parameters = pythonize(py, &tool_schema()).map_err(|error| {
-            PyValueError::new_err(format!(
-                "failed to convert test schema from JSON to Python: {error}"
-            ))
-        })?;
-        Py::new(
-            py,
-            PyTool::new(
-                "create_order".to_owned(),
-                Some("Create an order".to_owned()),
-                &parameters,
-                None,
-            )?,
-        )
-    }
-
-    #[test]
-    fn tool_round_trips_typed_fields() {
-        with_python(|py| {
-            let tool = make_py_tool(py)?;
-            let borrowed = tool.borrow(py);
-            assert_eq!(borrowed.name(), "create_order");
-            assert_eq!(borrowed.description(), Some("Create an order"));
-            assert_eq!(borrowed.strict(), None);
-
-            let parameters = borrowed.parameters(py)?;
-            let parameters = depythonize::<Value>(parameters.bind(py))?;
-            assert_eq!(parameters, tool_schema());
-            PyResult::Ok(())
-        })
-        .unwrap();
-    }
-
-    #[test]
-    fn output_append_and_coalesce_calls() {
-        with_python(|py| {
-            let first = Py::new(
-                py,
-                PyToolCallDelta::new(0, Some("create_order".to_owned()), "{\"a\"".to_owned()),
-            )?;
-            let second = Py::new(py, PyToolCallDelta::new(0, None, ":1}".to_owned()))?;
-            let mut output = PyToolParserOutput::new(py, "text", Some(vec![first]));
-            let other = Py::new(py, PyToolParserOutput::new(py, "", Some(vec![second])))?;
-            output.append(other.borrow(py));
-
-            let coalesced = output.coalesce_calls();
-            assert_eq!(coalesced.normal_text(), "text");
-            let calls = coalesced.calls();
-            assert_eq!(calls.len(), 1);
-            assert_eq!(calls[0].tool_index(), 0);
-            assert_eq!(calls[0].name(), Some("create_order"));
-            assert_eq!(calls[0].arguments(), "{\"a\":1}");
-            PyResult::Ok(())
-        })
-        .unwrap();
-    }
-
-    #[test]
-    fn parser_parse_finish_and_preserve_special_tokens() {
-        with_python(|py| {
-            let tool = make_py_tool(py)?;
-            let mut parser = PyToolParser::new(py, "DeepSeekV4ToolParser", vec![tool])?;
-            assert!(parser.preserve_special_tokens());
-
-            let mut output = PyToolParserOutput::new(py, "", None);
-            parser.parse_into_output(&build_call(), &mut output)?;
-            let finish = Py::new(py, parser.finish()?)?;
-            output.append(finish.borrow(py));
-            let output = output.coalesce_calls();
-
-            assert_eq!(output.normal_text(), "");
-            let calls = output.calls();
-            assert_eq!(calls.len(), 1);
-            assert_eq!(calls[0].name(), Some("create_order"));
-            assert_eq!(
-                serde_json::from_str::<Value>(calls[0].arguments()).unwrap(),
-                json!({
-                    "user_id": 42,
-                    "shipping": {
-                        "city": "Singapore",
-                        "zip": 18956
-                    }
-                })
-            );
-
-            assert_eq!(parser.reset(), "");
-            PyResult::Ok(())
-        })
-        .unwrap();
-    }
-
-    #[test]
-    fn parser_errors_for_unknown_name() {
-        with_python(|py| {
-            let tool = make_py_tool(py)?;
-            let error = match PyToolParser::new(py, "missing", vec![tool]) {
-                Ok(_) => panic!("missing parser name unexpectedly succeeded"),
-                Err(error) => error,
-            };
-            let message = format!("{error}");
-            assert!(message.contains("unsupported tool parser `missing`"));
-            PyResult::Ok(())
-        })
-        .unwrap();
-    }
-}

rust/src/tool-parser/src/deepseek_dsml/mod.rs

@@ -104,7 +104,7 @@ impl DeepSeekDsmlToolParser {
                         self.tool_parameters.convert_param_with_schema(
                             &name,
                             &param.name,
-                            param.value,
+                            &param.value,
                         )
                     };
                     arguments.insert(param.name, value);