[CI] Inline env vars lost by CONTINUE_ON_FAILURE subshell wrapping

Postmerge/nightly/daily builds set CONTINUE_ON_FAILURE=1, which wraps each YAML command in `(cmd) || CI_OVERALL_STATUS=1`. A standalone `export VAR=val` executes inside the subshell and the variable is immediately lost — subsequent commands never see it. This caused real failures (model_executor fastsafetensors crash from missing VLLM_WORKER_MULTIPROC_METHOD=spawn) and silent coverage gaps (rust_frontend tests silently running the Python frontend because VLLM_USE_RUST_FRONTEND=1 was never set). Fix: inline env vars as command prefixes (`VAR=val command`) instead of standalone `export` lines across all 11 affected CI YAML files. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tyler@tylermsmith.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-06-22 17:53:18 -04:00
93 changed files with 841 additions and 3935 deletions
@@ -136,7 +136,7 @@ steps:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 45m "
      pytest -x -v -s tests/models/multimodal/generation --ignore=tests/models/multimodal/generation/test_pixtral.py -m cpu_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB"
-  parallelism: 4
+  parallelism: 3

 - label: "Arm CPU Test"
  depends_on: []
@@ -21,13 +21,12 @@ else
  exit 0
 fi

-# build for arm64 GPU targets: Grace/GH200 (sm_90) and DGX Spark/GB10
-# (sm_121, family-covered by 12.0 under CUDA 13)
+# build (Grace/GH200 is the arm64 GPU target; sm_90)
 docker build --file docker/Dockerfile \
  --platform linux/arm64 \
  --build-arg max_jobs=16 \
  --build-arg nvcc_threads=4 \
-  --build-arg torch_cuda_arch_list="9.0 12.0" \
+  --build-arg torch_cuda_arch_list="9.0" \
  --build-arg USE_SCCACHE=1 \
  --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
  --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \
@@ -1,5 +1,5 @@
 group: Expert Parallelism
-depends_on:
+depends_on: 
  - image-build-xpu
 steps:
 - label: EPLB Algorithm
@@ -1,5 +1,5 @@
 group: Models - Multimodal
-depends_on:
+depends_on: 
  - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
@@ -12,7 +12,6 @@ steps:
  - tests/basic_correctness/test_cpu_offload
  - tests/basic_correctness/test_mem.py
  commands:
-  - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  - pytest -v -s basic_correctness/test_mem.py
-  - pytest -v -s basic_correctness/test_basic_correctness.py
-  - pytest -v -s basic_correctness/test_cpu_offload.py
+  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_mem.py
+  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_basic_correctness.py
+  - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s basic_correctness/test_cpu_offload.py
@@ -14,8 +14,7 @@ steps:
  - vllm/v1/cudagraph_dispatcher.py
  - tests/compile/correctness_e2e/test_sequence_parallel.py
  commands:
-  - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-  - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
+  - VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py

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

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

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

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

 - label: Fusion and Compile Unit Tests (2xB200)
  key: fusion-and-compile-unit-tests-2xb200
@@ -32,11 +32,10 @@ steps:
  - tests/entrypoints/openai/test_multi_api_servers.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
-  - TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
-  - DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
+  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_async_llm_dp.py
+  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
+  - NCCL_CUMEM_HOST_ENABLE=0 TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py
+  - NCCL_CUMEM_HOST_ENABLE=0 DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py

 - label: Distributed Compile + RPC Tests (2 GPUs)
  key: distributed-compile-rpc-tests-2-gpus
@@ -56,10 +55,9 @@ steps:
  - tests/entrypoints/llm/test_collective_rpc.py
  commands:
  # https://github.com/NVIDIA/nccl/issues/1838
-  - export NCCL_CUMEM_HOST_ENABLE=0
-  - pytest -v -s entrypoints/llm/test_collective_rpc.py
-  - pytest -v -s ./compile/fullgraph/test_basic_correctness.py
-  - pytest -v -s ./compile/test_wrapper.py
+  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s entrypoints/llm/test_collective_rpc.py
+  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/fullgraph/test_basic_correctness.py
+  - NCCL_CUMEM_HOST_ENABLE=0 pytest -v -s ./compile/test_wrapper.py

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

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

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

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

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

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

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

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

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

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

 - label: LM Eval Small Models (B200)
  key: lm-eval-small-models-b200
@@ -108,9 +107,7 @@ steps:
      depends_on:
      - image-build-amd
      commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
-      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt
+      - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_ROCM_ARCH=gfx942 pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt  # Limit Quark compilation to save time

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

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

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

- label: GPQA Eval (GPT-OSS) (DGX Spark)
-  key: gpqa-eval-gpt-oss-spark
-  timeout_in_minutes: 120
-  device: dgx-spark
-  optional: true
-  num_devices: 1
-  source_file_dependencies:
-  - csrc/
-  - vllm/model_executor/layers/quantization
-  - tests/evals/gpt_oss/
-  commands:
-    - uv pip install --system 'gpt-oss[eval]==0.0.5'
-    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
-
 - label: MRCR Eval Small Models
  device: h200_35gb
  timeout_in_minutes: 30
@@ -36,14 +36,14 @@ steps:
  commands:
    # FIXIT: find out which code initialize cuda before running the test
    # before the fix, we need to use spawn to test it
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    #
    # Alot of these tests are on the edge of OOMing
-    - export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+    #
    # There is some Tensor Parallelism related processing logic in LoRA that
    # requires multi-GPU testing for validation.
-    - pytest -v -s -x lora/test_chatglm3_tp.py
-    - pytest -v -s -x lora/test_llama_tp.py
-    - pytest -v -s -x lora/test_qwen3_with_multi_loras.py
-    - pytest -v -s -x lora/test_olmoe_tp.py
-    - pytest -v -s -x lora/test_gptoss_tp.py
-    - pytest -v -s -x lora/test_qwen35_densemodel_lora.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_chatglm3_tp.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_llama_tp.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen3_with_multi_loras.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_olmoe_tp.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_gptoss_tp.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True pytest -v -s -x lora/test_qwen35_densemodel_lora.py
@@ -18,9 +18,8 @@ steps:
    - vllm/v1/
    - tests/v1/spec_decode
  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # TODO: create another `optional` test group for slow tests
-    - pytest -v -s -m 'not slow_test' v1/spec_decode
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not slow_test' v1/spec_decode
  mirror:
    amd:
      device: mi300_1
@@ -50,12 +49,11 @@ steps:
    - tests/v1/test_request.py
    - tests/v1/test_outputs.py
  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s v1/sample
-    - pytest -v -s v1/logits_processors
-    - pytest -v -s v1/test_oracle.py
-    - pytest -v -s v1/test_request.py
-    - pytest -v -s v1/test_outputs.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/sample
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/logits_processors
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_oracle.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_request.py
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/test_outputs.py
  mirror:
    amd:
      device: mi325_1
@@ -93,18 +91,17 @@ steps:
    - tests/entrypoints/openai/correctness/test_lmeval.py
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    # split the test to avoid interference
-    - pytest -v -s -m 'not cpu_test' v1/core
-    - pytest -v -s v1/executor
-    - pytest -v -s v1/kv_offload
-    - pytest -v -s v1/simple_kv_offload
-    - pytest -v -s v1/worker
-    - pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
-    - pytest -v -s -m 'not cpu_test' v1/metrics
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/core
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/executor
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_offload
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/simple_kv_offload
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/worker
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'not cpu_test' v1/metrics
    # Integration test for streaming correctness (requires special branch).
    - pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
-    - pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
  mirror:
    amd:
      device: mi325_1
@@ -153,8 +150,7 @@ steps:
    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
    - tests/v1/kv_connector/extract_hidden_states_integration
  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s v1/kv_connector/extract_hidden_states_integration
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s v1/kv_connector/extract_hidden_states_integration

 - label: Extract Hidden States Integration (2 GPUs)
  key: extract-hidden-states-integration-2-gpus
@@ -167,8 +163,7 @@ steps:
    - vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
    - tests/v1/kv_connector/extract_hidden_states_integration
  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration
+    - VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s -m 'distributed' v1/kv_connector/extract_hidden_states_integration

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

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

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

 - label: Model Runner V2 Examples
  device: h200_35gb
@@ -42,26 +41,25 @@ steps:
    - examples/features/tensorize_vllm_model.py
  commands:
    - set -x
-    - export VLLM_USE_V2_MODEL_RUNNER=1
    - pip install tensorizer # for tensorizer test
-    - python3 basic/offline_inference/chat.py # for basic
-    - python3 basic/offline_inference/generate.py --model facebook/opt-125m
-    #- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
-    #- python3 basic/offline_inference/embed.py   # TODO
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/chat.py # for basic
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model facebook/opt-125m
+    #- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10  # TODO
+    #- VLLM_USE_V2_MODEL_RUNNER=1 python3 basic/offline_inference/embed.py   # TODO
    # for multi-modal models
-    - python3 generate/multimodal/audio_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_offline.py --seed 0
-    - python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
-    - python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/audio_language_offline.py --seed 0
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_offline.py --seed 0
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
    # for pooling models
-    - python3 pooling/embed/vision_embedding_offline.py --seed 0
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 pooling/embed/vision_embedding_offline.py --seed 0
    # for features demo
-    - python3 features/automatic_prefix_caching/prefix_caching_offline.py
-    - python3 deployment/llm_engine_example.py
-    - python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 features/automatic_prefix_caching/prefix_caching_offline.py
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 deployment/llm_engine_example.py
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && VLLM_USE_V2_MODEL_RUNNER=1 python3 features/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
    # https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
-    - python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
+    - VLLM_USE_V2_MODEL_RUNNER=1 python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536

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

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

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

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

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

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

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

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

 - label: Spec Decode MTP hybrid (B200)
  timeout_in_minutes: 30
@@ -8,73 +8,43 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR})
  set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
 endif()

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

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

 cuda_archs_loose_intersection(DEEPGEMM_ARCHS
@@ -6,48 +6,26 @@ if(DEFINED ENV{QUTLASS_SRC_DIR})
  set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
 endif()

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

+FetchContent_Populate(qutlass)
+
 if(NOT qutlass_SOURCE_DIR)
  message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")
 endif()
@@ -133,6 +133,16 @@ Priority is **1 = highest** (tried first).
 | 7 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 8 | `FLASHMLA_SPARSE` |

+**Ampere/Hopper (SM 8.x-9.x):**
+
+| Priority | Backend |
+| -------- | ------- |
+| 1 | `FLASH_ATTN_MLA` |
+| 2 | `FLASHMLA` |
+| 3 | `FLASHINFER_MLA` |
+| 4 | `TRITON_MLA` |
+| 5 | `FLASHMLA_SPARSE` |
+
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -221,8 +231,7 @@ MLA decode backends are selected using the standard
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
-| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
@@ -239,11 +248,10 @@ DeepSeek V4 sparse MLA uses its own decode backends, selected via
 `--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
 `FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
 pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
-default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and
-`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.
+default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.

 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
-| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x |
-| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
+| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
+| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
@@ -26,4 +26,4 @@ quack-kernels>=0.3.3
 tokenspeed-mla==0.1.2

 # Humming kernels for quantization gemm
-humming-kernels[cu13]==0.1.6
+humming-kernels[cu13]==0.1.4
@@ -59,8 +59,9 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
    return accuracy


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

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

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


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

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

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

    with RemoteOpenAIServer(
        MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
@@ -782,125 +782,3 @@ def test_rearrange_expert_weights_profile_mode(world_size):
        _test_rearrange_expert_weights_profile_mode,
        world_size,
    )
-
-
-def _test_nixl_deferred_init_worker(
-    env,
-    world_size: int,
-    num_layers: int,
-    num_local_experts: int,
-    num_logical_experts: int,
-) -> None:
-    """Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""
-    from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator
-
-    set_env_vars_and_device(env)
-
-    vllm_config = VllmConfig()
-    vllm_config.parallel_config.tensor_parallel_size = world_size
-
-    with set_current_vllm_config(vllm_config):
-        ensure_model_parallel_initialized(
-            tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
-        )
-
-        ep_group_coordinator = get_tp_group()
-        ep_group = ep_group_coordinator.cpu_group
-        ep_rank = torch.distributed.get_rank()
-        device = torch.device(f"cuda:{ep_rank}")
-
-        total_physical_experts = world_size * num_local_experts
-        hidden_sizes = [32, 64]
-
-        redundancy_config = create_redundancy_config(
-            num_logical_experts, total_physical_experts
-        )
-        old_indices = create_expert_indices_with_redundancy(
-            num_layers,
-            num_logical_experts,
-            total_physical_experts,
-            redundancy_config,
-        )
-
-        new_redundancy_config = create_redundancy_config(
-            num_logical_experts, total_physical_experts
-        )
-        new_indices = create_expert_indices_with_redundancy(
-            num_layers,
-            num_logical_experts,
-            total_physical_experts,
-            new_redundancy_config,
-        )
-
-        expert_weights = create_expert_weights(
-            num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
-        )
-
-        expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
-
-        communicator = NixlEplbCommunicator(
-            cpu_group=ep_group_coordinator.cpu_group,
-            all_expert_weights=expert_weights,
-            expert_buffer=expert_buffer,
-            defer_remote_setup=True,
-        )
-        assert not communicator._remote_state_initialized
-
-        rearrange_expert_weights_inplace(
-            old_indices,
-            new_indices,
-            expert_weights,
-            expert_buffer,
-            ep_group,
-            communicator,
-        )
-
-        assert communicator._remote_state_initialized
-
-    local_ok = verify_expert_weights_after_shuffle(
-        expert_weights,
-        new_indices,
-        hidden_sizes,
-        ep_rank,
-        num_local_experts,
-    )
-
-    local_ok = (
-        verify_redundant_experts_have_same_weights(
-            expert_weights,
-            new_indices,
-            hidden_sizes,
-            ep_rank,
-            world_size,
-            num_local_experts,
-        )
-        and local_ok
-    )
-    assert_verification_synced(
-        local_ok,
-        "Deferred NIXL init verification failed on at least one rank.",
-    )
-
-
-@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available")
-@pytest.mark.parametrize(
-    "world_size,num_layers,num_local_experts,num_logical_experts",
-    [(2, 2, 3, 4)],
-)
-def test_nixl_deferred_init(
-    world_size,
-    num_layers,
-    num_local_experts,
-    num_logical_experts,
-):
-    """Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""
-
-    if torch.accelerator.device_count() < world_size:
-        pytest.skip(f"Need at least {world_size} GPUs to run the test")
-    distributed_run(
-        _test_nixl_deferred_init_worker,
-        world_size,
-        num_layers,
-        num_local_experts,
-        num_logical_experts,
-    )
@@ -1,5 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-model_name: "openai/gpt-oss-20b"
-metric_threshold: 0.568
-reasoning_effort: "low"
@@ -1,2 +0,0 @@
-# DGX Spark model configurations for GPQA evaluation
-gpt-oss-20b-sm120.yaml
@@ -1,12 +0,0 @@
-model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
-accuracy_threshold: 0.86
-num_questions: 1319
-num_fewshot: 5
-server_args: >-
-  --enforce-eager
-  --max-model-len 8192
-  --tensor-parallel-size 1
-  --quantization humming
-  --kernel-config.enable_flashinfer_autotune=False
-env:
-  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'
@@ -1,10 +0,0 @@
-model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
-accuracy_threshold: 0.86
-num_questions: 1319
-num_fewshot: 5
-server_args: >-
-  --enforce-eager
-  --max-model-len 8192
-  --tensor-parallel-size 1
-  --quantization humming
-  --kernel-config.enable_flashinfer_autotune=False
@@ -1,2 +0,0 @@
-gpt-oss-20b-humming-act-fp8.yaml
-Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml
@@ -1,2 +0,0 @@
-gpt-oss-20b-humming.yaml
-Qwen3-30B-A3B-MXFP4A16-humming.yaml
@@ -1,11 +0,0 @@
-model_name: "openai/gpt-oss-20b"
-accuracy_threshold: 0.30
-num_questions: 1319
-num_fewshot: 5
-server_args: >-
-  --enforce-eager
-  --max-model-len 8192
-  --tensor-parallel-size 1
-  --moe-backend humming
-env:
-  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'
@@ -1,9 +0,0 @@
-model_name: "openai/gpt-oss-20b"
-accuracy_threshold: 0.30
-num_questions: 1319
-num_fewshot: 5
-server_args: >-
-  --enforce-eager
-  --max-model-len 8192
-  --tensor-parallel-size 1
-  --moe-backend humming
@@ -325,24 +325,3 @@ class TestWeightLoadingWithPaddedHiddenSize:
                shard_id="w2",
                expert_id=0,
            )
-
-
-class TestPerTensorScaleCoercion:
-    """Regression test for shape-(1,) per-tensor scales (issue #43297).
-
-    llm-compressor NVFP4 emits per-tensor weight and input scales as
-    shape-(1,) tensors. `_to_scalar` collapses them to a 0-D scalar so the
-    scalar-slot assignments in the weight loader neither broadcast nor raise.
-    """
-
-    def test_collapses_to_scalar(self):
-        # shape-(1,) and 0-D both reduce to a 0-D scalar.
-        for loaded_weight in (torch.tensor([0.5]), torch.tensor(0.5)):
-            scalar = RoutedExperts._to_scalar(loaded_weight)
-            assert scalar.shape == ()
-            assert scalar.item() == pytest.approx(0.5)
-
-    def test_rejects_non_scalar(self):
-        # numel > 1 must fail loudly instead of silently picking an element.
-        with pytest.raises(RuntimeError):
-            RoutedExperts._to_scalar(torch.tensor([0.1, 0.2]))
@@ -75,7 +75,7 @@ def get_ref_results(
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
-@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass", "cudnn", "trtllm", "b12x"])
+@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm", "b12x"])
@pytest.mark.parametrize("autotune", [False, True])
@torch.inference_mode()
 def test_flashinfer_nvfp4_gemm(
@@ -88,8 +88,6 @@ def test_flashinfer_nvfp4_gemm(
 ) -> None:
    if "trtllm" in backend and dtype == torch.float16:
        pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
-    if backend == "cute-dsl" and not current_platform.is_device_capability_family(100):
-        pytest.skip("FlashInfer cutedsl backend is only supported on SM10x")
    if backend == "b12x" and not current_platform.has_device_capability(120):
        pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
    if backend == "b12x" and not has_flashinfer_b12x_gemm():
@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+from hashlib import sha256
+from pathlib import Path
+from types import SimpleNamespace
+
+from vllm.model_executor.warmup import kernel_warmup
+
+
+def test_resolve_flashinfer_autotune_file_default_layout(
+    monkeypatch, tmp_path: Path
+) -> None:
+    fake_jit = SimpleNamespace(
+        env=SimpleNamespace(
+            FLASHINFER_WORKSPACE_DIR=Path("/flashinfer-cache/0.6.11.post2/103a")
+        )
+    )
+    fake_flashinfer = SimpleNamespace(jit=fake_jit)
+    monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
+    monkeypatch.setitem(sys.modules, "flashinfer.jit", fake_jit)
+    monkeypatch.setattr(
+        kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
+    )
+    monkeypatch.setattr(kernel_warmup.envs, "VLLM_CACHE_ROOT", str(tmp_path))
+    monkeypatch.setattr(kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None)
+
+    runner = SimpleNamespace(vllm_config=SimpleNamespace())
+    cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
+
+    path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
+
+    assert path == (
+        tmp_path
+        / "flashinfer_autotune_cache"
+        / "0.6.11.post2"
+        / "103a"
+        / cache_hash
+        / "autotune_configs.json"
+    )
+    assert path.parent.is_dir()
+
+
+def test_resolve_flashinfer_autotune_file_uses_override_dir(
+    monkeypatch, tmp_path: Path
+) -> None:
+    monkeypatch.setattr(
+        kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)
+    )
+    monkeypatch.setattr(
+        kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
+    )
+
+    runner = SimpleNamespace(vllm_config=SimpleNamespace())
+    cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
+
+    path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
+
+    assert path == tmp_path / cache_hash / "autotune_configs.json"
@@ -90,7 +90,6 @@ def test_models(example_prompts, model_name) -> None:
 EAGER = [True, False]

 SM_100_NVFP4_BACKENDS = [
-    "flashinfer_cutedsl",
    "flashinfer_cudnn",
    "flashinfer_trtllm",
    "flashinfer_cutlass",
@@ -103,18 +102,12 @@ SM_100_NVFP4_BACKENDS = [
    "backend",
    [
        "emulation",
-        "flashinfer_cutedsl",
        "flashinfer_cudnn",
        "flashinfer_trtllm",  # the small seq_len ensures trtllm_8x4_layout backend is used
        "flashinfer_cutlass",
    ],
 )
 def test_nvfp4(vllm_runner, model, eager, backend):
-    if backend == "flashinfer_cutedsl" and not (
-        current_platform.is_device_capability_family(100)
-    ):
-        pytest.skip("The flashinfer_cutedsl backend is only supported on SM10x")
-
    if (
        not current_platform.has_device_capability(100)
        and backend in SM_100_NVFP4_BACKENDS
@@ -1,54 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Behavior checks for FlashInfer SM120 sparse MLA backend selection."""
-
-from types import SimpleNamespace
-
-import torch
-
-from vllm.config import set_current_vllm_config
-from vllm.platforms.interface import DeviceCapability
-from vllm.utils import flashinfer as fi_utils
-from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
-    FlashInferMLASparseSM120Backend,
-)
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-
-
-def _fake_vllm_config(model_type: str) -> SimpleNamespace:
-    return SimpleNamespace(
-        model_config=SimpleNamespace(
-            hf_text_config=SimpleNamespace(model_type=model_type, index_topk=2048),
-        ),
-    )
-
-
-def test_sm120_backend_uses_dedicated_backend_name() -> None:
-    assert FlashInferMLASparseSM120Backend.get_name() == "FLASHINFER_MLA_SPARSE_SM120"
-    assert (
-        AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120.get_class()
-        is FlashInferMLASparseSM120Backend
-    )
-
-
-def test_v32_glm_sm120_backend_accepts_glm_block_size(
-    monkeypatch,
-) -> None:
-    monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True)
-
-    with set_current_vllm_config(_fake_vllm_config("glm4_moe")):
-        invalid_reasons = FlashInferMLASparseSM120Backend.validate_configuration(
-            head_size=576,
-            dtype=torch.bfloat16,
-            kv_cache_dtype="fp8",
-            block_size=256,
-            use_mla=True,
-            has_sink=False,
-            use_sparse=True,
-            use_mm_prefix=False,
-            use_per_head_quant_scales=False,
-            device_capability=DeviceCapability(12, 0),
-            attn_type="decoder",
-        )
-
-    assert invalid_reasons == []
@@ -36,7 +36,7 @@ if not current_platform.is_cuda():

 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
-    FlashInferMLASparseTRTLLMBackend,
+    FlashInferMLASparseBackend,
 )
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
    FlashMLASparseBackend,
@@ -174,8 +174,8 @@ def _quantize_dequantize_fp8_ds_mla(

@pytest.mark.parametrize(
    "backend_cls",
-    [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
-    ids=["FlashMLA", "FlashInferTRTLLM"],
+    [FlashMLASparseBackend, FlashInferMLASparseBackend],
+    ids=["FlashMLA", "FlashInfer"],
 )
@pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])
@@ -217,12 +217,9 @@ def test_sparse_backend_decode_correctness(
        ok, reason = flashmla.is_flashmla_sparse_supported()
        if not ok:
            pytest.skip(reason)
-    elif backend_cls == FlashInferMLASparseTRTLLMBackend:
-        device_capability = current_platform.get_device_capability()
-        if device_capability is None or not backend_cls.supports_compute_capability(
-            device_capability
-        ):
-            pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
+    elif backend_cls == FlashInferMLASparseBackend:
+        if not current_platform.has_device_capability(100):
+            pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher")

    batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
    use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
@@ -4708,38 +4708,6 @@ def test_free_encoder_inputs_respects_unconfirmed_placeholders():
    assert manager.get_cached_input_ids(request) == set()


-def test_free_encoder_inputs_defers_for_eagle_lookahead():
-    """With EAGLE speculative decoding, the encoder input is retained one extra
-    position so the drafter's +1 look-ahead mm-embedding gather (which reads one
-    position past the target's computed range) still finds it cached. This is
-    the primary mechanism that prevents the drafter "Encoder cache miss"; the
-    worker-side token-embedding fallback is only a backstop."""
-    scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
-    # create_scheduler only builds ngram spec configs; force the eagle path that
-    # _free_encoder_inputs keys off (self.use_eagle).
-    scheduler.use_eagle = True
-    mm_positions = [[PlaceholderRange(offset=50, length=100)]]
-    request = create_requests(
-        num_requests=1,
-        num_tokens=250,
-        mm_positions=mm_positions,
-    )[0]
-    manager = scheduler.encoder_cache_manager
-    manager.allocate(request, 0)
-    mm_end = 150  # offset + length
-
-    # Confirmed progress reaches the range end: without spec decode this frees
-    # (see test below), but the drafter's +1 look-ahead still needs it.
-    request.num_computed_tokens = mm_end
-    scheduler._free_encoder_inputs(request)
-    assert manager.get_cached_input_ids(request) == {0}
-
-    # One position past the range end: the +1 look-ahead has now passed it.
-    request.num_computed_tokens = mm_end + 1
-    scheduler._free_encoder_inputs(request)
-    assert manager.get_cached_input_ids(request) == set()
-
-
 def test_free_encoder_inputs_unchanged_without_spec_decode():
    """Without speculative decoding, encoder inputs are freed as soon as
    num_computed_tokens passes the placeholder range, as before."""
@@ -3,8 +3,6 @@

 import pytest

-from vllm.platforms import current_platform
-

 def test_mla_common_backend_rejects_cross_layer_kv_cache():
    """MLACommonBackend defaults to the identity permutation (layers dim
@@ -26,13 +24,8 @@ def test_mla_common_backend_rejects_cross_layer_kv_cache():

@pytest.mark.parametrize(
    "backend_path",
-    # See: https://github.com/vllm-project/vllm/issues/46411
    [
        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
-    ]
-    if current_platform.is_rocm()
-    else [
-        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
        "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
        "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
        "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
@@ -1,183 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Worker-side unit tests for SimpleCPUOffloadConnector.
-
-Covers the GPU->CPU store cross-stream synchronization: the store copy must be
-ordered after the compute stream that writes the KV blocks, otherwise it can
-read partially written / stale blocks and silently corrupt the CPU cache.
-"""
-
-from __future__ import annotations
-
-import time
-
-import pytest
-import torch
-
-from vllm.platforms import current_platform
-
-if not current_platform.is_cuda_alike():
-    pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
-
-from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
-from vllm.v1.simple_kv_offload.cuda_mem_ops import (
-    CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
-    CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
-    build_params,
-    pin_tensor,
-)
-from vllm.v1.simple_kv_offload.metadata import SimpleCPUOffloadMetadata
-from vllm.v1.simple_kv_offload.worker import SimpleCPUOffloadWorker
-
-NUM_BLOCKS = 64
-BLOCK_BYTES = 4096
-ITERS = 30
-# Keep the compute stream busy so the KV write lands late; this makes the
-# store-vs-compute race deterministic instead of timing-dependent.
-SLEEP_CYCLES = 50_000_000
-
-
-def _make_backend() -> tuple[DmaCopyBackend, torch.Tensor, torch.Tensor]:
-    gpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cuda")}
-    cpu = {"k": torch.zeros((NUM_BLOCKS, BLOCK_BYTES), dtype=torch.int8, device="cpu")}
-    pin_tensor(cpu["k"])
-    low_pri, _ = torch.cuda.Stream.priority_range()
-    backend = DmaCopyBackend()
-    backend.init(
-        gpu,
-        cpu,
-        gpu["k"].device,
-        torch.cuda.Stream(priority=low_pri),
-        torch.cuda.Stream(priority=low_pri),
-    )
-    return backend, gpu["k"], cpu["k"]
-
-
-def _drive_store(
-    backend: DmaCopyBackend,
-    gpu: torch.Tensor,
-    cpu: torch.Tensor,
-    *,
-    with_barrier: bool,
-) -> int:
-    """Run ITERS store cycles; return how many landed corrupted in the CPU pool.
-
-    Each cycle writes a unique value on a compute stream (after a deliberate
-    delay) and then issues the GPU->CPU store. The store is issued *after* the
-    write in host program order, mirroring the connector's deferred-store
-    assumption. Only the compute-done event creates a real device-side
-    happens-before edge.
-    """
-    block_ids = list(range(gpu.shape[0]))
-    compute_stream = torch.cuda.Stream()
-    corrupt = 0
-    for it in range(ITERS):
-        val = (it % 126) + 1  # 1..126; distinct from the zero-initialized pool
-        with torch.cuda.stream(compute_stream):
-            torch.cuda._sleep(SLEEP_CYCLES)
-            gpu.fill_(val)
-
-        wait_event = None
-        if with_barrier:
-            wait_event = torch.Event()
-            wait_event.record(compute_stream)
-
-        store_events: list[tuple[int, torch.Event]] = []
-        backend.launch_copy(
-            block_ids,
-            block_ids,
-            is_store=True,
-            event_idx=it,
-            events_list=store_events,
-            wait_event=wait_event,
-        )
-
-        deadline = time.time() + 10.0
-        while not store_events and time.time() < deadline:
-            time.sleep(0.0005)
-        assert store_events, "background copy was never enqueued"
-        store_events[0][1].synchronize()
-
-        if int((cpu[:, 0].to(torch.int32) != val).sum().item()):
-            corrupt += 1
-    return corrupt
-
-
-def test_store_orders_after_compute_write():
-    """The store must wait for the compute event; without it, it races.
-
-    Asserts both directions so the test is self-validating: the no-barrier
-    control must actually corrupt (proving the race window is exercised), and
-    the fixed path with the compute-done event must be clean.
-    """
-    backend, gpu, cpu = _make_backend()
-    try:
-        control = _drive_store(backend, gpu, cpu, with_barrier=False)
-        fixed = _drive_store(backend, gpu, cpu, with_barrier=True)
-    finally:
-        backend.shutdown()
-
-    assert control > 0, (
-        "no-barrier store did not race the compute write; the test no longer "
-        "exercises the hazard it is meant to guard"
-    )
-    assert fixed == 0, f"store raced compute even with the barrier: {fixed} corrupt"
-
-
-class _RecordingBackend:
-    """Captures launch_copy calls without touching the GPU."""
-
-    def __init__(self) -> None:
-        self.calls: list[dict] = []
-
-    def launch_copy(
-        self,
-        src_blocks,
-        dst_blocks,
-        is_store,
-        event_idx,
-        events_list,
-        wait_event=None,
-    ) -> None:
-        self.calls.append({"is_store": is_store, "wait_event": wait_event})
-
-
-def test_get_finished_passes_wait_event_for_store_only():
-    """get_finished gates stores on a compute-done event but not loads."""
-    worker = SimpleCPUOffloadWorker(
-        vllm_config=None, kv_cache_config=None, cpu_capacity_bytes=0
-    )
-    recording = _RecordingBackend()
-    worker._backend = recording
-    worker._connector_metadata = SimpleCPUOffloadMetadata(
-        load_event=0,
-        load_gpu_blocks=[0],
-        load_cpu_blocks=[0],
-        store_event=1,
-        store_gpu_blocks=[1],
-        store_cpu_blocks=[1],
-    )
-
-    worker.get_finished(set())
-
-    store_calls = [c for c in recording.calls if c["is_store"]]
-    load_calls = [c for c in recording.calls if not c["is_store"]]
-    assert len(store_calls) == 1
-    assert len(load_calls) == 1
-    assert isinstance(store_calls[0]["wait_event"], torch.Event)
-    assert load_calls[0]["wait_event"] is None
-
-
-def test_build_params_src_access_order():
-    """build_params defaults to ANY and honors an explicit STREAM override."""
-    gpu = {"k": torch.zeros((4, 64), dtype=torch.int8, device="cuda")}
-    cpu = {"k": torch.zeros((4, 64), dtype=torch.int8, device="cpu")}
-    stream = torch.cuda.Stream()
-
-    default = build_params(gpu, cpu, stream)
-    assert default.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_ANY
-
-    ordered = build_params(
-        gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
-    )
-    assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
@@ -132,9 +132,9 @@ def get_available_attention_backends() -> list[str]:
    )

    return [
-        candidate.backend.name
-        for candidate in valid_backends
-        if candidate.backend not in EXCLUDED_BACKENDS
+        backend.name
+        for backend, _ in valid_backends
+        if backend not in EXCLUDED_BACKENDS
    ]


@@ -1,157 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for EncoderRunner.gather_mm_embeddings (model runner V2).
-
-Covers the speculative-drafter encoder-cache handling: the drafter reads one
-position ahead of the target model (``draft_lookahead``). The +1 look-ahead
-feature past the processed boundary is used when its encoder output is present
-and tolerated (token-embedding fallback) when it is not, while a miss within
-the processed range still fails loudly.
-"""
-
-import numpy as np
-import pytest
-import torch
-
-from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
-from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
-from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
-
-pytestmark = pytest.mark.cpu_test
-
-HIDDEN = 4
-
-
-def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
-    return MultiModalFeatureSpec(
-        data=None,
-        modality="image",
-        identifier=identifier,
-        mm_position=PlaceholderRange(offset=offset, length=length),
-    )
-
-
-def _make_runner(
-    features: list[MultiModalFeatureSpec],
-    cached: list[MultiModalFeatureSpec],
-) -> EncoderRunner:
-    cache = EncoderCache()
-    cache.mm_features["req0"] = features
-    for f in cached:
-        length = f.mm_position.length
-        cache.encoder_outputs[f.identifier] = torch.arange(
-            length * HIDDEN, dtype=torch.float32
-        ).reshape(length, HIDDEN)
-    return EncoderRunner(
-        model=None,  # unused by gather_mm_embeddings
-        max_num_tokens=64,
-        hidden_size=HIDDEN,
-        encoder_cache=cache,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-    )
-
-
-def _gather(runner: EncoderRunner, *, num_scheduled: int, draft_lookahead: int):
-    # Single prefilling request, computed_prefill=0, prefill_len large.
-    return runner.gather_mm_embeddings(
-        req_ids=["req0"],
-        total_num_scheduled_tokens=num_scheduled,
-        num_scheduled_tokens=np.array([num_scheduled]),
-        query_start_loc=np.array([0]),
-        prefill_lens=np.array([1000]),
-        computed_prefill_lens=np.array([0]),
-        draft_lookahead=draft_lookahead,
-    )
-
-
-def test_draft_lookahead_uses_boundary_feature_when_cached():
-    """The drafter's +1 look-ahead can reach the feature at offset ==
-    processed_end (the next chunk). When its encoder output is already cached
-    (the scheduler encoded it ahead), it is used for the look-ahead position
-    rather than ignored."""
-    f0 = _feature("h0", offset=0, length=8)
-    f1 = _feature("h1", offset=8, length=8)  # starts exactly at processed_end
-    runner = _make_runner([f0, f1], cached=[f0, f1])
-
-    mm_embeds, is_mm_embed = _gather(runner, num_scheduled=8, draft_lookahead=1)
-
-    # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7.
-    assert len(mm_embeds) == 2
-    assert bool(is_mm_embed[7])
-    assert int(is_mm_embed.sum()) == 8
-
-
-def test_draft_lookahead_tolerates_missing_boundary_feature():
-    """When the +1 look-ahead feature past the processed boundary is not yet
-    encoded, fall back to the token embedding (the draft token is verified by
-    the target) instead of raising."""
-    f0 = _feature("h0", offset=0, length=8)
-    f1 = _feature("h1", offset=8, length=8)  # boundary feature, not cached
-    runner = _make_runner([f0, f1], cached=[f0])
-
-    mm_embeds, is_mm_embed = _gather(runner, num_scheduled=8, draft_lookahead=1)
-
-    # Only f0 is gathered; f1's boundary position falls back silently.
-    assert len(mm_embeds) == 1
-    assert not bool(is_mm_embed[7])
-    assert int(is_mm_embed.sum()) == 7
-
-
-def test_draft_lookahead_raises_on_interior_miss():
-    """A miss for a feature within the processed range (not the look-ahead
-    boundary) is a real invariant violation and must fail loudly, even on the
-    drafter path."""
-    f0 = _feature("h0", offset=0, length=8)  # interior, within processed range
-    runner = _make_runner([f0], cached=[])
-
-    with pytest.raises(RuntimeError, match="Encoder cache miss"):
-        _gather(runner, num_scheduled=8, draft_lookahead=1)
-
-
-def test_target_path_raises_on_encoder_cache_miss():
-    """On the target path (no look-ahead) a miss is a real invariant
-    violation and must fail loudly."""
-    f0 = _feature("h0", offset=0, length=8)
-    runner = _make_runner([f0], cached=[])
-
-    with pytest.raises(RuntimeError, match="Encoder cache miss"):
-        _gather(runner, num_scheduled=8, draft_lookahead=0)
-
-
-@pytest.mark.parametrize("draft_lookahead", [0, 1])
-def test_multi_request_batch_gathers_per_request(draft_lookahead):
-    """Two prefilling requests in one batch: per-request query bounds must be
-    indexed by request, not applied as whole arrays."""
-    a0 = _feature("a0", offset=0, length=8)
-    b0 = _feature("b0", offset=0, length=8)
-    cache = EncoderCache()
-    cache.mm_features["req0"] = [a0]
-    cache.mm_features["req1"] = [b0]
-    for f in (a0, b0):
-        cache.encoder_outputs[f.identifier] = torch.arange(
-            f.mm_position.length * HIDDEN, dtype=torch.float32
-        ).reshape(f.mm_position.length, HIDDEN)
-    runner = EncoderRunner(
-        model=None,
-        max_num_tokens=64,
-        hidden_size=HIDDEN,
-        encoder_cache=cache,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-    )
-
-    mm_embeds, is_mm_embed = runner.gather_mm_embeddings(
-        req_ids=["req0", "req1"],
-        total_num_scheduled_tokens=16,
-        num_scheduled_tokens=np.array([8, 8]),
-        query_start_loc=np.array([0, 8]),
-        prefill_lens=np.array([1000, 1000]),
-        computed_prefill_lens=np.array([0, 0]),
-        draft_lookahead=draft_lookahead,
-    )
-
-    # Both requests contribute a feature; with the +1 skew each marks 7 of its
-    # 8 positions (the skew drops one), otherwise all 8.
-    assert len(mm_embeds) == 2
-    assert int(is_mm_embed.sum()) == (14 if draft_lookahead else 16)
@@ -1,99 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for GPUModelRunner._gather_mm_embeddings (model runner V1).
-
-Mirrors tests/v1/worker/test_encoder_runner.py (the V2 runner): the EAGLE/MTP
-drafter reads one position ahead of the target (shift_computed_tokens=1). The
-+1 look-ahead feature past the processed boundary is used when its encoder
-output is present and tolerated (token-embedding fallback) when it is not,
-while a miss within the processed range still fails loudly.
-
-`_gather_mm_embeddings` only uses CPU-side state, so it is exercised against a
-lightweight stub for `self` instead of a full (CUDA-only) runner.
-"""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
-from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-
-pytestmark = pytest.mark.cpu_test
-
-HIDDEN = 4
-
-
-def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
-    return MultiModalFeatureSpec(
-        data=None,
-        modality="image",
-        identifier=identifier,
-        mm_position=PlaceholderRange(offset=offset, length=length),
-    )
-
-
-def _gather(features, cached, *, num_scheduled, shift, num_computed=0):
-    encoder_cache = {
-        f.identifier: torch.arange(
-            f.mm_position.length * HIDDEN, dtype=torch.float32
-        ).reshape(f.mm_position.length, HIDDEN)
-        for f in cached
-    }
-    req_state = SimpleNamespace(num_computed_tokens=num_computed, mm_features=features)
-    runner = SimpleNamespace(
-        input_batch=SimpleNamespace(req_ids=["req0"]),
-        requests={"req0": req_state},
-        encoder_cache=encoder_cache,
-        is_multimodal_pruning_enabled=False,
-        uses_mrope=False,
-    )
-    scheduler_output = SimpleNamespace(
-        total_num_scheduled_tokens=num_scheduled,
-        num_scheduled_tokens={"req0": num_scheduled},
-    )
-    return GPUModelRunner._gather_mm_embeddings(
-        runner, scheduler_output, shift_computed_tokens=shift
-    )
-
-
-def test_draft_shift_uses_boundary_feature_when_cached():
-    """The drafter's +1 look-ahead reaches the feature at offset ==
-    processed_end; when it is already cached it is used for the look-ahead
-    position rather than ignored."""
-    f0 = _feature("h0", offset=0, length=8)
-    f1 = _feature("h1", offset=8, length=8)  # starts exactly at processed_end
-    mm_embeds, is_mm_embed = _gather([f0, f1], [f0, f1], num_scheduled=8, shift=1)
-
-    # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7.
-    assert len(mm_embeds) == 2
-    assert bool(is_mm_embed[7])
-    assert int(is_mm_embed.sum()) == 8
-
-
-def test_draft_shift_tolerates_missing_boundary_feature():
-    """When the +1 look-ahead feature past the processed boundary is not yet
-    encoded, fall back to the token embedding instead of raising."""
-    f0 = _feature("h0", offset=0, length=8)
-    f1 = _feature("h1", offset=8, length=8)  # boundary feature, not cached
-    mm_embeds, is_mm_embed = _gather([f0, f1], [f0], num_scheduled=8, shift=1)
-
-    assert len(mm_embeds) == 1  # only f0; f1's boundary position falls back
-    assert not bool(is_mm_embed[7])
-    assert int(is_mm_embed.sum()) == 7
-
-
-def test_draft_shift_raises_on_interior_miss():
-    """A miss for a feature within the processed range (not the look-ahead
-    boundary) is a real invariant violation, even on the drafter path."""
-    f0 = _feature("h0", offset=0, length=8)  # interior, within processed range
-    with pytest.raises(RuntimeError, match="Encoder cache miss"):
-        _gather([f0], [], num_scheduled=8, shift=1)
-
-
-def test_target_path_raises_on_encoder_cache_miss():
-    """On the target path (no shift) a miss is a real invariant violation."""
-    f0 = _feature("h0", offset=0, length=8)
-    with pytest.raises(RuntimeError, match="Encoder cache miss"):
-        _gather([f0], [], num_scheduled=8, shift=0)
@@ -690,9 +690,7 @@ def parse_compute_capability(node: ast.ClassDef) -> str:
        major_list.sort()
        if len(major_list) == 1:
            return f"{major_list[0]}.x"
-        if major_list == list(range(major_list[0], major_list[-1] + 1)):
-            return f"{major_list[0]}.x-{major_list[-1]}.x"
-        return ", ".join(f"{major}.x" for major in major_list)
+        return f"{major_list[0]}.x-{major_list[-1]}.x"

    if min_cap:
        if max_cap:
@@ -1670,8 +1668,7 @@ def generate_mla_section(
                "`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,",
                "`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index",
                "pipeline (compressor + SWA + indexer, 256-token blocks, head 512);",
-                "default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and",
-                "`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.",
+                "default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.",
                "",
            ]
        )
@@ -956,15 +956,6 @@ class AsyncTPPass(VllmFusionPatternMatcherPass):
                                a_scale_view=a_scale_view,
                            )
                        )
-                self.register(
-                    FlashInferAllGatherFP4Pattern(
-                        self.model_dtype,
-                        self.device,
-                        "cute-dsl",
-                        use_8x4_sf_layout=False,
-                        a_scale_view="float8",
-                    )
-                )
                # NVFP4 reduce-scatter does not need scale communication: FP4
                # scales are consumed by the local GEMM and only BF16 partial
                # outputs are reduced. Keep this PR scoped to the all-gather
@@ -141,7 +141,6 @@ LinearBackend = Literal[
    "auto",
    "cutlass",
    "flashinfer_cutlass",
-    "flashinfer_cutedsl",
    "flashinfer_trtllm",
    "flashinfer_cudnn",
    "flashinfer_b12x",
@@ -199,7 +198,6 @@ class KernelConfig:
    - "auto": Automatically select the best backend based on model and hardware
    - "cutlass": Use CUTLASS-based kernels
    - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels
-    - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels
    - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
    - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
    - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)
@@ -803,6 +803,13 @@ class ParallelConfig:
        if self.enable_elastic_ep:
            if not self.enable_eplb:
                raise ValueError("Elastic EP is only supported with enable_eplb=True.")
+            if self.eplb_config.use_async:
+                raise ValueError(
+                    "Elastic EP requires the pynccl communicator, which is "
+                    "incompatible with async EPLB due to NCCL multi-stream "
+                    "conflicts. Disable async EPLB (eplb_config.use_async=False) "
+                    "to use elastic EP."
+                )
            if self.pipeline_parallel_size > 1:
                raise ValueError(
                    "Elastic EP is not supported with pipeline parallelism "
@@ -814,15 +821,6 @@ class ParallelConfig:
                    "or data_parallel_hybrid_lb. Elastic EP relies on a single API "
                    "server and core client to coordinate scale up/down."
                )
-            if self.eplb_config.use_async:
-                from vllm.distributed.nixl_utils import is_nixl_available
-
-                if not is_nixl_available():
-                    raise ValueError(
-                        "Elastic EP with async EPLB requires the NIXL "
-                        "package. Either install NIXL or set "
-                        "--eplb-config.use_async=false."
-                    )

        if self.data_parallel_size > 1 or self.data_parallel_size_local == 0:
            # Data parallel was specified in the engine args.
@@ -931,21 +929,23 @@ class ParallelConfig:
            )

        if self.enable_eplb and self.eplb_config.communicator is None:
-            # Prefer NIXL when available: zero-copy RDMA reads, compatible
-            # with both async EPLB and elastic EP (deferred remote setup).
-            # Fallbacks: pynccl for elastic EP (stateless groups need it),
-            # torch_gloo for static EP.  torch_nccl is avoided because NCCL
-            # is incompatible with async EPLB (multi-stream conflicts) and
-            # batched isend/irecv hangs under high load.
-            # See https://github.com/pytorch/pytorch/issues/174288
-            from vllm.distributed.nixl_utils import is_nixl_available
-
-            if is_nixl_available():
-                self.eplb_config.communicator = "nixl"
-            elif self.enable_elastic_ep:
+            if self.enable_elastic_ep:
+                # Elastic EP requires stateless mode
+                # (torch.distributed.batch_isend_irecv doesn't
+                # support stateless mode), so we use PyNCCL backend
                self.eplb_config.communicator = "pynccl"
            else:
-                self.eplb_config.communicator = "torch_gloo"
+                # Avoid torch_nccl: NCCL is fundamentally incompatible
+                # with async EPLB due to multi-stream conflicts, and
+                # batched isend/irecv hangs under high load.
+                # See https://github.com/pytorch/pytorch/issues/174288
+                # Prefer nixl when available; fall back to torch_gloo.
+                from vllm.distributed.nixl_utils import is_nixl_available
+
+                if is_nixl_available():
+                    self.eplb_config.communicator = "nixl"
+                else:
+                    self.eplb_config.communicator = "torch_gloo"

    @property
    def use_ray(self) -> bool:
@@ -207,9 +207,6 @@ class ElasticEPScalingExecutor:
            )
        if new_dp_size > old_dp_size:
            self._set_eplb_suppressed(True)
-            eplb_state = self.worker.model_runner.eplb_state
-            if eplb_state is not None:
-                eplb_state.drain_async()
        elif new_dp_size < old_dp_size:
            self._stage_standby_moe_quant_methods()

@@ -543,11 +540,6 @@ class ElasticEPScalingExecutor:
            eplb_model_state.physical_to_logical_map.shape[1]
        )
        eplb_state.is_async = is_async_enabled
-        # Start the async worker thread if it doesn't exist yet (idempotent).
-        # This is needed for new workers after scale-up: they create EplbState
-        # in setup_eplb_from_mapping() but don't start the thread there because
-        # groups aren't ready yet.
-        eplb_state.start_async_loop()
        if get_ep_group().rank == 0:
            logger.info("[Elastic EP] Expert resharding completed")

@@ -557,9 +549,6 @@ class ElasticEPScalingExecutor:

    def perform_scale_down_eplb_reshuffle(self, new_dp_size: int) -> None:
        self._set_eplb_suppressed(True)
-        eplb_state = self.worker.model_runner.eplb_state
-        if eplb_state is not None:
-            eplb_state.drain_async()
        parallel_config = self.worker.vllm_config.parallel_config
        tp_size = parallel_config.tensor_parallel_size
        old_ep_size = parallel_config.data_parallel_size * tp_size
@@ -8,6 +8,7 @@ import threading
 from typing import TYPE_CHECKING

 import torch
+from torch.distributed import ProcessGroup

 from vllm.distributed.parallel_state import get_eplb_group
 from vllm.logger import init_logger
@@ -25,7 +26,8 @@ def start_async_worker(
    state: "EplbState",
    is_profile: bool = False,
 ) -> threading.Thread:
-    rank = get_eplb_group().device_group.rank()
+    eplb_group = get_eplb_group().device_group
+    rank = eplb_group.rank()
    device_index = state.cuda_device_index
    assert state.is_async

@@ -36,6 +38,7 @@ def start_async_worker(
        try:
            transfer_run_periodically(
                state=state,
+                eplb_group=eplb_group,
                cuda_stream=cuda_stream,
                is_profile=is_profile,
            )
@@ -75,15 +78,13 @@ def run_rebalance_experts(

 def transfer_run_periodically(
    state: "EplbState",
+    eplb_group: ProcessGroup,
    cuda_stream: torch.cuda.Stream,
    is_profile: bool = False,
 ) -> None:
    while True:
        state.rearrange_event.wait(stream=cuda_stream)
-
-        eplb_group = get_eplb_group().device_group
-        eplb_cpu_group = get_eplb_group().cpu_group
-        ep_rank = eplb_group.rank()
+        logger.info("async worker woke up for EPLB transfer")

        assert state.is_async
        for model_state in state.model_states.values():
@@ -100,32 +101,16 @@ def transfer_run_periodically(
            new_physical_to_logical_map = run_rebalance_experts(
                model_state, state, physical_to_logical_map_cpu, cuda_stream
            )
+            logger.info(
+                "Async worker computed new indices for model %s",
+                model_state.model_name,
+            )

            # Execute one EPLB layer transfer per model forward pass. Each iteration
            # of this loop will copy the new set of expert weights into
            # model_state.expert_buffer, which will be consumed by the main thread in
-            # move_to_workspace.
-            # We sync the rebalanced flag across ranks before each iteration so
-            # all ranks make a coordinated decision to continue or stop.
-            while layer_idx < num_layers:
-                flag = torch.tensor(
-                    [int(model_state.rebalanced)],
-                    dtype=torch.int32,
-                    device="cpu",
-                )
-                torch.distributed.all_reduce(flag, group=eplb_cpu_group)
-                if int(flag.item()) != eplb_cpu_group.size():
-                    logger.warning(
-                        "async worker (rank=%d): layer %d coordinated stop "
-                        "(flag_sum=%d, group_size=%d)",
-                        ep_rank,
-                        layer_idx,
-                        int(flag.item()),
-                        eplb_cpu_group.size(),
-                    )
-                    model_state.rebalanced = False
-                    break
-
+            # move_to_workspace
+            while model_state.rebalanced and layer_idx < num_layers:
                transfer_metadata = transfer_layer(
                    old_layer_indices=physical_to_logical_map_cpu[layer_idx],
                    new_layer_indices=new_physical_to_logical_map[layer_idx],
@@ -158,5 +143,6 @@ def transfer_run_periodically(
                # finish copying model_state.expert_buffer into
                # model_state.model.expert_weights[layer_idx]
                consumed_event.wait(stream=cuda_stream)
+                logger.debug("Layer %d transfer complete", layer_idx)
                assert model_state.pending_result is None
                layer_idx += 1
@@ -246,20 +246,7 @@ class NixlEplbCommunicator(EplbCommunicator):
        cpu_group: ProcessGroup,
        all_expert_weights: Sequence[Sequence[torch.Tensor]],
        expert_buffer: Sequence[torch.Tensor],
-        defer_remote_setup: bool = False,
    ) -> None:
-        """Create a NIXL-backed EPLB communicator.
-
-        Args:
-            cpu_group: CPU process group for metadata exchange.
-            all_expert_weights: Expert weight tensors for all MoE layers.
-            expert_buffer: Pre-allocated receive buffer tensors.
-            defer_remote_setup: If True, postpone the collective
-                all-gather of NIXL agent metadata until the first
-                ``set_transfer_context`` call.  Required for elastic EP
-                where ranks join asynchronously and cannot participate
-                in collectives at construction time.
-        """
        assert all_expert_weights, (
            "NixlEplbCommunicator requires non-empty all_expert_weights."
        )
@@ -315,29 +302,10 @@ class NixlEplbCommunicator(EplbCommunicator):
        ] = {}

        self._cuda_device_id = int(self._device.index or 0)
-        self._remote_state_initialized = False
        self._init_step("buffers", self._init_registered_buffers)
-        if defer_remote_setup:
-            logger.info_once("NIXL EPLB: deferring remote agent setup (elastic EP).")
-        else:
-            self._init_remote_state()
-        self._log_initialized()
-
-    def _init_remote_state(self) -> None:
-        """Exchange NIXL agent metadata and RDMA pointer info with all peers.
-
-        This is a collective operation (uses ``all_gather_object`` twice).
-        Under elastic EP the call is deferred to the first
-        ``set_transfer_context`` invocation, where all ranks are
-        guaranteed to be synchronized.
-        """
        self._init_step("agents", self._init_remote_agents)
        self._init_step("send meta", self._exchange_remote_send_meta)
-        self._remote_state_initialized = True
-
-    def _ensure_remote_state(self) -> None:
-        if not self._remote_state_initialized:
-            self._init_remote_state()
+        self._log_initialized()

    @property
    def needs_profile_buffer_reservation(self) -> bool:
@@ -371,7 +339,8 @@ class NixlEplbCommunicator(EplbCommunicator):
        pass

    def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
-        self._ensure_remote_state()
+        # Pre-compute expert_id -> src_row mapping for every rank so that
+        # add_recv can immediately issue NIXL READs.
        assert not self._xfer_entries, (
            f"set_transfer_context() called with {len(self._xfer_entries)} "
            f"pending transfers from layer {self._layer_idx}; "
@@ -554,21 +523,6 @@ class NixlEplbCommunicator(EplbCommunicator):
        )
        return (local_handle, remote_handle, xfer_handle)

-    def _post_read_barrier(self) -> None:
-        """Correctness fence: prevents overwrite-while-remote-read race.
-
-        We avoid ``torch.distributed.monitored_barrier`` because it
-        calls ``get_backend(group)`` which fails for stateless groups
-        (elastic EP).  An async ``all_reduce`` + ``wait(timeout)``
-        works with both regular and stateless groups and provides
-        equivalent timeout detection.
-        """
-        _dummy = torch.zeros(1, dtype=torch.int32)
-        work = torch.distributed.all_reduce(
-            _dummy, group=self._cpu_group, async_op=True
-        )
-        work.wait(timeout=timedelta(minutes=5))
-
    def execute(self) -> None:
        assert self._layer_idx is not None or not self._xfer_entries, (
            "set_transfer_context() must be called before execute() "
@@ -577,7 +531,13 @@ class NixlEplbCommunicator(EplbCommunicator):
        try:
            self._wait_for_all_transfers([x[2] for x in self._xfer_entries])

-            self._post_read_barrier()
+            # Post-READ barrier.
+            # Correctness fence for zero-copy: prevents overwrite-while-
+            # remote-read race.
+            torch.distributed.monitored_barrier(
+                group=self._cpu_group,
+                timeout=timedelta(minutes=5),
+            )
        finally:
            for local_h, remote_h, xfer_h in self._xfer_entries:
                with contextlib.suppress(Exception):
@@ -668,13 +628,11 @@ def create_eplb_communicator(
            device and CPU communication groups.
        backend: Communicator backend name (``"torch_nccl"``,
            ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).
-            Falls back to ``"torch_nccl"`` when *None*.
-            Stateless (elastic EP) groups support ``"torch_nccl"``,
-            ``"pynccl"``, and ``"nixl"``; ``"torch_nccl"`` is silently
-            promoted to ``"pynccl"``.  ``"nixl"`` uses deferred remote
-            agent setup to avoid collective deadlocks during elastic
-            scaling.  When tensors reside on CPU, ``"torch_gloo"`` or
-            ``"torch_nccl"`` are used via the CPU process group.
+            Stateless (elastic EP) groups only support ``"torch_nccl"``
+            and ``"pynccl"``; ``"torch_nccl"`` is silently promoted to
+            ``"pynccl"`` in that case.  When tensors reside on CPU,
+            ``"torch_gloo"`` or ``"torch_nccl"`` are used via the CPU
+            process group.
        expert_weights: Expert weight tensors for *all* MoE layers.
            Shape ``(num_layers)(num_tensors_per_layer)``.
            NixlEplbCommunicator registers all layers with NIXL for
@@ -728,21 +686,18 @@ def create_eplb_communicator(

    is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)
    if is_stateless:
-        if backend == "nixl":
-            pass  # handled below with defer_remote_setup=True
-        elif backend not in ("torch_nccl", "pynccl"):
+        if backend not in ("torch_nccl", "pynccl"):
            raise ValueError(
-                f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' "
-                f"EPLB communicator (got '{backend}')."
+                f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator "
+                f"(got '{backend}')."
            )
-        else:
-            if backend == "torch_nccl":
-                logger.warning(
-                    "Stateless elastic EP requires PyNCCL backend. "
-                    "Forcing EPLB communicator to 'pynccl'."
-                )
-                backend = "pynccl"
-            return _create_pynccl()
+        if backend == "torch_nccl":
+            logger.warning(
+                "Stateless elastic EP requires PyNCCL backend. "
+                "Forcing EPLB communicator to 'pynccl'."
+            )
+            backend = "pynccl"
+        return _create_pynccl()

    if backend == "nixl":
        if not has_nixl():
@@ -759,7 +714,6 @@ def create_eplb_communicator(
                cpu_group=group_coordinator.cpu_group,
                all_expert_weights=expert_weights,
                expert_buffer=expert_buffer,
-                defer_remote_setup=is_stateless,
            )
        except Exception as exc:
            raise RuntimeError(
@@ -27,7 +27,6 @@ physical experts.
 """

 import threading
-import time
 from collections.abc import Sequence
 from dataclasses import dataclass

@@ -826,45 +825,6 @@ class EplbState:
                is_profile=is_profile,
            )

-    def drain_async(self) -> None:
-        """Drain in-flight async EPLB by consuming all remaining layer results.
-
-        Each pending result is acknowledged (consumed_event recorded) so the
-        async worker can proceed, but the transferred weights are intentionally
-        NOT applied — a full synchronous rearrange is expected to follow.
-
-        Ranks are kept in lockstep via _all_ranks_result_ready (all_reduce
-        on the EP CPU group).  The async worker's coordinated-stop collectives
-        use the separate EPLB group, so the two sets of collectives do not
-        interfere.
-
-        No-op when no async cycle is in progress (rebalanced=False).
-        """
-        if not self.is_async:
-            return
-        for model_key, ms in self.model_states.items():
-            needs_drain = ms.rebalanced
-            if needs_drain:
-                logger.info(
-                    "Draining async EPLB worker for model %s",
-                    model_key,
-                )
-            while ms.rebalanced:
-                if self._all_ranks_result_ready(ms):
-                    result = ms.pending_result
-                    assert result is not None
-                    if result.layer_idx == ms.model.num_moe_layers - 1:
-                        ms.rebalanced = False
-                    ms.pending_result = None
-                    result.consumed_event.record()
-                else:
-                    time.sleep(0.001)
-            if needs_drain:
-                logger.info(
-                    "Async EPLB worker drained for model %s",
-                    model_key,
-                )
-
    def _all_ranks_result_ready(self, model_state: EplbModelState) -> bool:
        parallel_state = get_ep_group()
        has_result = int(model_state.pending_result is not None)
@@ -890,9 +850,8 @@ class EplbState:
        """
        All-reduce a list of tensors.
        """
-        ep_group = get_ep_group().device_group
        if len(tensor_list) == 1:
-            all_reduce(tensor_list[0], group=ep_group)
+            all_reduce(tensor_list[0], group=get_ep_group().device_group)
            return tensor_list
        assert all(t.dim() == 2 for t in tensor_list), "All tensors must be 2D."
        assert all(t.shape[1] == tensor_list[0].shape[1] for t in tensor_list), (
@@ -904,6 +863,7 @@ class EplbState:
        shapes = [t.shape for t in tensor_list]
        concat_tensor = torch.cat(tensor_list, dim=0)

+        ep_group = get_ep_group().device_group
        all_reduce(concat_tensor, group=ep_group)

        all_reduce_list = []
@@ -387,29 +387,41 @@ class DecodeBenchConnectorWorker:

            kv_cache = self.kv_caches[layer_name]

-            # Attention layers store KV as a single block-indexed tensor whose
-            # first dim is num_blocks; fill the requested block rows. Hybrid /
-            # linear-attention layers (e.g. Mamba, Kimi Delta Attention) store
-            # their state as a list/tuple of tensors that are NOT block-indexed
-            # — each tensor is a single state buffer with no num_blocks
-            # dimension — so fill each tensor in its entirety with the same
-            # dummy values.
-            if isinstance(kv_cache, torch.Tensor):
-                self._fill_block_tensor(kv_cache, block_ids)
-            elif isinstance(kv_cache, (list, tuple)) and all(
-                isinstance(t, torch.Tensor) for t in kv_cache
-            ):
-                for state_tensor in kv_cache:
-                    self._fill_state_tensor(state_tensor)
-            else:
-                logger.warning_once(
-                    "DecodeBenchConnector: skipping fill for layer %s whose KV "
-                    "cache is %s, not a tensor or a list/tuple of tensors.",
-                    layer_name,
-                    type(kv_cache).__name__,
-                )
+            # Convert block_ids to tensor on device
+            block_ids_tensor = torch.tensor(
+                block_ids, dtype=torch.long, device=kv_cache.device
+            )
+
+            # Filter invalid block IDs
+            valid_mask = block_ids_tensor < kv_cache.shape[0]
+            valid_block_ids = block_ids_tensor[valid_mask]
+
+            if len(valid_block_ids) == 0:
                continue

+            # Create fill values - either constant or random
+            block_shape = kv_cache.shape[1:]
+            if self.fill_std > 0:
+                # Random normal sampling
+                fill_values = torch.normal(
+                    mean=self.fill_mean,
+                    std=self.fill_std,
+                    size=(len(valid_block_ids),) + block_shape,
+                    dtype=kv_cache.dtype,
+                    device=kv_cache.device,
+                )
+            else:
+                # Constant fill value
+                fill_values = torch.full(
+                    (len(valid_block_ids),) + block_shape,
+                    self.fill_mean,
+                    dtype=kv_cache.dtype,
+                    device=kv_cache.device,
+                )
+
+            # Batch fill operation
+            kv_cache[valid_block_ids] = fill_values
+
        logger.debug(
            "DecodeBenchConnector: Filled %d blocks in group %d with %s values "
            "(mean=%.3f, std=%.3f)",
@@ -419,62 +431,3 @@ class DecodeBenchConnectorWorker:
            self.fill_mean,
            self.fill_std,
        )
-
-    def _fill_block_tensor(self, kv_cache: torch.Tensor, block_ids: list[int]):
-        """Fill the requested block rows of a block-indexed KV cache tensor.
-
-        Args:
-            kv_cache: A KV cache tensor whose first dim is num_blocks.
-            block_ids: Block IDs to fill. IDs that are out of range for this
-                tensor's first dim are ignored.
-        """
-        # Convert block_ids to tensor on device
-        block_ids_tensor = torch.tensor(
-            block_ids, dtype=torch.long, device=kv_cache.device
-        )
-
-        # Filter invalid block IDs
-        valid_mask = block_ids_tensor < kv_cache.shape[0]
-        valid_block_ids = block_ids_tensor[valid_mask]
-
-        if len(valid_block_ids) == 0:
-            return
-
-        # Create fill values - either constant or random
-        block_shape = kv_cache.shape[1:]
-        if self.fill_std > 0:
-            # Random normal sampling
-            fill_values = torch.normal(
-                mean=self.fill_mean,
-                std=self.fill_std,
-                size=(len(valid_block_ids),) + block_shape,
-                dtype=kv_cache.dtype,
-                device=kv_cache.device,
-            )
-        else:
-            # Constant fill value
-            fill_values = torch.full(
-                (len(valid_block_ids),) + block_shape,
-                self.fill_mean,
-                dtype=kv_cache.dtype,
-                device=kv_cache.device,
-            )
-
-        # Batch fill operation
-        kv_cache[valid_block_ids] = fill_values
-
-    def _fill_state_tensor(self, kv_cache: torch.Tensor):
-        """Fill an entire non-block-indexed state tensor with dummy values.
-
-        Hybrid / linear-attention layers (e.g. Mamba, Kimi Delta Attention)
-        store their per-layer state as tensors with no num_blocks dimension,
-        so the whole tensor is filled with the same constant or random values
-        used for block fills, rather than selected block rows.
-
-        Args:
-            kv_cache: A state tensor to fill in its entirety.
-        """
-        if self.fill_std > 0:
-            kv_cache.normal_(mean=self.fill_mean, std=self.fill_std)
-        else:
-            kv_cache.fill_(self.fill_mean)
@@ -115,7 +115,6 @@ from vllm.model_executor.kernels.linear.nvfp4.fbgemm import (
 from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
    FlashInferB12xNvFp4LinearKernel,
    FlashInferCudnnNvFp4LinearKernel,
-    FlashInferCuteDslNvFp4LinearKernel,
    FlashInferCutlassNvFp4LinearKernel,
    FlashInferTrtllmNvFp4LinearKernel,
 )
@@ -210,9 +209,6 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
        FlashInferCutlassNvFp4LinearKernel,
        FlashInferMxFp4LinearKernel,
    },
-    "flashinfer_cutedsl": {
-        FlashInferCuteDslNvFp4LinearKernel,
-    },
    "flashinfer_trtllm": {
        FlashInferTrtllmNvFp4LinearKernel,
    },
@@ -403,7 +399,6 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {

 _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
    PlatformEnum.CUDA: [
-        FlashInferCuteDslNvFp4LinearKernel,
        # FlashInferB12xNvFp4LinearKernel excluded from auto-selection until
        # upstream CUTLASS SM121 MMA op guard is resolved; use
        # --linear-backend flashinfer_b12x to opt in explicitly.
@@ -1043,7 +1038,6 @@ __all__ = [
    "CutlassNvFp4LinearKernel",
    "EmulationNvFp4LinearKernel",
    "FbgemmNvFp4LinearKernel",
-    "FlashInferCuteDslNvFp4LinearKernel",
    "FlashInferB12xNvFp4LinearKernel",
    "FlashInferCutlassNvFp4LinearKernel",
    "FlashInferTrtllmNvFp4LinearKernel",
@@ -28,72 +28,6 @@ from vllm.utils.flashinfer import (
 from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig


-class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
-    """NVFP4 GEMM via FlashInfer's cutedsl backend."""
-
-    @classmethod
-    def is_supported(
-        cls, compute_capability: int | None = None
-    ) -> tuple[bool, str | None]:
-        if not current_platform.is_device_capability_family(100):
-            return False, "FlashInfer cutedsl requires sm_10x"
-        if not has_flashinfer():
-            return False, "FlashInfer required"
-        return True, None
-
-    @classmethod
-    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
-        return True, None
-
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        # cutedsl uses the same swizzled + padded layout as cutlass.
-        layer.weight_scale = torch.nn.Parameter(
-            swizzle_blockscale(layer.weight_scale.data), requires_grad=False
-        )
-        padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
-            layer.weight.data
-        )
-        layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
-        layer.weights_padding_cols = weights_padding_cols
-
-    def apply_weights(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        output_size = layer.output_size_per_partition
-        output_dtype = x.dtype
-        output_shape = [*x.shape[:-1], output_size]
-
-        x_fp4, x_blockscale = scaled_fp4_quant(
-            x,
-            layer.input_global_scale_inv,
-            is_sf_swizzled_layout=True,
-            backend="flashinfer-cutedsl",
-        )
-
-        x_fp4 = pad_nvfp4_activation_for_cutlass(
-            x_fp4, getattr(layer, "weights_padding_cols", 0)
-        )
-
-        out = flashinfer_scaled_fp4_mm(
-            x_fp4,
-            layer.weight,
-            x_blockscale,
-            layer.weight_scale,
-            layer.alpha,
-            output_dtype,
-            backend="cute-dsl",
-        )
-
-        out = slice_nvfp4_output(out, output_size)
-
-        if bias is not None:
-            out = out + bias
-        return out.view(*output_shape)
-
-
 class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
    """NVFP4 GEMM via FlashInfer's CUTLASS wrapper."""

@@ -208,7 +208,6 @@ from vllm.config import (
    get_current_vllm_config,
    get_current_vllm_config_or_none,
 )
-from vllm.config.cache import CacheDType
 from vllm.distributed.parallel_state import (
    get_dcp_group,
    is_global_first_rank,
@@ -320,22 +319,6 @@ def _detect_output_quant_key(
    return kFp8StaticTensorSym


-def _canonicalize_sparse_mla_kv_cache_dtype(
-    attn_backend: type[AttentionBackend],
-    kv_cache_dtype: CacheDType,
-) -> CacheDType:
-    backend_name = attn_backend.get_name()
-    if backend_name == "FLASHMLA_SPARSE" and is_quantized_kv_cache(kv_cache_dtype):
-        return "fp8_ds_mla"
-    if backend_name == "FLASHINFER_MLA_SPARSE_SM120" and kv_cache_dtype in (
-        "auto",
-        "fp8",
-        "fp8_e4m3",
-    ):
-        return "fp8_ds_mla"
-    return kv_cache_dtype
-
-
 class MLAAttention(nn.Module, AttentionLayerBase):
    """Multi-Head Latent Attention layer.

@@ -386,7 +369,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

        if cache_config is not None:
-            kv_cache_dtype: CacheDType = cache_config.cache_dtype
+            kv_cache_dtype = cache_config.cache_dtype
            calculate_kv_scales = cache_config.calculate_kv_scales
        else:
            kv_cache_dtype = "auto"
@@ -410,22 +393,24 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                num_heads=self.num_heads,
            )

-        normalized_kv_cache_dtype = _canonicalize_sparse_mla_kv_cache_dtype(
-            self.attn_backend, kv_cache_dtype
-        )
-        if normalized_kv_cache_dtype != kv_cache_dtype:
-            if cache_config is not None:
-                cache_config.cache_dtype = normalized_kv_cache_dtype
-            kv_cache_dtype = normalized_kv_cache_dtype
+        # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format
+        # Automatically convert fp8 kv-cache format to "fp8_ds_mla"
+        if (
+            self.attn_backend.get_name() == "FLASHMLA_SPARSE"
+            and is_quantized_kv_cache(kv_cache_dtype)
+            and kv_cache_dtype != "fp8_ds_mla"
+        ):
+            assert cache_config is not None
+            cache_config.cache_dtype = "fp8_ds_mla"
+            kv_cache_dtype = "fp8_ds_mla"
            logger.info_once(
-                "Using %s KV cache format for %s backend.",
-                kv_cache_dtype,
-                self.attn_backend.get_name(),
+                "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "
+                "fp8 kv-cache format, please set `--attention-backend "
+                "FLASHINFER_MLA_SPARSE`"
            )

        if (
            self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE"
-            and kv_cache_dtype != "fp8_ds_mla"
            and is_quantized_kv_cache(kv_cache_dtype)
        ):
            logger.info_once(
@@ -23,82 +23,24 @@ def expert_num_tokens_round_up_and_sum(
    return torch.sum(ent).item()


-def compute_aligned_M_and_alignment(
-    M: int,
-    num_topk: int,
-    local_num_experts: int,
-    alignment: int,
-    expert_tokens_meta: mk.ExpertTokensMetadata | None,
-) -> tuple[int, int]:
-    """Return (M_sum, alignment_used).
-
-    `alignment_used` may be smaller than the caller-supplied `alignment` on
-    SM100/SM120 when DeepGEMM can JIT a smaller BLOCK_M for the per-call
-    expected_m. Callers that index by block size (e.g. ``M_sum // block_m``)
-    or assert workspace alignment must use the returned `alignment_used`,
-    not their original `alignment` argument.
-
-    Prefer this over the int-returning :func:`compute_aligned_M` when the
-    GEMM call site needs to wrap itself in ``mk_alignment_scope`` or
-    otherwise reason about the actual per-expert padding.
-    """
-    if (expert_tokens_meta is not None) and (
-        expert_tokens_meta.expert_num_tokens_cpu is not None
-    ):
-        return (
-            expert_num_tokens_round_up_and_sum(
-                expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment
-            ),
-            alignment,
-        )
-
-    # expert_num_tokens not on cpu. Cap padding by min(M*num_topk,
-    # local_num_experts) — at batch=1 decode only `num_topk` experts can be
-    # active, so the worst-case `local_num_experts*(align-1)` is too loose.
-    # Also shrink `alignment` to DeepGEMM's per-call theoretical BLOCK_M on
-    # SM100/SM120 when smaller.
-    expected_m = M * num_topk
-    try:
-        from vllm.utils.deep_gemm import (
-            get_theoretical_mk_alignment_for_contiguous_layout,
-        )
-
-        # num_groups=local_num_experts so the helper recovers per-expert em;
-        # omitting it over-picks BLOCK_M on SM120 (heuristic assumes em is
-        # already per-expert).
-        per_call_align = get_theoretical_mk_alignment_for_contiguous_layout(
-            expected_m=expected_m,
-            num_groups=local_num_experts,
-        )
-        if per_call_align and per_call_align <= alignment:
-            alignment = per_call_align
-    except Exception:
-        pass
-
-    max_active_experts = min(M * num_topk, local_num_experts)
-    M_sum = (M * num_topk) + max_active_experts * (alignment - 1)
-    M_sum = round_up(M_sum, alignment)
-    return M_sum, alignment
-
-
 def compute_aligned_M(
    M: int,
    num_topk: int,
    local_num_experts: int,
    alignment: int,
    expert_tokens_meta: mk.ExpertTokensMetadata | None,
-) -> int:
-    """Return ``M_sum`` only (backward-compat wrapper).
+):
+    if (expert_tokens_meta is not None) and (
+        expert_tokens_meta.expert_num_tokens_cpu is not None
+    ):
+        return expert_num_tokens_round_up_and_sum(
+            expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment
+        )

-    Equivalent to :func:`compute_aligned_M_and_alignment`'s first return
-    value. Existing downstream callers and the warmup path that only size
-    a workspace use this. Call sites that need the actual per-expert
-    alignment (to wrap GEMMs in ``mk_alignment_scope``) should use
-    :func:`compute_aligned_M_and_alignment` instead.
-    """
-    M_sum, _ = compute_aligned_M_and_alignment(
-        M, num_topk, local_num_experts, alignment, expert_tokens_meta
-    )
+    # expert_num_tokens information is not available on the cpu.
+    # compute the max required size.
+    M_sum = (M * num_topk) + local_num_experts * (alignment - 1)
+    M_sum = round_up(M_sum, alignment)
    return M_sum


@@ -109,6 +51,12 @@ def apply_expert_map(expert_id, expert_map):
    return expert_id


+@triton.jit
+def round_up_128(x: int) -> int:
+    y = 128
+    return ((x + y - 1) // y) * y
+
+
@triton.jit
 def _fwd_kernel_ep_scatter_1(
    num_recv_tokens_per_expert,
@@ -117,7 +65,6 @@ def _fwd_kernel_ep_scatter_1(
    num_experts: tl.constexpr,
    BLOCK_E: tl.constexpr,
    BLOCK_EXPERT_NUM: tl.constexpr,
-    ALIGN_M: tl.constexpr,
 ):
    cur_expert = tl.program_id(0)

@@ -127,8 +74,7 @@ def _fwd_kernel_ep_scatter_1(
        mask=offset_cumsum < num_experts,
        other=0,
    )
-    # Round up to ALIGN_M so cumsum matches the workspace's per-expert slices.
-    tokens_per_expert = ((tokens_per_expert + ALIGN_M - 1) // ALIGN_M) * ALIGN_M
+    tokens_per_expert = round_up_128(tokens_per_expert)
    cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert

    # Extract this block's offset from the register vector (warp shuffle,
@@ -281,12 +227,10 @@ def ep_scatter(
    output_tensor_scale: torch.Tensor,
    m_indices: torch.Tensor,
    output_index: torch.Tensor,
-    align_m: int = 128,
    block_size: int = 128,
    pack_ue8m0: bool = False,
 ):
-    # BLOCK_E is the m_indices fill-loop tile (masked), independent of align_m.
-    BLOCK_E = 128
+    BLOCK_E = 128  # token num of per expert is aligned to 128
    BLOCK_D = block_size  # block size of activation-scale quantization
    num_warps = 8
    num_experts = num_recv_tokens_per_expert.shape[0]
@@ -294,7 +238,7 @@ def ep_scatter(
    # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
    grid = num_experts

-    assert m_indices.shape[0] % align_m == 0
+    assert m_indices.shape[0] % BLOCK_E == 0
    assert expert_start_loc.shape[0] == num_experts

    # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
@@ -309,7 +253,6 @@ def ep_scatter(
        num_warps=num_warps,
        BLOCK_E=BLOCK_E,
        BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
-        ALIGN_M=align_m,
    )

    grid = min(recv_topk.shape[0], 1024 * 8)
@@ -475,7 +418,7 @@ def deepgemm_moe_permute(
    if block_size is not None:
        block_k = block_size

-    M_sum, align_used = compute_aligned_M_and_alignment(
+    M_sum = compute_aligned_M(
        M=topk_ids.size(0),
        num_topk=topk_ids.size(1),
        local_num_experts=local_num_experts,
@@ -539,12 +482,11 @@ def deepgemm_moe_permute(
        output_tensor_scale=aq_scale_out,
        m_indices=expert_ids,
        output_index=inv_perm,
-        align_m=align_used,
        block_size=block_k,
        pack_ue8m0=pack_ue8m0,
    )

-    return aq_out, aq_scale_out, expert_ids, inv_perm, align_used
+    return aq_out, aq_scale_out, expert_ids, inv_perm


 def deepgemm_unpermute_and_reduce(
@@ -318,12 +318,11 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):

    def supports_packed_ue8m0_act_scales(self) -> bool:
        """
-        DeepGemm supports packed ue8m0 activation scales on Blackwell-family
-        GPUs (SM100 datacenter and SM120 consumer).
+        DeepGemm supports packed ue8m0 activation scales format in devices == sm100
        """
-        return is_deep_gemm_e8m0_used() and (
-            current_platform.is_device_capability_family(100)
-            or current_platform.is_device_capability_family(120)
+        return (
+            is_deep_gemm_e8m0_used()
+            and current_platform.is_device_capability_family(100)
        )

    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEQuantConfig,
 )
 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
-    compute_aligned_M_and_alignment,
+    compute_aligned_M,
    deepgemm_moe_permute,
    deepgemm_unpermute_and_reduce,
 )
@@ -43,7 +43,6 @@ from vllm.utils.deep_gemm import (
    is_deep_gemm_supported,
    m_grouped_fp8_fp4_gemm_nt_contiguous,
    m_grouped_fp8_gemm_nt_contiguous,
-    mk_alignment_scope,
 )
 from vllm.utils.import_utils import has_deep_gemm

@@ -211,10 +210,10 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
        # Use the contiguous-layout M alignment (matches apply()); block_shape[0]
        # is the quant block (1 for MXFP8) and would under-size the workspace.
        block_m = get_mk_alignment_for_contiguous_layout()[0]
-        M_sum, align_used = compute_aligned_M_and_alignment(
+        M_sum = compute_aligned_M(
            M, topk, local_num_experts, block_m, expert_tokens_meta
        )
-        assert M_sum % align_used == 0
+        assert M_sum % block_m == 0

        activation_out_dim = self.adjust_N_for_activation(N, activation)
        workspace1 = (M_sum, max(activation_out_dim, K))
@@ -317,7 +316,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):

        assert w2.size(1) == K

-        M_sum, _ = compute_aligned_M_and_alignment(
+        M_sum = compute_aligned_M(
            M=topk_ids.size(0),
            num_topk=topk_ids.size(1),
            local_num_experts=local_num_experts,
@@ -328,7 +327,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
        a1q_perm = _resize_cache(
            workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)
        )
-        a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
+        a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
            aq=a1q,
            aq_scale=a1q_scale,
            topk_ids=topk_ids,
@@ -350,35 +349,23 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
            else {}
        )

-        # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment;
-        # otherwise the scheduler can pick the wrong expert id from m_indices
-        # under cudagraph replay.
-        with mk_alignment_scope(align_used):
-            mm1_out = _resize_cache(workspace2, (M_sum, N))
-            m_grouped_fp8_gemm_nt_contiguous(
-                (a1q, a1q_scale),
-                (w1, self.w1_scale),
-                mm1_out,
-                expert_ids,
-                **gemm_kwargs,
-            )
+        mm1_out = _resize_cache(workspace2, (M_sum, N))
+        m_grouped_fp8_gemm_nt_contiguous(
+            (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs
+        )

-            activation_out_dim = self.adjust_N_for_activation(N, activation)
-            quant_out = _resize_cache(
-                workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
-            )
-            a2q, a2q_scale = self._act_mul_quant(
-                input=mm1_out.view(-1, N), output=quant_out, activation=activation
-            )
+        activation_out_dim = self.adjust_N_for_activation(N, activation)
+        quant_out = _resize_cache(
+            workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
+        )
+        a2q, a2q_scale = self._act_mul_quant(
+            input=mm1_out.view(-1, N), output=quant_out, activation=activation
+        )

-            mm2_out = _resize_cache(workspace2, (M_sum, K))
-            m_grouped_fp8_gemm_nt_contiguous(
-                (a2q, a2q_scale),
-                (w2, self.w2_scale),
-                mm2_out,
-                expert_ids,
-                **gemm_kwargs,
-            )
+        mm2_out = _resize_cache(workspace2, (M_sum, K))
+        m_grouped_fp8_gemm_nt_contiguous(
+            (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, **gemm_kwargs
+        )

        if apply_router_weight_on_input:
            topk_weights = torch.ones_like(topk_weights)
@@ -397,8 +384,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
    """DeepGemm-based fused MoE expert implementation for FP4 weights.

    Uses m_grouped_fp8_fp4_gemm_nt_contiguous with FP8 activations and
-    MXFP4 (FP4 E2M1 packed as uint8) weights. Requires Blackwell-family
-    GPUs (SM100 datacenter or SM120 consumer).
+    MXFP4 (FP4 E2M1 packed as uint8) weights. Requires SM100+ (Blackwell).
    """

    # FP8 activation block size (hardcoded since mxfp4_w4a8 quant config
@@ -423,9 +409,9 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
    def _supports_current_device() -> bool:
        from vllm.platforms import current_platform

-        return is_deep_gemm_supported() and (
-            current_platform.is_device_capability_family(100)
-            or current_platform.is_device_capability_family(120)
+        return (
+            is_deep_gemm_supported()
+            and current_platform.is_device_capability_family(100)
        )

    @staticmethod
@@ -468,10 +454,10 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
        activation: MoEActivation,
    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
        block_m = get_mk_alignment_for_contiguous_layout()[0]
-        M_sum, align_used = compute_aligned_M_and_alignment(
+        M_sum = compute_aligned_M(
            M, topk, local_num_experts, block_m, expert_tokens_meta
        )
-        assert M_sum % align_used == 0
+        assert M_sum % block_m == 0

        activation_out_dim = self.adjust_N_for_activation(N, activation)
        workspace1 = (M_sum, max(activation_out_dim, K))
@@ -547,7 +533,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
        if global_num_experts == -1:
            global_num_experts = local_num_experts

-        M_sum, _ = compute_aligned_M_and_alignment(
+        M_sum = compute_aligned_M(
            M=topk_ids.size(0),
            num_topk=topk_ids.size(1),
            local_num_experts=local_num_experts,
@@ -558,7 +544,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
        a1q_perm = _resize_cache(
            workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)
        )
-        a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
+        a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
            aq=a1q,
            aq_scale=a1q_scale,
            topk_ids=topk_ids,
@@ -569,40 +555,37 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
        )
        assert a1q.size(0) == M_sum

-        # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment;
-        # see DeepGemmExperts.apply for rationale.
-        with mk_alignment_scope(align_used):
-            # FC1: FP8 activations x FP4 weights
-            # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4).
-            mm1_out = _resize_cache(workspace2, (M_sum, N))
-            m_grouped_fp8_fp4_gemm_nt_contiguous(
-                (a1q, a1q_scale),
-                (w1.view(torch.int8), self.w1_scale),
-                mm1_out,
-                expert_ids,
-                recipe_a=(1, self._ACT_BLOCK_K),
-                recipe_b=(1, self._WEIGHT_BLOCK_K),
-            )
+        # FC1: FP8 activations x FP4 weights
+        # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4).
+        mm1_out = _resize_cache(workspace2, (M_sum, N))
+        m_grouped_fp8_fp4_gemm_nt_contiguous(
+            (a1q, a1q_scale),
+            (w1.view(torch.int8), self.w1_scale),
+            mm1_out,
+            expert_ids,
+            recipe_a=(1, self._ACT_BLOCK_K),
+            recipe_b=(1, self._WEIGHT_BLOCK_K),
+        )

-            # SwiGLU activation + FP8 requant
-            activation_out_dim = self.adjust_N_for_activation(N, activation)
-            quant_out = _resize_cache(
-                workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
-            )
-            a2q, a2q_scale = self._act_mul_quant(
-                input=mm1_out.view(-1, N), output=quant_out, activation=activation
-            )
+        # SwiGLU activation + FP8 requant
+        activation_out_dim = self.adjust_N_for_activation(N, activation)
+        quant_out = _resize_cache(
+            workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
+        )
+        a2q, a2q_scale = self._act_mul_quant(
+            input=mm1_out.view(-1, N), output=quant_out, activation=activation
+        )

-            # FC2: FP8 activations x FP4 weights
-            mm2_out = _resize_cache(workspace2, (M_sum, K))
-            m_grouped_fp8_fp4_gemm_nt_contiguous(
-                (a2q, a2q_scale),
-                (w2.view(torch.int8), self.w2_scale),
-                mm2_out,
-                expert_ids,
-                recipe_a=(1, self._ACT_BLOCK_K),
-                recipe_b=(1, self._WEIGHT_BLOCK_K),
-            )
+        # FC2: FP8 activations x FP4 weights
+        mm2_out = _resize_cache(workspace2, (M_sum, K))
+        m_grouped_fp8_fp4_gemm_nt_contiguous(
+            (a2q, a2q_scale),
+            (w2.view(torch.int8), self.w2_scale),
+            mm2_out,
+            expert_ids,
+            recipe_a=(1, self._ACT_BLOCK_K),
+            recipe_b=(1, self._WEIGHT_BLOCK_K),
+        )

        if apply_router_weight_on_input:
            topk_weights = torch.ones_like(topk_weights)
@@ -57,40 +57,6 @@ if has_triton_kernels():
        )


-def _pack_deepgemm_mxfp4_scales(
-    w13_weight: torch.Tensor,
-    w2_weight: torch.Tensor,
-    w13_weight_scale: torch.Tensor,
-    w2_weight_scale: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-        deepgemm_post_process_weight_scale_block,
-    )
-
-    num_experts = w13_weight.shape[0]
-    intermediate_size_2 = w13_weight.shape[1]  # = intermediate*2
-    hidden_size = w13_weight.shape[2] * 2  # weight is FP4-packed
-    intermediate_size = w2_weight.shape[2] * 2  # weight is FP4-packed
-    block_shape = (1, 32)  # MXFP4 block (per-row, K=32)
-
-    return (
-        deepgemm_post_process_weight_scale_block(
-            ws=w13_weight_scale.data,
-            mn=intermediate_size_2,
-            k=hidden_size,
-            quant_block_shape=block_shape,
-            num_groups=num_experts,
-        ),
-        deepgemm_post_process_weight_scale_block(
-            ws=w2_weight_scale.data,
-            mn=hidden_size,
-            k=intermediate_size,
-            quant_block_shape=block_shape,
-            num_groups=num_experts,
-        ),
-    )
-
-
 class Mxfp4MoeBackend(Enum):
    NONE = "None"
    # DeepGEMM FP8xFP4 backend (SM100+)
@@ -686,18 +652,15 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
    """Convert loaded weights into backend-specific kernel format."""

    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
-        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
-            w13_weight,
-            w2_weight,
-            w13_weight_scale,
-            w2_weight_scale,
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            _upcast_e8m0_to_fp32,
        )

        return (
            w13_weight.data,
            w2_weight.data,
-            w13_weight_scale,
-            w2_weight_scale,
+            _upcast_e8m0_to_fp32(w13_weight_scale.data),
+            _upcast_e8m0_to_fp32(w2_weight_scale.data),
            w13_bias,
            w2_bias,
        )
@@ -1232,18 +1195,17 @@ def convert_weight_to_mxfp4_moe_kernel_format(
    """

    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
-        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
-            w13_weight,
-            w2_weight,
-            w13_weight_scale,
-            w2_weight_scale,
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            _upcast_e8m0_to_fp32,
        )

+        # Weights stay as uint8 packed FP4 — no layout change needed.
+        # Convert E8M0 uint8 scales to float32.
        return (
            w13_weight.data,
            w2_weight.data,
-            w13_weight_scale,
-            w2_weight_scale,
+            _upcast_e8m0_to_fp32(w13_weight_scale.data),
+            _upcast_e8m0_to_fp32(w2_weight_scale.data),
            w13_bias,
            w2_bias,
        )
@@ -275,12 +275,6 @@ class RoutedExperts(PluggableLayer):
    # Weight Loading Methods
    #

-    @staticmethod
-    def _to_scalar(loaded_weight: torch.Tensor) -> torch.Tensor:
-        # Per-tensor scales arrive 0-D or as shape-(1,) (llm-compressor NVFP4);
-        # reduce to a 0-D scalar. numel > 1 raises instead of broadcasting.
-        return loaded_weight.reshape(())
-
    def _load_per_tensor_weight_scale(
        self,
        shard_id: str,
@@ -294,10 +288,10 @@ class RoutedExperts(PluggableLayer):
            # We have to keep the weight scales of w1 and w3 because
            # we need to re-quantize w1/w3 weights after weight loading.
            idx = 0 if shard_id == "w1" else 1
-            param_data[expert_id][idx] = self._to_scalar(loaded_weight)
+            param_data[expert_id][idx] = loaded_weight
        # If we are in the row parallel case (down_proj)
        elif shard_id == "w2":
-            param_data[expert_id] = self._to_scalar(loaded_weight)
+            param_data[expert_id] = loaded_weight

    def _load_combined_w13_weight_scale(
        self,
@@ -531,7 +525,7 @@ class RoutedExperts(PluggableLayer):
        param_data = param.data

        # Input scales can be loaded directly and should be equal.
-        param_data[expert_id] = self._to_scalar(loaded_weight)
+        param_data[expert_id] = loaded_weight

    def _load_g_idx(
        self,
@@ -698,9 +692,7 @@ class RoutedExperts(PluggableLayer):
            ):
                scale_expert_id = global_expert_id if use_global_sf else expert_id
                scale_shard_id = 0 if shard_id == "w1" else 1
-                param.data[scale_expert_id][scale_shard_id] = self._to_scalar(
-                    loaded_weight
-                )
+                param.data[scale_expert_id][scale_shard_id] = loaded_weight.reshape(())
                return True if return_success else None

            if (
@@ -852,7 +852,7 @@ class HummingMoEMethod(FusedMoEMethodBase):

        # use moe modular
        experts: HummingIndexedExperts | HummingGroupedExperts
-        layer._ensure_moe_quant_config_init()
+        layer.ensure_moe_quant_config_init()
        assert self.moe_quant_config is not None
        if get_humming_moe_gemm_type() == "indexed":
            experts = HummingIndexedExperts(layer, self.moe, self.moe_quant_config)
@@ -1058,34 +1058,6 @@ def _upcast_e8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor:
    return fp32_bits.view(torch.float32)


-def deepgemm_post_process_weight_scale_block(
-    ws: torch.Tensor,
-    mn: int,
-    k: int,
-    quant_block_shape: tuple[int, ...],
-    num_groups: int,
-    is_sfa: bool = False,
-) -> torch.Tensor:
-    if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
-        # Scales already in E8M0 from checkpoint; upcast to fp32 and let
-        # DeepGEMM pack the layout expected by the target architecture.
-        ws = _upcast_e8m0_to_fp32(ws)
-    else:
-        assert ws.dtype == torch.float32, (
-            f"Expected tensor scales dtype to be torch.float32 or "
-            f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
-        )
-
-    return transform_sf_into_required_layout(
-        sf=ws,
-        mn=mn,
-        k=k,
-        recipe=(1, quant_block_shape[0], quant_block_shape[1]),
-        num_groups=num_groups,
-        is_sfa=is_sfa,
-    )
-
-
 def deepgemm_post_process_fp8_weight_block(
    wq: torch.Tensor,
    ws: torch.Tensor,
@@ -1101,13 +1073,13 @@ def deepgemm_post_process_fp8_weight_block(

    if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
        # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
-        # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization
+        # bits as uint8 for MXFP8) — upcast to fp32 and skip requantization
        # (weights already have power-of-two scales).
        ws = _upcast_e8m0_to_fp32(ws)
    else:
        assert ws.dtype == torch.float32, (
            f"Expected tensor scales dtype to be torch.float32 or "
-            f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
+            f"torch.float8_e8m0fnu, got {ws.dtype} instead"
        )
        if use_e8m0:
            requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)
@@ -1122,12 +1094,16 @@ def deepgemm_post_process_fp8_weight_block(
        r = wq.size(0) // g
        wq = wq.view(g, r, d)
        ws = ws.view(g, r // quant_block_shape[0], d // quant_block_shape[1])
-        dg_ws = deepgemm_post_process_weight_scale_block(
-            ws=ws,
+        # Pre-transform scale with recipe=(1, 128, 128) to broadcast + pack
+        # into TMA-aligned UE8M0 (INT32) layout. At runtime fp8_einsum uses
+        # recipe=(1, 1, 128) which sees INT dtype and skips re-transform.
+        dg_ws = transform_sf_into_required_layout(
+            sf=ws,
            mn=r,
            k=d,
-            quant_block_shape=quant_block_shape,
+            recipe=(1, quant_block_shape[0], quant_block_shape[1]),
            num_groups=g,
+            is_sfa=False,
        )
        return wq, dg_ws

@@ -1137,12 +1113,22 @@ def deepgemm_post_process_fp8_weight_block(
        wq = wq.unsqueeze(0)
        ws = ws.unsqueeze(0)

-    dg_ws = deepgemm_post_process_weight_scale_block(
-        ws=ws,
+    # From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46
+    # (1, block_n, block_k): (1, 128, 128) for FP8 block, (1, 1, 32) for MXFP8.
+    recipe = (1, quant_block_shape[0], quant_block_shape[1])
+
+    # Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp
+    # DeepGemm uses the `transform_sf_into_required_layout` function to
+    # represent scales in the correct format.
+    dg_ws = transform_sf_into_required_layout(
+        sf=ws,
        mn=wq.size(1),
        k=wq.size(2),
-        quant_block_shape=quant_block_shape,
+        recipe=recipe,
        num_groups=wq.size(0),
+        # is the scale factors for A in (Refers to the argument A in A @ B).
+        # Weights are B.
+        is_sfa=False,
    )

    if original_ndim == 2:
@@ -12,9 +12,7 @@ from tqdm import tqdm
 import vllm.envs as envs
 from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank
 from vllm.model_executor.layers.fused_moe import MoERunner
-from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
-    compute_aligned_M_and_alignment,
-)
+from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M
 from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts
 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
    TritonOrDeepGemmExperts,
@@ -27,7 +25,6 @@ from vllm.utils.deep_gemm import (
    fp8_gemm_nt,
    get_mk_alignment_for_contiguous_layout,
    m_grouped_fp8_gemm_nt_contiguous,
-    mk_alignment_scope,
 )
 from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import num_compute_units
@@ -241,7 +238,7 @@ def _get_grouped_gemm_params(
    w2: torch.Tensor,
    num_topk: int,
    max_tokens: int,
-) -> tuple[int, int, list[tuple[int, int, torch.Tensor]]]:
+) -> tuple[int, int, torch.Tensor]:
    assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts"

    block_m = get_mk_alignment_for_contiguous_layout()[0]
@@ -251,46 +248,19 @@ def _get_grouped_gemm_params(
    # Assumes all ranks have the same max_num_batched_tokens
    max_tokens = get_dp_group().world_size * max_tokens

-    request_m_values = _generate_optimal_warmup_m_values(
-        max_tokens,
-        max(w1.size(1), w2.size(1)),
-        device,
+    # This is the maximum GroupedGemm M size that we expect to run
+    # the grouped_gemm with.
+    MAX_M = compute_aligned_M(
+        max_tokens, num_topk, num_experts, block_m, expert_tokens_meta=None
    )
-    request_m_values = sorted({m for m in (*request_m_values, max_tokens) if m > 0})
-    if not request_m_values:
-        return 0, block_m, []
+    # Distribute expert-ids evenly.
+    MAX_BLOCKS = MAX_M // block_m
+    expert_ids_block = torch.randint(
+        low=0, high=num_experts, size=(MAX_BLOCKS,), device=device, dtype=torch.int32
+    )
+    expert_ids = torch.repeat_interleave(expert_ids_block, block_m, dim=0)

-    cases_by_shape: dict[tuple[int, int], torch.Tensor] = {}
-    for request_m in request_m_values:
-        M_sum, align_used = compute_aligned_M_and_alignment(
-            M=request_m,
-            num_topk=num_topk,
-            local_num_experts=num_experts,
-            alignment=block_m,
-            expert_tokens_meta=None,
-        )
-        if (M_sum, align_used) in cases_by_shape:
-            continue
-
-        num_blocks = M_sum // align_used
-        expert_ids_block = torch.randint(
-            low=0,
-            high=num_experts,
-            size=(num_blocks,),
-            device=device,
-            dtype=torch.int32,
-        )
-        cases_by_shape[(M_sum, align_used)] = torch.repeat_interleave(
-            expert_ids_block, align_used, dim=0
-        )
-
-    max_m = max(M_sum for M_sum, _ in cases_by_shape)
-    warmup_cases = [
-        (M_sum, align_used, expert_ids)
-        for (M_sum, align_used), expert_ids in sorted(cases_by_shape.items())
-    ]
-
-    return max_m, block_m, warmup_cases
+    return MAX_M, block_m, expert_ids


 def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
@@ -308,11 +278,7 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
    ):
        return

-    MAX_M, block_m, warmup_cases = _get_grouped_gemm_params(
-        w1, w2, num_topk, max_tokens
-    )
-    if not warmup_cases:
-        return
+    MAX_M, block_m, expert_ids = _get_grouped_gemm_params(w1, w2, num_topk, max_tokens)
    device = w1.device

    def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
@@ -323,14 +289,15 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
        )
        out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)

-        for num_tokens, align_used, expert_ids in warmup_cases:
-            with mk_alignment_scope(align_used):
-                m_grouped_fp8_gemm_nt_contiguous(
-                    (a1q[:num_tokens], a1q_scales[:num_tokens]),
-                    (w, w_scale),
-                    out[:num_tokens],
-                    expert_ids,
-                )
+        m_values = list(range(block_m, MAX_M + 1, block_m))
+
+        for num_tokens in m_values:
+            m_grouped_fp8_gemm_nt_contiguous(
+                (a1q[:num_tokens], a1q_scales[:num_tokens]),
+                (w, w_scale),
+                out[:num_tokens],
+                expert_ids[:num_tokens],
+            )
            if pbar is not None:
                pbar.update(1)

@@ -383,8 +350,8 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
            w13, _, w2, _, num_topk = _extract_data_from_fused_moe_module(m)
            if w13.size() in seen_grouped_sizes and w2.size() in seen_grouped_sizes:
                continue
-            _, _, warmup_cases = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
-            n_values = len(warmup_cases)
+            MAX_M, block_m, _ = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
+            n_values = (MAX_M - block_m) // block_m + 1
            if w13.size() not in seen_grouped_sizes:
                total += n_values
                seen_grouped_sizes.add(w13.size())
@@ -1,226 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up DeepSeek V4 mHC TileLang kernels before serving requests.
-
-Ported from lucifer1004/vllm-jasl with the two env-var knobs removed
-(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`).
-Gating is intrinsic: non-DSv4 models and layers without hc_* attributes
-return early, so the warmup is a no-op except where it's needed.
-"""
-
-import time
-from collections.abc import Iterable
-
-import torch
-
-from vllm.logger import init_logger
-from vllm.tracing import instrument
-from vllm.utils.math_utils import cdiv
-
-logger = init_logger(__name__)
-
-_AUTO_WARMUP_MAX_TOKENS = 16_384
-_DEFAULT_TOKEN_SIZE_CANDIDATES = (
-    1,
-    2,
-    4,
-    8,
-    16,
-    32,
-    64,
-    128,
-    256,
-    512,
-    1024,
-    2048,
-    4096,
-    8192,
-    16_384,
-)
-
-
-def _compute_mhc_pre_num_split(
-    *,
-    num_tokens: int,
-    hidden_size: int,
-    hc_mult: int,
-    num_sms: int,
-) -> int:
-    block_k = 64
-    block_m = 64
-    k = hc_mult * hidden_size
-    grid_size = cdiv(num_tokens, block_m)
-    split_k = num_sms // grid_size
-    num_block_k = cdiv(k, block_k)
-    split_k = min(split_k, num_block_k // 4)
-    return max(split_k, 1)
-
-
-def _normalize_token_sizes(
-    token_sizes: Iterable[int],
-    *,
-    max_tokens: int,
-) -> list[int]:
-    return sorted({size for size in token_sizes if 1 <= size <= max_tokens})
-
-
-def _select_mhc_warmup_token_sizes(
-    *,
-    max_tokens: int,
-    cudagraph_capture_sizes: list[int],
-) -> list[int]:
-    if max_tokens <= 0:
-        return []
-
-    max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS)
-    candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES)
-    candidates.extend(cudagraph_capture_sizes)
-    candidates.append(max_auto_tokens)
-    return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens)
-
-
-def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None:
-    for module in model.modules():
-        if module.__class__.__name__ != "DeepseekV4DecoderLayer":
-            continue
-        if all(
-            hasattr(module, attr)
-            for attr in (
-                "hc_pre",
-                "hc_post",
-                "hc_attn_fn",
-                "hc_attn_scale",
-                "hc_attn_base",
-                "hc_ffn_fn",
-                "hc_ffn_scale",
-                "hc_ffn_base",
-            )
-        ):
-            return module
-    return None
-
-
-def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None:
-    for module in model.modules():
-        if module.__class__.__name__ != "DeepseekV4Model":
-            continue
-        if all(
-            hasattr(module, attr)
-            for attr in ("hc_head_fn", "hc_head_scale", "hc_head_base")
-        ):
-            return module
-    return None
-
-
-def _warmup_layer_mhc(
-    layer: torch.nn.Module,
-    token_sizes: list[int],
-) -> None:
-    max_tokens = max(token_sizes)
-    hidden_size = int(layer.hidden_size)
-    hc_mult = int(layer.hc_mult)
-    device = layer.hc_attn_fn.device
-    residual = torch.zeros(
-        max_tokens,
-        hc_mult,
-        hidden_size,
-        dtype=torch.bfloat16,
-        device=device,
-    )
-
-    for size in token_sizes:
-        residual_slice = residual[:size]
-        for fn, scale, base in (
-            (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base),
-            (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base),
-        ):
-            layer_input, post_mix, comb_mix = layer.hc_pre(
-                residual_slice,
-                fn,
-                scale,
-                base,
-            )
-            layer.hc_post(layer_input, residual_slice, post_mix, comb_mix)
-
-
-def _warmup_hc_head(
-    model: torch.nn.Module,
-    token_sizes: list[int],
-) -> None:
-    # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored
-    # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp
-    # instance attached to the model as ``hc_head_op``. We call through
-    # that instance so the warmup exercises the same dispatched
-    # implementation as the inference path.
-    hc_head_op = getattr(model, "hc_head_op", None)
-    if hc_head_op is None:
-        return
-
-    max_tokens = max(token_sizes)
-    hidden_size = int(model.config.hidden_size)
-    hc_mult = int(model.hc_mult)
-    device = model.hc_head_fn.device
-    hidden_states = torch.zeros(
-        max_tokens,
-        hc_mult,
-        hidden_size,
-        dtype=torch.bfloat16,
-        device=device,
-    )
-
-    for size in token_sizes:
-        hc_head_op(
-            hidden_states[:size],
-            model.hc_head_fn,
-            model.hc_head_scale,
-            model.hc_head_base,
-            model.rms_norm_eps,
-            model.hc_eps,
-        )
-
-
-@instrument(span_name="DeepSeek V4 mHC warmup")
-def deepseek_v4_mhc_warmup(
-    model: torch.nn.Module,
-    *,
-    max_tokens: int,
-    cudagraph_capture_sizes: list[int] | None = None,
-) -> None:
-    # Cheap model-type gate before walking ``model.modules()``. The class
-    # walk below is O(num_layers) and shows up in startup time on very
-    # large checkpoints; bail out for any model that is not DeepSeek V4.
-    config = getattr(model, "config", None)
-    model_type = getattr(config, "model_type", None) if config is not None else None
-    if model_type is not None and model_type != "deepseek_v4":
-        return
-
-    layer = _find_first_mhc_layer(model)
-    if layer is None:
-        return
-
-    device = layer.hc_attn_fn.device
-    if device.type != "cuda":
-        return
-
-    deepseek_model = _find_deepseek_v4_model(model)
-    token_sizes = _select_mhc_warmup_token_sizes(
-        max_tokens=max_tokens,
-        cudagraph_capture_sizes=cudagraph_capture_sizes or [],
-    )
-    if not token_sizes:
-        return
-
-    started = time.perf_counter()
-    logger.info(
-        "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s",
-        token_sizes,
-    )
-    with torch.inference_mode():
-        _warmup_layer_mhc(layer, token_sizes)
-        if deepseek_model is not None:
-            _warmup_hc_head(deepseek_model, token_sizes)
-        torch.accelerator.synchronize()
-    logger.info(
-        "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.",
-        time.perf_counter() - started,
-    )
@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FlashInfer autotune cache helpers."""
-
-import hashlib
-import os
-import tempfile
-from contextlib import suppress
-from pathlib import Path
-from typing import TYPE_CHECKING
-
-import vllm.envs as envs
-from vllm.compilation.caching import aot_compile_hash_factors
-
-if TYPE_CHECKING:
-    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-
-
-def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str:
-    factors = aot_compile_hash_factors(runner.vllm_config)
-    return hashlib.sha256(str(factors).encode()).hexdigest()
-
-
-def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
-    override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR
-    if override_dir:
-        root = Path(override_dir).expanduser()
-    else:
-        from flashinfer.jit import env as flashinfer_jit_env
-
-        flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR
-        root = (
-            Path(envs.VLLM_CACHE_ROOT)
-            / "flashinfer_autotune_cache"
-            / flashinfer_workspace.parent.name
-            / flashinfer_workspace.name
-        )
-
-    output_dir = root / flashinfer_autotune_cache_hash(runner)
-    output_dir.mkdir(parents=True, exist_ok=True)
-    return output_dir / "autotune_configs.json"
-
-
-def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None:
-    cache_path.parent.mkdir(parents=True, exist_ok=True)
-    fd, tmp_path = tempfile.mkstemp(
-        dir=cache_path.parent, suffix=".tmp", prefix=f".{cache_path.name}."
-    )
-    try:
-        with os.fdopen(fd, "wb") as f:
-            f.write(contents)
-        os.replace(tmp_path, cache_path)
-    except BaseException:
-        with suppress(OSError):
-            os.unlink(tmp_path)
-        raise
@@ -1,255 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warmup and autotune helpers for FlashInfer sparse MLA backends."""
-
-from typing import TYPE_CHECKING, cast
-
-import torch
-
-from vllm.logger import init_logger
-from vllm.model_executor.warmup.flashinfer_autotune_cache import (
-    resolve_flashinfer_autotune_file,
-    write_flashinfer_autotune_cache,
-)
-from vllm.platforms import current_platform
-from vllm.utils.flashinfer import autotune as flashinfer_autotune
-from vllm.utils.flashinfer import has_flashinfer
-from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
-
-if TYPE_CHECKING:
-    from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2GPUModelRunner
-    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-    from vllm.v1.worker.gpu_worker import Worker
-
-logger = init_logger(__name__)
-
-_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
-    {
-        "FLASHMLA_SPARSE_DSV4",
-        "FLASHINFER_MLA_SPARSE_DSV4",
-        "ROCM_FLASHMLA_SPARSE_DSV4",
-        "DEEPSEEK_SPARSE_SWA",
-    }
-)
-_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"})
-_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"})
-
-_FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
-    "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2",
-    "FLASHINFER_MLA_SPARSE_DSV4": "DSv4",
-}
-
-_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
-
-
-def _attention_backend_name(backend: object) -> str | None:
-    get_name = getattr(backend, "get_name", None)
-    if get_name is None:
-        return None
-    try:
-        return get_name()
-    except NotImplementedError:
-        return None
-
-
-def _has_deepseek_v4_sparse_mla_backend(runner: "GPUModelRunner") -> bool:
-    for groups in getattr(runner, "attn_groups", []) or ():
-        for group in groups:
-            name = _attention_backend_name(getattr(group, "backend", None))
-            if name in _DEEPSEEK_V4_SPARSE_MLA_BACKENDS:
-                return True
-    return False
-
-
-def _flashinfer_sparse_mla_decode_label(
-    runner: "GPUModelRunner",
-    allowed_backends: frozenset[str],
-) -> str | None:
-    for groups in getattr(runner, "attn_groups", []) or ():
-        for group in groups:
-            name = _attention_backend_name(getattr(group, "backend", None))
-            if name in allowed_backends:
-                return _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS.get(name)
-    return None
-
-
-def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int:
-    return max(0, min(num_tokens, max_tokens))
-
-
-def _uses_v2_model_runner(runner: "GPUModelRunner") -> bool:
-    vllm_config = getattr(runner, "vllm_config", None)
-    return bool(getattr(vllm_config, "use_v2_model_runner", False))
-
-
-def _run_flashinfer_sparse_mla_decode_autotune(
-    worker: "Worker",
-    num_tokens: int,
-    allowed_backends: frozenset[str],
-) -> bool:
-    """Autotune FlashInfer's SM120 sparse-MLA decode path."""
-    runner = worker.model_runner
-    log_label = _flashinfer_sparse_mla_decode_label(runner, allowed_backends)
-    if log_label is None:
-        return False
-    if worker.vllm_config.kernel_config.enable_flashinfer_autotune is not True:
-        return False
-    if not has_flashinfer() or not current_platform.is_device_capability_family(120):
-        return False
-
-    try:
-        from flashinfer.autotuner import AutoTuner
-    except ImportError:
-        logger.warning(
-            "Skipping FlashInfer SM120 sparse MLA decode autotune because "
-            "FlashInfer autotuner is unavailable."
-        )
-        return False
-
-    from vllm.distributed.parallel_state import get_world_group
-
-    world = get_world_group()
-    is_leader = world.rank_in_group == 0
-    cache_path = resolve_flashinfer_autotune_file(runner)
-
-    dummy_run_kwargs = dict(
-        num_tokens=num_tokens,
-        skip_eplb=True,
-        is_profile=True,
-        force_attention=True,
-        create_mixed_batch=True,
-    )
-
-    if is_leader:
-        logger.info(
-            "Autotuning FlashInfer SM120 sparse MLA %s decode with cache: %s",
-            log_label,
-            cache_path,
-        )
-
-    with torch.inference_mode():
-        warmup_executed = True
-        if is_leader:
-            if _uses_v2_model_runner(runner):
-                v2_runner = cast("V2GPUModelRunner", runner)
-                warmup_executed = run_mixed_prefill_decode_warmup(
-                    v2_runner,
-                    worker.execute_model,
-                    worker.sample_tokens,
-                    num_tokens,
-                    mixed_step_context=flashinfer_autotune(True, cache=str(cache_path)),
-                    req_id_prefix="_sparse_mla_v2_warmup",
-                )
-            else:
-                with flashinfer_autotune(True, cache=str(cache_path)):
-                    runner._dummy_run(**dummy_run_kwargs)
-        else:
-            if _uses_v2_model_runner(runner):
-                v2_runner = cast("V2GPUModelRunner", runner)
-                warmup_executed = run_mixed_prefill_decode_warmup(
-                    v2_runner,
-                    worker.execute_model,
-                    worker.sample_tokens,
-                    num_tokens,
-                    req_id_prefix="_sparse_mla_v2_warmup",
-                )
-            else:
-                runner._dummy_run(**dummy_run_kwargs)
-
-    if not warmup_executed:
-        return False
-
-    tune_results: bytes | None = None
-    if is_leader and cache_path.exists():
-        with open(cache_path, "rb") as f:
-            tune_results = f.read()
-
-    tune_results = world.broadcast_object(tune_results, src=0)
-    if tune_results is None:
-        logger.warning(
-            "No FlashInfer SM120 sparse MLA %s decode autotune cache entries found. "
-            "Falling back to FlashInfer's default tactic heuristic.",
-            log_label,
-        )
-        world.barrier()
-        return True
-
-    write_flashinfer_autotune_cache(cache_path, tune_results)
-    world.barrier()
-
-    AutoTuner.get().load_configs(str(cache_path))
-    logger.info(
-        "FlashInfer SM120 sparse MLA %s decode autotune cache loaded on rank %d "
-        "from %s.",
-        log_label,
-        world.rank_in_group,
-        cache_path,
-    )
-    return True
-
-
-def _flashinfer_sparse_mla_decode_autotune(
-    worker: "Worker",
-    num_tokens: int,
-) -> bool:
-    return _run_flashinfer_sparse_mla_decode_autotune(
-        worker, num_tokens, _FLASHINFER_MLA_SPARSE_BACKENDS
-    )
-
-
-def _deepseek_v4_sparse_mla_decode_autotune(
-    worker: "Worker",
-    num_tokens: int,
-) -> bool:
-    return _run_flashinfer_sparse_mla_decode_autotune(
-        worker, num_tokens, _DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS
-    )
-
-
-def flashinfer_sparse_mla_decode_autotune_warmup(worker: "Worker") -> None:
-    """Autotune generic FlashInfer sparse MLA decode when selected."""
-    runner = worker.model_runner
-    if runner.is_pooling_model:
-        return
-
-    max_tokens = worker.scheduler_config.max_num_batched_tokens
-    mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens)
-    if mixed_tokens <= 0:
-        return
-    _flashinfer_sparse_mla_decode_autotune(worker, mixed_tokens)
-
-
-def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None:
-    """Warm DSv4 sparse-MLA mixed prefill+decode attention."""
-    runner = worker.model_runner
-    if runner.is_pooling_model or not _has_deepseek_v4_sparse_mla_backend(runner):
-        return
-
-    max_tokens = worker.scheduler_config.max_num_batched_tokens
-    mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens)
-    if mixed_tokens <= 0:
-        return
-
-    logger.info(
-        "Warming up DeepSeek V4 sparse MLA attention for mixed tokens=%s.",
-        mixed_tokens,
-    )
-    mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens)
-    if not mixed_warmup_done:
-        if _uses_v2_model_runner(runner):
-            v2_runner = cast("V2GPUModelRunner", runner)
-            run_mixed_prefill_decode_warmup(
-                v2_runner,
-                worker.execute_model,
-                worker.sample_tokens,
-                mixed_tokens,
-                req_id_prefix="_sparse_mla_v2_warmup",
-            )
-        else:
-            runner._dummy_run(
-                num_tokens=mixed_tokens,
-                skip_eplb=True,
-                is_profile=True,
-                force_attention=True,
-                create_mixed_batch=True,
-            )
@@ -6,24 +6,16 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
 happen during model execution.
 """

+import hashlib
+from pathlib import Path
 from typing import TYPE_CHECKING

 import torch

 import vllm.envs as envs
+from vllm.compilation.caching import aot_compile_hash_factors
 from vllm.logger import init_logger
 from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
-from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
-    deepseek_v4_mhc_warmup,
-)
-from vllm.model_executor.warmup.flashinfer_autotune_cache import (
-    resolve_flashinfer_autotune_file,
-    write_flashinfer_autotune_cache,
-)
-from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
-    deepseek_v4_sparse_mla_attention_warmup,
-    flashinfer_sparse_mla_decode_autotune_warmup,
-)
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
 from vllm.utils.flashinfer import has_flashinfer
@@ -35,26 +27,36 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)


+def _flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str:
+    factors = aot_compile_hash_factors(runner.vllm_config)
+    return hashlib.sha256(str(factors).encode()).hexdigest()
+
+
+def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
+    override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR
+    if override_dir:
+        root = Path(override_dir).expanduser()
+    else:
+        from flashinfer.jit import env as flashinfer_jit_env
+
+        flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR
+        root = (
+            Path(envs.VLLM_CACHE_ROOT)
+            / "flashinfer_autotune_cache"
+            / flashinfer_workspace.parent.name
+            / flashinfer_workspace.name
+        )
+
+    output_dir = root / _flashinfer_autotune_cache_hash(runner)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    return output_dir / "autotune_configs.json"
+
+
 def kernel_warmup(worker: "Worker"):
    from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
        minimax_m3_msa_warmup,
    )

-    # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
-    # layer per token; warm them across token sizes first so the first real
-    # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside).
-    deepseek_v4_mhc_warmup(
-        worker.get_model(),
-        max_tokens=worker.scheduler_config.max_num_batched_tokens,
-        cudagraph_capture_sizes=(
-            worker.vllm_config.compilation_config.cudagraph_capture_sizes or []
-        ),
-    )
-
-    # Run next so input-prep kernels JIT against pristine runner state.
-    flashinfer_sparse_mla_decode_autotune_warmup(worker)
-    deepseek_v4_sparse_mla_attention_warmup(worker)
-
    # Deep GEMM warmup
    do_deep_gemm_warmup = (
        envs.VLLM_USE_DEEP_GEMM
@@ -145,7 +147,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
    world = get_world_group()
    is_leader = world.rank_in_group == 0

-    cache_path = resolve_flashinfer_autotune_file(runner)
+    cache_path = _resolve_flashinfer_autotune_file(runner)
    if is_leader:
        logger.info("Using FlashInfer autotune cache file: %s", cache_path)

@@ -181,7 +183,9 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
            "Falling back to default tactics."
        )
    else:
-        write_flashinfer_autotune_cache(cache_path, tune_results)
+        if not is_leader and world.local_rank == 0:
+            with open(cache_path, "wb") as f:
+                f.write(tune_results)
        world.barrier()
        from flashinfer.autotuner import AutoTuner

@@ -62,22 +62,23 @@ logger = init_logger(__name__)


 def _resolve_dsv4_kv_cache_dtype(
-    use_fp8_ds_mla_layout: bool,
+    use_flashmla_fp8_layout: bool,
    kv_cache_dtype: str,
    cache_config: CacheConfig | None,
 ) -> tuple[str, torch.dtype]:
    """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.

    Both layouts are paged; they differ in the per-token block format. The
-    ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
-    canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
-    page-size specs pick the 576B per-token slot). Plain-row backends store each
-    token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
+    FlashMLA fp8 layout (FlashMLA / ROCm Aiter) is the ``fp8_ds_mla`` format:
+    UE8M0 block-scaled fp8 packed as ``uint8`` (the canonical ``fp8_ds_mla``
+    string is written back onto ``cache_config`` so the page-size specs pick
+    the 576B per-token slot). Otherwise (FlashInfer) each token's KV row is
+    stored in its plain element dtype — bf16 or per-tensor FP8 E4M3.
    """
-    if use_fp8_ds_mla_layout:
+    if use_flashmla_fp8_layout:
        # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
        assert kv_cache_dtype.startswith("fp8"), (
-            f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, "
+            f"DeepseekV4 FlashMLA fp8 layout only supports fp8 kv-cache, "
            f"got {kv_cache_dtype}"
        )
        if kv_cache_dtype != "fp8_ds_mla":
@@ -99,20 +100,18 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

    The platform-specific sparse-MLA forward (``forward_mqa`` /
    ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a
-    subclass — ``DeepseekV4FlashMLAAttention`` /
-    ``DeepseekV4FlashInferSM120Attention`` /
-    ``DeepseekV4FlashInferMLAAttention`` (CUDA) or
-    ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific
-    deepseek_v4 model module. The base is never instantiated directly.
+    subclass — ``DeepseekV4FlashMLAAttention`` / ``DeepseekV4FlashInferMLAAttention``
+    (CUDA) or ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the
+    platform-specific deepseek_v4 model module. The base is never instantiated
+    directly.
    """

    # Provided by the platform subclass.
    backend_cls: ClassVar[type[AttentionBackend]]
    # KV-cache per-token block format (both layouts are paged). True (default)
-    # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain
-    # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when
-    # a single attention class dispatches across arch-specific layouts.
-    use_fp8_ds_mla_layout: ClassVar[bool] = True
+    # = FlashMLA / ROCm fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8);
+    # False = FlashInfer plain bf16 / per-tensor fp8 KV row.
+    use_flashmla_fp8_layout: ClassVar[bool] = True
    # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
    # workspace allocated in _forward_prefill and is also read by the dummy-run
    # path to pre-reserve that workspace.
@@ -146,10 +145,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
        """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
        raise NotImplementedError

-    def _uses_fp8_ds_mla_layout(self) -> bool:
-        """Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
-        return self.use_fp8_ds_mla_layout
-
    def __init__(
        self,
        vllm_config: VllmConfig,
@@ -281,10 +276,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
        )
        self.max_model_len = vllm_config.model_config.max_model_len

-        # Resolve the kv-cache dtype from this backend's block format. The same
-        # resolution drives the SWA cache tensor dtype below.
+        # Resolve the kv-cache dtype from this backend's block format (a
+        # ClassVar set by the subclass): fp8_ds_mla (UE8M0 block-scaled fp8 as
+        # uint8) for FlashMLA / ROCm, vs a plain bf16 / per-tensor fp8 row for
+        # FlashInfer. The same resolution drives the SWA cache tensor dtype
+        # below.
        self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
-            self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
+            self.use_flashmla_fp8_layout, cache_config.cache_dtype, cache_config
        )

        self.swa_cache_layer = DeepseekV4SWACache(
@@ -541,7 +539,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

        # kv is unchanged; attention reads kv solely via swa_kv_cache.
        if cache_dtype == torch.uint8:
-            # fp8_ds_mla UE8M0 paged path. Horizontally fused:
+            # Legacy FlashMLA UE8M0 paged path. Horizontally fused:
            #   Q side:  per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
            #            the padding head slots; the kernel allocates and returns
            #            the padded q tensor.
@@ -559,10 +557,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
                swa_metadata.block_size,
            )

-        # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
-        # row in its element dtype (no Q padding). bf16 rewrites q in place;
-        # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
-        # KV row.
+        # FlashInfer full-cache path: the [num_blocks, block_size, 512] cache
+        # stores the KV row in its plain dtype (no Q padding). bf16 rewrites q
+        # in place; per-tensor fp8 writes a separately-allocated fp8 q and
+        # quantizes the KV row.
        block_size = swa_metadata.block_size
        swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
        if cache_dtype == torch.bfloat16:
@@ -603,18 +601,18 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
            self.compress_ratio <= 1
        ):  # SWA part. Allocated separately as DeepseekV4SWACache.
            return None
-        # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
-        # alignment; plain bf16 / per-tensor fp8 rows use natural element-size
-        # pages.
-        uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
+        # FlashMLA uses the fp8_ds_mla block format (UE8M0 block-scaled fp8 as
+        # uint8, 576B aligned); FlashInfer stores a plain bf16 / per-tensor fp8
+        # row with no extra alignment.
+        is_flashmla = self.kv_cache_dtype == "fp8_ds_mla"
        return MLAAttentionSpec(
            block_size=vllm_config.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_dim,
-            dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
+            dtype=torch.uint8 if is_flashmla else self.kv_cache_torch_dtype,
            compress_ratio=self.compress_ratio,
            cache_dtype_str=self.kv_cache_dtype,
-            alignment=576 if uses_fp8_ds_mla_layout else None,
+            alignment=576 if is_flashmla else None,  # FlashMLA needs 576B
            model_version="deepseek_v4",
        )

@@ -40,18 +40,9 @@ def _fused_inv_rope_fp8_quant_per_head(
    USE_GDC: tl.constexpr,
    launch_pdl: tl.constexpr,  # triton metadata
 ):
-    # Cast every stride to int64 — without this, Python-int strides are
-    # inferred as int32 and `pid_token(int64) × stride(int32)` can lower to
-    # int32 arithmetic, wrapping past 2³¹ for large prefill batches → IMA.
+    # int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
    pid_token = tl.program_id(0).to(tl.int64)
    pid_gh = tl.program_id(1).to(tl.int64)
-    o_stride_token = o_stride_token.to(tl.int64)
-    o_stride_head = o_stride_head.to(tl.int64)
-    cache_stride_pos = cache_stride_pos.to(tl.int64)
-    fp8_stride_group = fp8_stride_group.to(tl.int64)
-    fp8_stride_token = fp8_stride_token.to(tl.int64)
-    scale_stride_group = scale_stride_group.to(tl.int64)
-    scale_stride_k = scale_stride_k.to(tl.int64)

    g = pid_gh // heads_per_group
    head_in_group = pid_gh % heads_per_group
@@ -155,17 +155,17 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
            raise ValueError(f"Invalid compress ratio: {compress_ratio}")

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
-        # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain
-        # full-cache rows share state pages with contiguous KV pages, so padding
-        # would break page matching.
-        uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
+        # FlashMLA's UE8M0 paged layout needs 576B alignment; the FlashInfer
+        # full-cache path shares state pages with contiguous KV pages, so
+        # padding would break page matching.
+        is_flashmla = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
        return SlidingWindowMLASpec(  # only has one vector instead of K + V
            block_size=self.block_size,
            num_kv_heads=1,
            head_size=self.state_dim,
            dtype=self.dtype,
            sliding_window=self.sliding_window,
-            alignment=576 if uses_fp8_ds_mla_layout else None,
+            alignment=576 if is_flashmla else None,
        )

    def forward(self): ...
@@ -340,8 +340,8 @@ class DeepseekCompressor(nn.Module):
        k_cache_layer = self._static_forward_context[self.k_cache_prefix]
        kv_cache = k_cache_layer.kv_cache

-        # Plain-row V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
-        # fp8_ds_mla path uses the UE8M0 paged uint8 layout.
+        # FlashInfer V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
+        # legacy FlashMLA path uses the UE8M0 paged uint8 layout.
        store_full_kv = self.head_dim == 512 and kv_cache.dtype != torch.uint8
        store_full_fp8 = kv_cache.dtype == torch.float8_e4m3fn
        fp8_scale = (
@@ -358,8 +358,8 @@ class DeepseekCompressor(nn.Module):
                compress_norm_rope_store_cutedsl,
            )

-            # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla
-            # layout and the plain full-cache layout. The full-cache flags
+            # head=512 on CUDA always uses cutedsl, for both the legacy UE8M0
+            # layout and the FlashInfer full-cache layout. The full-cache flags
            # are consumed only here.
            compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
            extra_kwargs: dict[str, Any] = dict(
@@ -1,6 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""DeepSeek V4 FlashInfer sparse MLA backend."""
+"""DeepSeek V4 FlashInfer TRTLLM-gen sparse MLA backend.
+
+Uses FlashInfer's public ``trtllm_batch_decode_sparse_mla_dsv4`` launcher with a
+plain bf16 / per-tensor FP8 KV row (vs FlashMLA's packed ``fp8_ds_mla`` block
+format). Shares the V4 sparse-index pipeline (SWA cache + compressor + indexer,
+256-token blocks, head_size 512) with the FlashMLA V4 backend; only the
+attention forward differs.
+"""

 from typing import TYPE_CHECKING, ClassVar, cast

@@ -11,7 +18,6 @@ from vllm.forward_context import get_forward_context
 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
 from vllm.models.deepseek_v4.common.ops import (
    build_flashinfer_mixed_sparse_indices,
-    compute_global_topk_indices_and_lens,
 )
 from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
    compute_fp8_einsum_recipe,
@@ -21,14 +27,13 @@ from vllm.models.deepseek_v4.sparse_mla import (
    DeepseekV4FlashMLABackend,
    DeepseekV4FlashMLAMetadata,
 )
-from vllm.platforms import current_platform
-from vllm.platforms.interface import DeviceCapability
 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
-from vllm.v1.attention.backend import MultipleOf

 if TYPE_CHECKING:
    from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata

+# 128 MB TRTLLM-gen workspace, allocated once per device and zero-initialized
+# (required for first use). Reused across all FlashInfer V4 layers.
 _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}

@@ -46,113 +51,34 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:


 class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
-    """FlashInfer backend using the DSv4 sparse metadata/cache layout.
+    """Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl.

-    Inheriting from the FlashMLA V4 backend reuses its
-    ``DeepseekV4FlashMLAMetadata`` builder.
+    Inheriting from the FlashMLA V4 backend reuses its ``DeepseekV4FlashMLAMetadata``
+    builder.
    """

    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
-    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
-        "auto",
-        "bfloat16",
-        "fp8",
-        "fp8_e4m3",
-        "fp8_ds_mla",
-    ]
-
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
-        return [256]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16", "fp8"]

    @staticmethod
    def get_name() -> str:
        return "FLASHINFER_MLA_SPARSE_DSV4"

-    @classmethod
-    def get_supported_head_sizes(cls) -> list[int]:
-        return [512]
-
-    @classmethod
-    def supports_sink(cls) -> bool:
-        return True
-
-    @classmethod
-    def is_sparse(cls) -> bool:
-        return True
-
-    @classmethod
-    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
-        return capability.major in [10, 12]
-
-    @classmethod
-    def supports_combination(
-        cls,
-        head_size: int,
-        dtype: torch.dtype,
-        kv_cache_dtype: CacheDType | None,
-        block_size: int | None,
-        use_mla: bool,
-        has_sink: bool,
-        use_sparse: bool,
-        use_mm_prefix: bool,
-        device_capability: DeviceCapability,
-    ) -> str | None:
-        if device_capability.major == 10:
-            if kv_cache_dtype == "fp8_ds_mla":
-                return (
-                    "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain "
-                    "per-tensor FP8 KV layout, not fp8_ds_mla"
-                )
-            if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"):
-                return "kv_cache_dtype not supported"
-            return None
-        if device_capability.major == 12:
-            if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"):
-                return "kv_cache_dtype not supported"
-            from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
-
-            if not has_flashinfer_sparse_mla_sm120():
-                return (
-                    "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's "
-                    "sparse MLA decode API"
-                )
-            return None
-        return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        device_capability = current_platform.get_device_capability()
-        if device_capability is not None and device_capability.major == 12:
-            return DeepseekV4FlashMLABackend.get_kv_cache_shape(
-                num_blocks,
-                block_size,
-                num_kv_heads,
-                head_size,
-                cache_dtype_str,
-            )
-        assert num_kv_heads == 1
-        return (num_blocks, block_size, head_size)
-

 class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
-    """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
+    """FlashInfer TRTLLM-gen sparse MLA attention layer for DeepSeek V4."""

    backend_cls = DeepseekV4FlashInferMLASparseBackend
-    use_fp8_ds_mla_layout: ClassVar[bool] = False
+    # FlashInfer stores a plain bf16 / per-tensor fp8 KV row, not the FlashMLA
+    # packed fp8_ds_mla block format (UE8M0 block-scaled fp8 as uint8).
+    use_flashmla_fp8_layout: ClassVar[bool] = False

    @classmethod
    def get_padded_num_q_heads(cls, num_heads: int) -> int:
        # FP8 decode kernel only supports h_q = 64 or 128.
        if num_heads > 128:
            raise ValueError(
-                f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
+                f"DeepseekV4 Flashinfer MLA Sparse does not support {num_heads} heads "
                "(FP8 decode kernel requires h_q in {64, 128})."
            )
        return 64 if num_heads <= 64 else 128
@@ -180,6 +106,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
        # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
        if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
            return
+        # TODO: load real per-tensor Q/KV scales from the checkpoint; unit
+        # scales until the scale tensor names are wired.
        fp8_q_scale = 1.0
        fp8_kv_scale = 1.0
        self.register_buffer(
@@ -197,8 +125,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
            torch.tensor([fp8_kv_scale], dtype=torch.float32),
            persistent=False,
        )
-        # TRTLLM-gen takes scalar scale args on a distinct C++ path vs
-        # one-element tensors, so these are Python floats.
+        # TRTLLM-gen takes scalar scale args on a distinct (correct) C++ path
+        # vs 1-elem tensors, so these are Python floats. bmm1 folds the softmax
+        # scale and the Q/KV per-tensor scales; bmm2 is the KV scale.
        self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
        self._flashinfer_fp8_bmm2_scale = fp8_kv_scale

@@ -458,8 +387,9 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
        assert query_start_loc is not None and query_start_loc_cpu is not None

-        # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
-        # uniform-q batches, and this avoids flattening mixed batches into one call.
+        # Keep Perkz's two-call decode/prefill split: the TRTLLM-gen launcher is
+        # tuned for uniform-q batches, and collapsing the mixed batch into a
+        # single call is the suspected source of the prior IMA.
        if num_decode_tokens > 0:
            decode_cu = query_start_loc[: num_decodes + 1]
            decode_cu_cpu = query_start_loc_cpu[: num_decodes + 1]
@@ -504,379 +434,3 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
                cum_seq_lens_q=prefill_cu,
                max_q_len=int(prefill_lens_cpu.max().item()),
            )
-
-
-class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
-    """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels."""
-
-    backend_cls = DeepseekV4FlashInferMLASparseBackend
-    use_fp8_ds_mla_layout: ClassVar[bool] = True
-
-    @staticmethod
-    def _get_workspace(device: torch.device) -> torch.Tensor:
-        return _get_flashinfer_dsv4_workspace(device)
-
-    @staticmethod
-    def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor:
-        if kv_cache.dtype == torch.float8_e4m3fn:
-            kv_cache = kv_cache.view(torch.uint8)
-        if kv_cache.dim() == 4:
-            return kv_cache
-        return kv_cache.unsqueeze(-2)
-
-    @classmethod
-    def get_padded_num_q_heads(cls, num_heads: int) -> int:
-        if num_heads <= 16:
-            return 16
-        if num_heads <= 32:
-            return 32
-        if num_heads <= 64:
-            return 64
-        if num_heads <= 128:
-            return 128
-        raise ValueError(
-            f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
-            "(SM120 kernel requires h_q in {16, 32, 64, 128})."
-        )
-
-    def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
-        return deep_gemm_fp8_o_proj(
-            o,
-            positions,
-            self.rotary_emb.cos_sin_cache,
-            self.wo_a,
-            self.wo_b,
-            n_groups=self.n_local_groups,
-            heads_per_group=self.n_local_heads // self.n_local_groups,
-            nope_dim=self.nope_head_dim,
-            rope_dim=self.rope_head_dim,
-            o_lora_rank=self.o_lora_rank,
-            einsum_recipe=self._einsum_recipe,
-            tma_aligned_scales=self._tma_aligned_scales,
-        )
-
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
-
-        if not has_flashinfer_sparse_mla_sm120():
-            raise RuntimeError(
-                "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires FlashInfer's "
-                "sparse MLA decode API."
-            )
-        self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
-        # Per-tensor FP8 cache path scales.
-        if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
-            return
-        fp8_q_scale = 1.0
-        fp8_kv_scale = 1.0
-        self.register_buffer(
-            "_flashinfer_fp8_q_scale",
-            torch.tensor([fp8_q_scale], dtype=torch.float32),
-            persistent=False,
-        )
-        self.register_buffer(
-            "_flashinfer_fp8_q_scale_inv",
-            torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
-            persistent=False,
-        )
-        self.register_buffer(
-            "_flashinfer_fp8_kv_scale",
-            torch.tensor([fp8_kv_scale], dtype=torch.float32),
-            persistent=False,
-        )
-        # FlashInfer expects scalar scale arguments for this path.
-        self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
-        self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
-
-    def _reserve_empty_forward_workspace(self) -> None:
-        self._get_workspace(
-            torch.device("cuda", torch.accelerator.current_device_index())
-        )
-
-    def _forward_sparse_impl(
-        self,
-        q: torch.Tensor,
-        output: torch.Tensor,
-        flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
-        swa_metadata: "DeepseekSparseSWAMetadata",
-        self_kv_cache: torch.Tensor | None,
-        swa_kv_cache: torch.Tensor,
-        swa_only: bool,
-    ) -> None:
-        num_decode_tokens = swa_metadata.num_decode_tokens
-        if swa_metadata.num_prefills > 0:
-            self._forward_prefill(
-                q=q[num_decode_tokens:],
-                compressed_k_cache=self_kv_cache,
-                swa_k_cache=swa_kv_cache,
-                output=output[num_decode_tokens:],
-                attn_metadata=flashmla_metadata,
-                swa_metadata=swa_metadata,
-            )
-        if swa_metadata.num_decodes > 0:
-            self._forward_decode(
-                q=q[:num_decode_tokens],
-                kv_cache=self_kv_cache,
-                swa_metadata=swa_metadata,
-                attn_metadata=flashmla_metadata,
-                swa_only=swa_only,
-                output=output[:num_decode_tokens],
-            )
-
-    def forward_mqa(
-        self,
-        q: torch.Tensor,
-        kv: torch.Tensor,
-        positions: torch.Tensor,
-        output: torch.Tensor,
-    ) -> None:
-        # Output may be padded to backend-supported head counts.
-        assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
-            f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
-        )
-        assert output.shape[1] >= q.shape[1], (
-            f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
-        )
-        # Per-tensor FP8 q produces a bf16 attention output.
-        expected_output_dtype = (
-            torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
-        )
-        assert output.dtype == expected_output_dtype, (
-            f"output dtype {output.dtype} must match expected {expected_output_dtype} "
-            f"for q dtype {q.dtype}"
-        )
-
-        forward_context = get_forward_context()
-        attn_metadata = forward_context.attn_metadata
-        if attn_metadata is None:
-            self._reserve_empty_forward_workspace()
-            output.zero_()
-            return
-
-        assert isinstance(attn_metadata, dict)
-        flashmla_metadata = cast(
-            DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
-        )
-        swa_metadata = cast(
-            "DeepseekSparseSWAMetadata | None",
-            attn_metadata.get(self.swa_cache_layer.prefix),
-        )
-        assert swa_metadata is not None
-
-        swa_only = self.compress_ratio <= 1
-        # SWA-only layers don't allocate their own compressed KV cache.
-        self_kv_cache = self.kv_cache if not swa_only else None
-        swa_kv_cache = self.swa_cache_layer.kv_cache
-
-        self._forward_sparse_impl(
-            q=q,
-            output=output,
-            flashmla_metadata=flashmla_metadata,
-            swa_metadata=swa_metadata,
-            self_kv_cache=self_kv_cache,
-            swa_kv_cache=swa_kv_cache,
-            swa_only=swa_only,
-        )
-
-    def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
-        if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
-            assert q.dtype == torch.float8_e4m3fn
-            q = q.to(torch.bfloat16)
-        else:
-            assert q.dtype == torch.bfloat16
-        padded_heads = output.shape[1]
-        if q.shape[1] < padded_heads:
-            padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2]))
-            padded_query[:, : q.shape[1], :] = q
-            q = padded_query
-        return q.contiguous()
-
-    def _forward_decode(
-        self,
-        q: torch.Tensor,
-        kv_cache: torch.Tensor | None,
-        swa_metadata: "DeepseekSparseSWAMetadata",
-        attn_metadata: DeepseekV4FlashMLAMetadata | None,
-        swa_only: bool,
-        output: torch.Tensor,
-    ) -> None:
-        num_decodes = swa_metadata.num_decodes
-        num_decode_tokens = swa_metadata.num_decode_tokens
-
-        extra_sparse_indices = None
-        extra_sparse_lengths = None
-        if not swa_only:
-            if attn_metadata is None:
-                raise RuntimeError(
-                    "Sparse MLA metadata is required for compressed layers."
-                )
-            if swa_metadata.is_valid_token is None:
-                raise RuntimeError(
-                    "SWA validity metadata is required for compressed layers."
-                )
-            is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
-            if self.compress_ratio == 4:
-                if self.topk_indices_buffer is None:
-                    raise RuntimeError(
-                        "C4A decode requires top-k indices from the indexer."
-                    )
-                block_size = attn_metadata.block_size // self.compress_ratio
-                global_indices, extra_sparse_lengths = (
-                    compute_global_topk_indices_and_lens(
-                        self.topk_indices_buffer[:num_decode_tokens],
-                        swa_metadata.token_to_req_indices,
-                        attn_metadata.block_table[:num_decodes],
-                        block_size,
-                        is_valid,
-                    )
-                )
-                extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)
-            else:
-                extra_sparse_indices = attn_metadata.c128a_global_decode_topk_indices
-                extra_sparse_lengths = attn_metadata.c128a_decode_topk_lens
-
-        swa_indices = swa_metadata.decode_swa_indices
-        swa_lens = swa_metadata.decode_swa_lens
-        assert swa_indices is not None
-        assert swa_lens is not None
-        q = self._prepare_query(q, output)
-        swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache)
-        extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None
-        if extra_cache is not None and extra_sparse_indices is None:
-            raise RuntimeError(
-                "Compressed sparse MLA decode requires compressed sparse indices."
-            )
-        flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
-            query=q,
-            swa_kv_cache=swa_cache,
-            workspace_buffer=self._get_workspace(q.device),
-            sparse_indices=swa_indices,
-            compressed_kv_cache=extra_cache,
-            out=output,
-            bmm1_scale=self.scale,
-            sinks=self.attn_sink,
-            kv_layout="NHD",
-            swa_topk_lens=swa_lens,
-            extra_sparse_indices=extra_sparse_indices,
-            extra_sparse_topk_lens=extra_sparse_lengths,
-        )
-
-    def _forward_prefill(
-        self,
-        q: torch.Tensor,
-        compressed_k_cache: torch.Tensor | None,
-        swa_k_cache: torch.Tensor,
-        output: torch.Tensor,
-        attn_metadata: DeepseekV4FlashMLAMetadata | None,
-        swa_metadata: "DeepseekSparseSWAMetadata",
-    ) -> None:
-        swa_only = self.compress_ratio <= 1
-
-        num_prefills = swa_metadata.num_prefills
-        num_decodes = swa_metadata.num_decodes
-        num_decode_tokens = swa_metadata.num_decode_tokens
-        num_prefill_tokens = swa_metadata.num_prefill_tokens
-
-        query_start_loc_cpu = swa_metadata.query_start_loc_cpu
-        assert query_start_loc_cpu is not None
-        prefill_token_base = query_start_loc_cpu[num_decodes]
-
-        local_topk_indices: torch.Tensor | None
-        if swa_only:
-            local_topk_indices = None
-        elif self.compress_ratio == 4:
-            if self.topk_indices_buffer is None:
-                raise RuntimeError(
-                    "C4A prefill requires top-k indices from the indexer."
-                )
-            local_topk_indices = self.topk_indices_buffer[
-                num_decode_tokens : num_decode_tokens + num_prefill_tokens
-            ]
-        else:
-            if attn_metadata is None:
-                raise RuntimeError("C128A prefill metadata is missing.")
-            local_topk_indices = attn_metadata.c128a_prefill_topk_indices
-
-        extra_sparse_indices: torch.Tensor | None = None
-        extra_sparse_lengths: torch.Tensor | None = None
-        if local_topk_indices is not None:
-            if attn_metadata is None:
-                raise RuntimeError("C4A prefill metadata is missing.")
-            if swa_metadata.token_to_req_indices is None:
-                raise RuntimeError("C4A prefill request mapping is missing.")
-            if swa_metadata.is_valid_token is None:
-                raise RuntimeError("C4A prefill validity metadata is missing.")
-            prefill_token_slice = slice(
-                num_decode_tokens, num_decode_tokens + num_prefill_tokens
-            )
-            block_size = attn_metadata.block_size // self.compress_ratio
-            extra_sparse_indices, extra_sparse_lengths = (
-                compute_global_topk_indices_and_lens(
-                    local_topk_indices,
-                    swa_metadata.token_to_req_indices[prefill_token_slice],
-                    attn_metadata.block_table,
-                    block_size,
-                    swa_metadata.is_valid_token[prefill_token_slice],
-                )
-            )
-
-        assert swa_metadata.prefill_swa_indices is not None
-        assert swa_metadata.prefill_swa_lens is not None
-
-        q = self._prepare_query(q, output)
-        swa_kv_paged = self._as_sparse_cache(swa_k_cache)
-        if swa_only:
-            extra_kv_paged = None
-        else:
-            if compressed_k_cache is None:
-                raise RuntimeError(
-                    "Compressed sparse MLA layers require their compressed KV cache."
-                )
-            extra_kv_paged = self._as_sparse_cache(compressed_k_cache)
-
-        num_chunks = (
-            num_prefills + self.PREFILL_CHUNK_SIZE - 1
-        ) // self.PREFILL_CHUNK_SIZE
-        for chunk_idx in range(num_chunks):
-            chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
-            chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
-            query_start = (
-                query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
-            )
-            query_end = (
-                query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
-            )
-
-            extra_sparse_indices_chunk = (
-                extra_sparse_indices[query_start:query_end]
-                if extra_sparse_indices is not None
-                else None
-            )
-            extra_sparse_lengths_chunk = (
-                extra_sparse_lengths[query_start:query_end]
-                if extra_sparse_lengths is not None
-                else None
-            )
-
-            q_chunk = q[query_start:query_end]
-            swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end]
-            swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end]
-            if extra_kv_paged is not None and extra_sparse_indices_chunk is None:
-                raise RuntimeError(
-                    "Compressed sparse MLA prefill requires compressed sparse indices."
-                )
-            flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
-                query=q_chunk,
-                swa_kv_cache=swa_kv_paged,
-                workspace_buffer=self._get_workspace(q.device),
-                sparse_indices=swa_indices_chunk,
-                compressed_kv_cache=extra_kv_paged,
-                out=output[query_start:query_end],
-                bmm1_scale=self.scale,
-                sinks=self.attn_sink,
-                kv_layout="NHD",
-                swa_topk_lens=swa_lens_chunk,
-                extra_sparse_indices=extra_sparse_indices_chunk,
-                extra_sparse_topk_lens=extra_sparse_lengths_chunk,
-            )
@@ -60,11 +60,9 @@ from vllm.model_executor.utils import set_weight_attrs
 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
 from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
    DeepseekV4FlashInferMLAAttention,
-    DeepseekV4FlashInferSM120Attention,
 )
 from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
-from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -738,34 +736,14 @@ class DeepseekV4MoE(nn.Module):
 def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
    """Pick the CUDA sparse-MLA attention class for the configured backend.

-    The generic CUDA backend selector does not instantiate DSv4 layers directly,
-    so map generic sparse-MLA choices to the DSv4-specialized attention class.
-    Without an explicit backend, SM12 defaults to FlashInfer while the other
-    CUDA arches keep the FlashMLA path.
+    An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the
+    FlashInfer TRTLLM-gen path; otherwise the FlashMLA path is used.
    """
-    backend = vllm_config.attention_config.backend
-    device_capability = current_platform.get_device_capability()
-    if backend in (
-        AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
-        AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
+    if (
+        vllm_config.attention_config.backend
+        == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4
    ):
-        raise ValueError(
-            f"{backend.name} is not a DeepSeek V4 attention backend. "
-            "Use FLASHINFER_MLA_SPARSE_DSV4 for DeepSeek V4 FlashInfer "
-            "sparse MLA."
-        )
-    if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4:
-        if device_capability is not None and device_capability.major == 12:
-            return DeepseekV4FlashInferSM120Attention
        return DeepseekV4FlashInferMLAAttention
-    if backend in (
-        AttentionBackendEnum.FLASHMLA_SPARSE,
-        AttentionBackendEnum.FLASHMLA_SPARSE_DSV4,
-    ):
-        return DeepseekV4FlashMLAAttention
-
-    if device_capability is not None and device_capability.major == 12:
-        return DeepseekV4FlashInferSM120Attention
    return DeepseekV4FlashMLAAttention


@@ -86,10 +86,6 @@ class DeepseekV4FlashMLABackend(AttentionBackend):
    def is_sparse(cls) -> bool:
        return True

-    @classmethod
-    def supports_sink(cls) -> bool:
-        return True
-
    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability.major in [9, 10]
@@ -11,7 +11,7 @@ import platform
 from collections.abc import Callable
 from datetime import timedelta
 from functools import cache, lru_cache, wraps
-from typing import TYPE_CHECKING, NamedTuple, TypeVar
+from typing import TYPE_CHECKING, TypeVar

 import torch
 from torch.distributed import PrefixStore, ProcessGroup
@@ -31,7 +31,6 @@ if TYPE_CHECKING:
    from vllm.config import VllmConfig
    from vllm.config.cache import CacheDType
    from vllm.config.kernel import IrOpPriorityConfig
-    from vllm.v1.attention.backend import AttentionBackend
    from vllm.v1.attention.selector import AttentionSelectorConfig
 else:
    VllmConfig = None
@@ -127,11 +126,6 @@ def _get_backend_priorities(
                AttentionBackendEnum.TRITON_MLA,
                *sparse_backends,
            ]
-        elif device_capability.major == 12:
-            return [
-                AttentionBackendEnum.TRITON_MLA,
-                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
-            ]
        else:
            return [
                AttentionBackendEnum.FLASH_ATTN_MLA,
@@ -159,21 +153,6 @@ def _get_backend_priorities(
            ]


-def _backend_cls_path(backend_cls: type[AttentionBackend]) -> str:
-    module, qualname = backend_cls.full_cls_name()
-    return f"{module}.{qualname}"
-
-
-def _get_attn_backend_class(backend: AttentionBackendEnum) -> type[AttentionBackend]:
-    return backend.get_class()
-
-
-class _BackendCandidate(NamedTuple):
-    backend_class: type[AttentionBackend]
-    backend: AttentionBackendEnum
-    priority: int
-
-
 def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
    @wraps(fn)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
@@ -361,7 +340,7 @@ class CudaPlatformBase(Platform):
        attn_selector_config: AttentionSelectorConfig,
        num_heads: int | None = None,
    ) -> tuple[
-        list[_BackendCandidate],
+        list[tuple[AttentionBackendEnum, int]],
        dict[AttentionBackendEnum, tuple[int, list[str]]],
    ]:
        valid_backends_priorities = []
@@ -375,7 +354,7 @@ class CudaPlatformBase(Platform):
        )
        for priority, backend in enumerate(backend_priorities):
            try:
-                backend_class = _get_attn_backend_class(backend)
+                backend_class = backend.get_class()
                invalid_reasons_i = backend_class.validate_configuration(
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
@@ -385,9 +364,7 @@ class CudaPlatformBase(Platform):
            if invalid_reasons_i:
                invalid_reasons[backend] = (priority, invalid_reasons_i)
            else:
-                valid_backends_priorities.append(
-                    _BackendCandidate(backend_class, backend, priority)
-                )
+                valid_backends_priorities.append((backend, priority))

        return valid_backends_priorities, invalid_reasons

@@ -404,7 +381,7 @@ class CudaPlatformBase(Platform):
        # First try checking just the selected backend, if there is one.
        if selected_backend is not None:
            try:
-                backend_class = _get_attn_backend_class(selected_backend)
+                backend_class = selected_backend.get_class()
                invalid_reasons = backend_class.validate_configuration(
                    device_capability=device_capability,
                    **attn_selector_config._asdict(),
@@ -418,7 +395,7 @@ class CudaPlatformBase(Platform):
                )
            else:
                logger.info("Using %s backend.", selected_backend)
-                return _backend_cls_path(backend_class)
+                return selected_backend.get_path()

        # No selected backend or the selected backend is invalid,
        # so we try finding a valid backend.
@@ -448,13 +425,13 @@ class CudaPlatformBase(Platform):

        # We have found some valid backends. Select the one with the
        # highest priority.
-        selected_candidate = min(
-            valid_backends_priorities,
-            key=lambda candidate: candidate.priority,
+        sorted_indices = sorted(
+            range(len(valid_backends_priorities)),
+            key=lambda i: valid_backends_priorities[i][1],
        )
-        selected_backend_class = selected_candidate.backend_class
-        selected_backend = selected_candidate.backend
-        selected_priority = selected_candidate.priority
+        selected_index = sorted_indices[0]
+        selected_backend = valid_backends_priorities[selected_index][0]
+        selected_priority = valid_backends_priorities[selected_index][1]

        # If the user specified --block-size (but not --attention-backend),
        # check whether that constraint precluded any higher-priority backends.
@@ -480,14 +457,10 @@ class CudaPlatformBase(Platform):
        logger.info_once(
            "Using %s attention backend out of potential backends: %s.",
            selected_backend.name,
-            "["
-            + ", ".join(
-                f"'{candidate.backend.name}'" for candidate in valid_backends_priorities
-            )
-            + "]",
+            "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
        )

-        return _backend_cls_path(selected_backend_class)
+        return selected_backend.get_path()

    @classmethod
    def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]:
@@ -662,11 +635,7 @@ class CudaPlatformBase(Platform):
    @classmethod
    def support_deep_gemm(cls) -> bool:
        """Currently, only Hopper and Blackwell GPUs are supported."""
-        return (
-            cls.is_device_capability(90)
-            or cls.is_device_capability_family(100)
-            or cls.is_device_capability_family(120)
-        )
+        return cls.is_device_capability(90) or cls.is_device_capability_family(100)

    @classmethod
    def is_integrated_gpu(cls, device_id: int = 0) -> bool:
@@ -27,10 +27,8 @@ logger = init_logger(__name__)
 try:
    from amdsmi import (
        AmdSmiException,
-        AmdSmiMemoryType,
        amdsmi_get_gpu_asic_info,
        amdsmi_get_gpu_device_uuid,
-        amdsmi_get_gpu_memory_total,
        amdsmi_get_processor_handles,
        amdsmi_init,
        amdsmi_shut_down,
@@ -169,14 +167,6 @@ def _query_gcn_arch_from_amdsmi() -> str:
    raise RuntimeError("amdsmi did not return valid GCN arch")


-@with_amdsmi_context
-def _query_total_memory_from_amdsmi(physical_device_id: int) -> int:
-    """Query total VRAM (bytes) from amdsmi. Raises if not available."""
-    handles = amdsmi_get_processor_handles()
-    handle = handles[physical_device_id]
-    return amdsmi_get_gpu_memory_total(handle, AmdSmiMemoryType.VRAM)
-
-
 def _get_gcn_arch() -> str:
    """
    Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda.
@@ -736,22 +726,8 @@ class RocmPlatform(Platform):

    @classmethod
    def get_device_total_memory(cls, device_id: int = 0) -> int:
-        # Query total VRAM via amdsmi so we don't initialize a HIP context in
-        # the calling process. torch.cuda.get_device_properties() creates a
-        # HIP context, which makes vLLM fall back from `fork` to `spawn` for
-        # worker processes. Keeping this query context-free preserves `fork`
-        # where it is otherwise valid (e.g. out-of-tree models registered in
-        # the parent process).
-        try:
-            physical_device_id = cls.device_id_to_physical_device_id(device_id)
-            return _query_total_memory_from_amdsmi(physical_device_id)
-        except Exception as e:
-            logger.debug("Failed to get total memory via amdsmi: %s", e)
-            logger.warning_once(
-                "Failed to get total memory via amdsmi, falling back to "
-                "torch.cuda. This will initialize CUDA."
-            )
-        return torch.cuda.get_device_properties(device_id).total_memory
+        device_props = torch.cuda.get_device_properties(device_id)
+        return device_props.total_memory

    @classmethod
    def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
@@ -5,7 +5,6 @@
 Users of vLLM should always import **only** these wrappers.
 """

-import contextlib
 import functools
 import importlib
 import os
@@ -38,10 +37,7 @@ def should_auto_disable_deep_gemm(model_type: str | None) -> bool:
    """
    if model_type is None:
        return False
-    if not (
-        current_platform.is_device_capability_family(100)
-        or current_platform.is_device_capability_family(120)
-    ):
+    if not current_platform.is_device_capability_family(100):
        return False
    return model_type in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES

@@ -75,10 +71,7 @@ class DeepGemmQuantScaleFMT(Enum):

        cls._oracle_cache = (  # type: ignore
            cls.UE8M0
-            if (
-                current_platform.is_device_capability_family(100)
-                or current_platform.is_device_capability_family(120)
-            )
+            if current_platform.is_device_capability_family(100)
            else cls.FLOAT32_CEIL_UE8M0
        )

@@ -145,15 +138,7 @@ _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
 _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
 _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
-_get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = (
-    None
-)
 _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None
-_pack_ue8m0_to_int_impl: Callable[..., Any] | None = None
-_get_mn_major_tma_aligned_packed_ue8m0_tensor_impl: Callable[..., Any] | None = None
-_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl: (
-    Callable[..., Any] | None
-) = None


@functools.cache
@@ -218,11 +203,7 @@ def _lazy_init() -> None:
    global _tf32_hc_prenorm_gemm_impl
    global _get_mn_major_tma_aligned_tensor_impl
    global _get_mk_alignment_for_contiguous_layout_impl
-    global _get_theoretical_mk_alignment_for_contiguous_layout_impl
    global _transform_sf_into_required_layout_impl
-    global _pack_ue8m0_to_int_impl
-    global _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl
-    global _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl
    # fast path
    if (
        _cublaslt_gemm_nt_impl is not None
@@ -237,9 +218,6 @@ def _lazy_init() -> None:
        or _tf32_hc_prenorm_gemm_impl is not None
        or _get_mk_alignment_for_contiguous_layout_impl is not None
        or _transform_sf_into_required_layout_impl is not None
-        or _pack_ue8m0_to_int_impl is not None
-        or _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None
-        or _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None
    ):
        return

@@ -280,19 +258,9 @@ def _lazy_init() -> None:
    _get_mk_alignment_for_contiguous_layout_impl = getattr(
        _dg, "get_mk_alignment_for_contiguous_layout", None
    )
-    _get_theoretical_mk_alignment_for_contiguous_layout_impl = getattr(
-        _dg, "get_theoretical_mk_alignment_for_contiguous_layout", None
-    )
    _transform_sf_into_required_layout_impl = getattr(
        _dg, "transform_sf_into_required_layout", None
    )
-    _pack_ue8m0_to_int_impl = getattr(_dg, "pack_ue8m0_to_int", None)
-    _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr(
-        _dg, "get_mn_major_tma_aligned_packed_ue8m0_tensor", None
-    )
-    _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr(
-        _dg, "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", None
-    )
    DeepGemmQuantScaleFMT.init_oracle_cache()


@@ -312,6 +280,7 @@ def set_num_sms(num_sms: int) -> None:
    dg.set_num_sms(num_sms)


+@functools.cache
 def get_mk_alignment_for_contiguous_layout() -> list[int]:
    _lazy_init()
    if _get_mk_alignment_for_contiguous_layout_impl is None:
@@ -320,70 +289,6 @@ def get_mk_alignment_for_contiguous_layout() -> list[int]:
    return [mk_align_size, mk_align_size]


-def get_theoretical_mk_alignment_for_contiguous_layout(
-    expected_m: int | None = None,
-    num_groups: int | None = None,
-) -> int:
-    """Per-call optimal M alignment for grouped contiguous GEMMs.
-
-    `expected_m` is the TOTAL routed tokens (sum across experts, typically
-    M × num_topk). `num_groups` is the number of experts on this rank.
-    The helper divides to recover per-expert em and picks an alignment based
-    on data-driven thresholds (see deep_gemm runtime.hpp comments).
-
-    Older callers that omit `num_groups` are interpreted as passing already
-    per-expert em (legacy behaviour preserved for backward compat).
-    """
-    _lazy_init()
-    if _get_theoretical_mk_alignment_for_contiguous_layout_impl is None:
-        return _missing()
-    if num_groups is None:
-        return _get_theoretical_mk_alignment_for_contiguous_layout_impl(expected_m)
-    if num_groups <= 0:
-        raise ValueError(f"num_groups must be positive, got {num_groups}")
-    try:
-        return _get_theoretical_mk_alignment_for_contiguous_layout_impl(
-            expected_m, num_groups
-        )
-    except TypeError:
-        per_group_m = None if expected_m is None else cdiv(expected_m, num_groups)
-        return _get_theoretical_mk_alignment_for_contiguous_layout_impl(per_group_m)
-
-
-def set_mk_alignment_for_contiguous_layout(value: int) -> None:
-    """Set DeepGEMM's BLOCK_M cap for grouped contiguous GEMMs.
-
-    The DG heuristic constrains BLOCK_M ≤ this value when picking a kernel
-    layout. Use this in concert with `compute_aligned_M_and_alignment`'s
-    per-call alignment so the workspace's per-expert padding matches the
-    kernel's BLOCK_M; a mismatch leads to the scheduler reading the wrong
-    expert_id from `m_indices` at `m_block_idx * BLOCK_M` stride and
-    OOB-indexing the B-weights tensor (manifests as IMA under CUDA-graph
-    replay).
-    """
-    _lazy_init()
-    dg = _import_deep_gemm()
-    if dg is None:
-        raise RuntimeError("DeepGEMM is not available")
-    dg.set_mk_alignment_for_contiguous_layout(value)
-
-
-@contextlib.contextmanager
-def mk_alignment_scope(value: int):
-    """Temporarily set DeepGEMM's BLOCK_M cap, restoring on exit.
-
-    Use around a sequence of grouped-contiguous GEMM calls whose workspace
-    is padded to `value` (typically the per_call_align returned by
-    `compute_aligned_M_and_alignment`).
-    """
-    prev = get_mk_alignment_for_contiguous_layout()[0]
-    set_mk_alignment_for_contiguous_layout(value)
-    try:
-        yield
-    finally:
-        set_mk_alignment_for_contiguous_layout(prev)
-
-
 def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
    """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
    _lazy_init()
@@ -392,48 +297,6 @@ def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
    return _get_mn_major_tma_aligned_tensor_impl(x)


-def pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor:
-    """Pack 4 UE8M0 (uint8) scales into one int32.
-
-    DeepGEMM's SM100/SM120 FP8/FP4 kernels accept either ``float32`` scales
-    (legacy format, 4 B/scale) or ``int32`` packed UE8M0 scales (1 B/scale
-    after 4:1 packing — 4× smaller than the legacy fp32 representation).
-    """
-    _lazy_init()
-    if _pack_ue8m0_to_int_impl is None:
-        return _missing()
-    return _pack_ue8m0_to_int_impl(x)
-
-
-def get_mn_major_tma_aligned_packed_ue8m0_tensor(x: torch.Tensor) -> torch.Tensor:
-    """Pack UE8M0 (uint8) → int32 with the MN-major TMA-aligned layout the
-    DeepGEMM kernels consume directly. 16× smaller than the fp32 legacy SF
-    format. Use for non-grouped 2D scale tensors.
-    """
-    _lazy_init()
-    if _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None:
-        return _missing()
-    return _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl(x)
-
-
-def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(
-    sf: torch.Tensor,
-    ks_tensor: torch.Tensor,
-    ks: list[int],
-    gran_k: int,
-) -> torch.Tensor:
-    """Grouped (3D, expert-batched) variant of
-    ``get_mn_major_tma_aligned_packed_ue8m0_tensor``. Use for MoE weight
-    scale tensors of shape ``(num_experts, mn, k_scale)``.
-    """
-    _lazy_init()
-    if _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None:
-        return _missing()
-    return _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl(
-        sf, ks_tensor, ks, gran_k
-    )
-
-
 def cublaslt_gemm_nt(*args, **kwargs):
    _lazy_init()
    if _cublaslt_gemm_nt_impl is None:
@@ -738,8 +601,4 @@ __all__ = [
    "should_use_deepgemm_for_fp8_linear",
    "get_col_major_tma_aligned_tensor",
    "get_mk_alignment_for_contiguous_layout",
-    "get_theoretical_mk_alignment_for_contiguous_layout",
-    "pack_ue8m0_to_int",
-    "get_mn_major_tma_aligned_packed_ue8m0_tensor",
-    "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor",
 ]
@@ -72,10 +72,11 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
    )


-def _missing_sparse_mla(*_: Any, **__: Any) -> NoReturn:
+def _missing_dsv4_sparse_mla(*_: Any, **__: Any) -> NoReturn:
    raise RuntimeError(
-        "FlashInfer sparse MLA decode APIs are not available. "
-        "Install a FlashInfer build that includes sparse MLA decode support."
+        "flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4 is not available. "
+        "Install a FlashInfer build that includes DeepSeek V4 sparse MLA "
+        "TRTLLM-GEN support."
    )


@@ -148,18 +149,14 @@ flashinfer_b12x_fused_moe = _lazy_import_wrapper(
 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
    "flashinfer", "trtllm_fp4_block_scale_moe"
 )
-flashinfer_trtllm_batch_decode_with_kv_cache_mla = _lazy_import_wrapper(
-    "flashinfer.decode",
-    "trtllm_batch_decode_with_kv_cache_mla",
-    fallback_fn=_missing_sparse_mla,
-)
+# DeepSeek V4 sparse MLA TRTLLM-GEN decode launcher (public wrapper). Handles
+# the SWA + compressed KV pools, the concatenated sparse-index matrix, and
+# per-tensor FP8 / BF16 inputs with BF16 output.
 flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper(
-    "flashinfer.decode",
+    "flashinfer.mla",
    "trtllm_batch_decode_sparse_mla_dsv4",
-    fallback_fn=_missing_sparse_mla,
+    fallback_fn=_missing_dsv4_sparse_mla,
 )
-
-
 # Special case for autotune since it returns a context manager
 autotune = _lazy_import_wrapper(
    "flashinfer.autotuner",
@@ -212,26 +209,6 @@ def has_flashinfer_moe() -> bool:
    )


-@functools.cache
-def has_flashinfer_sparse_mla_sm120() -> bool:
-    """Return ``True`` if FlashInfer sparse MLA decode support is available."""
-    if not has_flashinfer():
-        return False
-    try:
-        from flashinfer.autotuner import autotune
-        from flashinfer.decode import (
-            trtllm_batch_decode_sparse_mla_dsv4,
-            trtllm_batch_decode_with_kv_cache_mla,
-        )
-    except ImportError:
-        return False
-    return (
-        callable(trtllm_batch_decode_sparse_mla_dsv4)
-        and callable(trtllm_batch_decode_with_kv_cache_mla)
-        and callable(autotune)
-    )
-
-
@functools.cache
 def has_flashinfer_cutedsl() -> bool:
    """Return ``True`` if FlashInfer cutedsl module is available."""
@@ -1011,7 +988,6 @@ __all__ = [
    "flashinfer_b12x_fused_moe",
    "flashinfer_convert_sf_to_mma_layout",
    "trtllm_fp4_block_scale_moe",
-    "flashinfer_trtllm_batch_decode_with_kv_cache_mla",
    "flashinfer_trtllm_batch_decode_sparse_mla_dsv4",
    "autotune",
    "has_flashinfer_moe",
@@ -436,7 +436,7 @@ class CommonAttentionMetadata:
    positions: torch.Tensor | None = None
    """(num_actual_tokens,) token positions.  Optional; set when the caller
    has positions available so that builders can pre-compute position-dependent
-    sparse metadata for DeepSeek V4 C128A layers."""
+    metadata (e.g. C128A topk indices for DeepSeek V4)."""

    is_prefilling: torch.Tensor | None = None
    """(batch_size,) bool tensor: True if request is still in prefill phase
@@ -1,12 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FlashInfer sparse MLA attention backend."""
+"""FlashInfer MLA Sparse Attention Backend.
+
+This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k
+for models like DeepSeek-V3.2 that use index-based sparse attention.
+
+For sparse MLA:
+- block_tables shape changes from [batch_size, max_num_blocks] (dense)
+  to [batch_size, q_len_per_request, sparse_mla_top_k] (sparse)
+- The sparse indices represent physical cache slot positions to attend to
+- sparse_mla_top_k parameter must be set to the topk value
+"""

 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar

 import numpy as np
 import torch
+from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla

 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
@@ -41,13 +52,34 @@ logger = init_logger(__name__)
 FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024


-class _FlashInferMLASparseBackendBase(AttentionBackend):
-    """Common metadata for concrete FlashInfer sparse MLA backends."""
+class FlashInferMLASparseBackend(AttentionBackend):
+    """FlashInfer MLA backend with sparse attention support.
+
+    This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k
+    for models like DeepSeek-V3.2 that use index-based sparse attention.
+    """
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
+        "auto",
+        "float16",
+        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
+    ]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [32, 64]

    @staticmethod
    def get_name() -> str:
        return "FLASHINFER_MLA_SPARSE"

+    @staticmethod
+    def get_impl_cls() -> type["FlashInferMLASparseImpl"]:
+        return FlashInferMLASparseImpl
+
    @staticmethod
    def get_builder_cls() -> type["FlashInferMLASparseMetadataBuilder"]:
        return FlashInferMLASparseMetadataBuilder
@@ -64,29 +96,9 @@ class _FlashInferMLASparseBackendBase(AttentionBackend):
    def is_sparse(cls) -> bool:
        return True

-
-class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
-    """FlashInfer sparse MLA backend using the TRTLLM-gen launcher."""
-
-    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
-    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
-        "auto",
-        "float16",
-        "bfloat16",
-        "fp8",
-        "fp8_e4m3",
-    ]
-
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
-        return [32, 64]
-
-    @staticmethod
-    def get_impl_cls() -> type[SparseMLAAttentionImpl]:
-        return FlashInferMLASparseImpl
-
    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        # FlashInfer sparse MLA targets Blackwell (SM 10.x)
        return capability.major == 10

    @classmethod
@@ -102,15 +114,10 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
        use_mm_prefix: bool,
        device_capability: DeviceCapability,
    ) -> str | None:
+        # FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192]
        from vllm.config import get_current_vllm_config

        vllm_config = get_current_vllm_config()
-        if kv_cache_dtype == "fp8_ds_mla":
-            return (
-                "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype"
-            )
-
-        # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192].
        if vllm_config.model_config is not None:
            hf_text_config = vllm_config.model_config.hf_text_config
            qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
@@ -139,102 +146,6 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
        return "HND"


-class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
-    """FlashInfer sparse MLA backend for SM120."""
-
-    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
-    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
-        "auto",
-        "fp8",
-        "fp8_e4m3",
-        "fp8_ds_mla",
-    ]
-
-    @staticmethod
-    def get_name() -> str:
-        return "FLASHINFER_MLA_SPARSE_SM120"
-
-    @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
-        return [64, 256]
-
-    @staticmethod
-    def get_impl_cls() -> type[SparseMLAAttentionImpl]:
-        from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
-            FlashInferMLASparseSM120Impl,
-        )
-
-        return FlashInferMLASparseSM120Impl
-
-    @classmethod
-    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
-        return capability.major == 12
-
-    @classmethod
-    def supports_combination(
-        cls,
-        head_size: int,
-        dtype: torch.dtype,
-        kv_cache_dtype: CacheDType | None,
-        block_size: int | None,
-        use_mla: bool,
-        has_sink: bool,
-        use_sparse: bool,
-        use_mm_prefix: bool,
-        device_capability: DeviceCapability,
-    ) -> str | None:
-        from vllm.config import get_current_vllm_config
-        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
-
-        if not has_flashinfer_sparse_mla_sm120():
-            return (
-                "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
-                "sparse MLA decode API"
-            )
-        if dtype != torch.bfloat16:
-            return "dtype not supported"
-        if kv_cache_dtype not in (
-            None,
-            "auto",
-            "fp8",
-            "fp8_e4m3",
-            "fp8_ds_mla",
-        ):
-            return "kv_cache_dtype not supported"
-        vllm_config = get_current_vllm_config()
-        if vllm_config.model_config is not None:
-            hf_text_config = vllm_config.model_config.hf_text_config
-            index_topk = getattr(hf_text_config, "index_topk", None)
-            if index_topk is None:
-                return (
-                    "FLASHINFER_MLA_SPARSE_SM120 requires a model with "
-                    "index_topk config"
-                )
-            if int(index_topk) != 2048:
-                return (
-                    "FLASHINFER_MLA_SPARSE_SM120 requires index_topk=2048; "
-                    f"got {index_topk}"
-                )
-        return None
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,  # assumed to be 1 for MLA
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"):
-            # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE.
-            return (num_blocks, block_size, 656)
-        return (num_blocks, block_size, head_size)
-
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return None
-
-
@dataclass
 class FlashInferMLASparseMetadata(AttentionMetadata):
    """Attention metadata for FlashInfer MLA Sparse backend."""
@@ -442,8 +353,6 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
            if is_quantized_kv_cache(self.kv_cache_dtype):
                self.bmm2_scale *= layer._k_scale_float

-        from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
-
        o = trtllm_batch_decode_with_kv_cache_mla(
            query=q.unsqueeze(1),
            kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -1,155 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""SM120 implementation variant for ``FLASHINFER_MLA_SPARSE_SM120``."""
-
-from typing import TYPE_CHECKING, cast
-
-import torch
-
-from vllm.v1.attention.backend import (
-    AttentionLayer,
-    AttentionType,
-    SparseMLAAttentionImpl,
-)
-from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
-    FlashInferMLASparseMetadata,
-    _get_workspace_buffer,
-)
-from vllm.v1.attention.backends.mla.sparse_utils import (
-    triton_convert_req_index_to_global_index,
-)
-
-if TYPE_CHECKING:
-    from vllm.model_executor.models.deepseek_v2 import Indexer
-
-
-def _kv_scale_format_for_model(model_type: str | None) -> str:
-    if model_type is not None and model_type.startswith("glm"):
-        return "arbitrary_fp32"
-    return "pow2_fp32"
-
-
-class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
-    """SM120 FlashInfer sparse-MLA implementation."""
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        scale: float,
-        num_kv_heads: int,
-        alibi_slopes: list[float] | None,
-        sliding_window: int | None,
-        kv_cache_dtype: str,
-        logits_soft_cap: float | None,
-        attn_type: str,
-        kv_sharing_target_layer_name: str | None,
-        indexer: "Indexer | None" = None,
-        **mla_args,
-    ) -> None:
-        if any([alibi_slopes, sliding_window, logits_soft_cap]):
-            raise NotImplementedError(
-                "FLASHINFER_MLA_SPARSE_SM120 does not support alibi_slopes / "
-                "sliding_window / logits_soft_cap"
-            )
-        if attn_type != AttentionType.DECODER:
-            raise NotImplementedError(
-                "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
-            )
-
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.scale = float(scale)
-        self.num_kv_heads = num_kv_heads
-        self.kv_cache_dtype = kv_cache_dtype
-        if self.kv_cache_dtype != "fp8_ds_mla":
-            raise NotImplementedError(
-                "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
-                f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
-            )
-
-        self.kv_lora_rank: int = mla_args["kv_lora_rank"]
-        self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
-        self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
-        from vllm.config import get_current_vllm_config
-
-        vllm_config = get_current_vllm_config()
-        model_type = None
-        if vllm_config.model_config is not None:
-            model_type = getattr(
-                vllm_config.model_config.hf_text_config, "model_type", None
-            )
-        self.kv_scale_format = _kv_scale_format_for_model(model_type)
-
-        assert indexer is not None, (
-            "FLASHINFER_MLA_SPARSE_SM120 requires a sparse-MLA indexer "
-            "(model with index_topk in its config)."
-        )
-        self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
-        from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
-
-        if not has_flashinfer_sparse_mla_sm120():
-            raise RuntimeError(
-                "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
-                "sparse MLA decode API."
-            )
-        assert self.topk_indices_buffer is not None
-
-        self.supports_quant_query_input = False
-        self._workspace_buffer: torch.Tensor | None = None
-
-    def forward_mqa(
-        self,
-        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-        kv_c_and_k_pe_cache: torch.Tensor,
-        attn_metadata: FlashInferMLASparseMetadata,
-        layer: AttentionLayer,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if isinstance(q, tuple):
-            q = torch.cat(q, dim=-1)
-
-        num_actual_toks = q.shape[0]
-
-        assert self.topk_indices_buffer is not None
-        topk_indices = self.topk_indices_buffer[:num_actual_toks]
-
-        topk_indices_physical = cast(
-            torch.Tensor,
-            triton_convert_req_index_to_global_index(
-                attn_metadata.req_id_per_token[:num_actual_toks],
-                attn_metadata.block_table,
-                topk_indices,
-                BLOCK_SIZE=attn_metadata.block_size,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
-            ),
-        )
-
-        output = q.new_empty(
-            (num_actual_toks, self.num_heads, self.kv_lora_rank),
-            dtype=q.dtype,
-        )
-
-        if self._workspace_buffer is None:
-            self._workspace_buffer = _get_workspace_buffer(q.device)
-
-        from vllm.utils.flashinfer import (
-            flashinfer_trtllm_batch_decode_with_kv_cache_mla,
-        )
-
-        out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
-            query=q.unsqueeze(1),
-            kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
-            workspace_buffer=self._workspace_buffer,
-            qk_nope_head_dim=self.qk_nope_head_dim,
-            kv_lora_rank=self.kv_lora_rank,
-            qk_rope_head_dim=self.qk_rope_head_dim,
-            block_tables=topk_indices_physical.unsqueeze(1),
-            seq_lens=None,
-            max_seq_len=attn_metadata.topk_tokens,
-            out=output.unsqueeze(1),
-            bmm1_scale=self.scale,
-            bmm2_scale=1.0,
-            sparse_mla_top_k=attn_metadata.topk_tokens,
-            kv_scale_format=self.kv_scale_format,
-        )
-        return out.squeeze(1), None
@@ -74,14 +74,14 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
        # determines the SWA block size of 64 tokens per block.
        # TODO(yifan): make SWA block size automatically determined and configurable.
        self.block_size = 64
-        # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
-        # contiguous full-cache layout.
+        # uint8: legacy FlashMLA UE8M0 paged layout. bfloat16 / float8_e4m3fn:
+        # FlashInfer contiguous full-cache layout.
        assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
-        # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
-        # bf16/fp8 cache uses the natural element-size page.
-        uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
+        # FlashMLA's UE8M0 paged layout needs 576B alignment; FlashInfer's
+        # contiguous bf16/fp8 cache uses the natural element-size page.
+        is_flashmla = self.cache_config.cache_dtype == "fp8_ds_mla"
        return SlidingWindowMLASpec(
            block_size=self.block_size,
            num_kv_heads=1,
@@ -89,7 +89,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
            dtype=self.dtype,
            sliding_window=self.window_size,
            cache_dtype_str=self.cache_config.cache_dtype,
-            alignment=576 if uses_fp8_ds_mla_layout else None,
+            alignment=576 if is_flashmla else None,
            model_version="deepseek_v4",
        )

@@ -164,11 +164,6 @@ class DeepseekSparseSWAMetadata:
    token_to_req_indices: torch.Tensor | None = None  # [num_tokens]
    decode_swa_indices: torch.Tensor | None = None  # [num_decode_tokens, window_size]
    decode_swa_lens: torch.Tensor | None = None  # [num_decode_tokens]
-    # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill).
-    prefill_swa_indices: torch.Tensor | None = (
-        None  # [num_prefill_tokens, 1, window_size]
-    )
-    prefill_swa_lens: torch.Tensor | None = None  # [num_prefill_tokens]

    # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
    num_decodes: int = 0
@@ -348,20 +343,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
            dtype=torch.int32,
            device=self.device,
        )
-        # Allocated unconditionally — consumer picks paged-direct vs dequant
-        # at call time.
-        self.prefill_swa_indices = torch.zeros(
-            max_tokens,
-            1,
-            self.window_size,
-            dtype=torch.int32,
-            device=self.device,
-        )
-        self.prefill_swa_lens = torch.zeros(
-            max_tokens,
-            dtype=torch.int32,
-            device=self.device,
-        )
        self.is_valid_token = torch.zeros(
            max_tokens,
            dtype=torch.bool,
@@ -421,29 +402,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
                block_table,
                block_table.stride(0),
                self.block_size,
-                token_offset=0,
-                TRITON_BLOCK_SIZE=1024,
-            )
-
-        # Prefill SWA indices live in paged coordinates. `token_offset` lets
-        # the kernel read is_valid_token / token_to_req_indices at absolute
-        # prefill positions while writing output starting at index 0.
-        if num_prefill_tokens > 0:
-            prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
-            prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
-            _compute_swa_indices_and_lens_kernel[(num_prefill_tokens,)](
-                prefill_swa_indices,
-                prefill_swa_indices.stride(0),
-                prefill_swa_lens,
-                self.window_size,
-                query_start_loc,
-                seq_lens,
-                token_to_req_indices,
-                is_valid_token,
-                block_table,
-                block_table.stride(0),
-                self.block_size,
-                token_offset=num_decode_tokens,
                TRITON_BLOCK_SIZE=1024,
            )

@@ -473,16 +431,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
            token_to_req_indices=token_to_req_indices,
            decode_swa_indices=self.decode_swa_indices[:num_decode_tokens],
            decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
-            prefill_swa_indices=(
-                self.prefill_swa_indices[:num_prefill_tokens]
-                if num_prefill_tokens > 0
-                else None
-            ),
-            prefill_swa_lens=(
-                self.prefill_swa_lens[:num_prefill_tokens]
-                if num_prefill_tokens > 0
-                else None
-            ),
            block_size=self.block_size,
            num_decodes=num_decodes,
            num_prefills=num_prefills,
@@ -517,7 +465,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
            num_decode_tokens == 0
            or current_platform.is_rocm()
            or current_platform.is_xpu()
-            or current_platform.is_device_capability_family(120)
        ):
            return out
        for layer_type in self._layer_types:
@@ -542,7 +489,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
        Returns a dict of keyword arguments to pass to the
        DeepseekSparseSWAMetadata constructor.

-        Note: C128A sparse metadata is computed by the FlashMLASparse builder
+        Note: C128A topk indices are computed by the FlashMLASparse builder
        (which owns the C128A block_table), not here.
        """
        result: dict[str, torch.Tensor | int | None] = {}
@@ -592,14 +539,10 @@ def _compute_prefill_metadata_kernel(
    """Compute prefill gather_lens in a single pass."""
    offset = tl.arange(0, BLOCK_SIZE)
    mask = offset < num_prefills
-    # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
-    # masked-off lanes even though the load mask gates the actual read, so
-    # clamp the offset. Caller guarantees num_prefills > 0.
-    safe_offset = tl.minimum(offset, num_prefills - 1)

-    seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
-    qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
-    qsl_end = tl.load(query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask)
+    seq_len = tl.load(seq_lens_ptr + num_decodes + offset, mask=mask)
+    qsl_start = tl.load(query_start_loc_ptr + num_decodes + offset, mask=mask)
+    qsl_end = tl.load(query_start_loc_ptr + num_decodes + offset + 1, mask=mask)

    query_len = qsl_end - qsl_start
    prefix_len = seq_len - query_len
@@ -608,7 +551,7 @@ def _compute_prefill_metadata_kernel(
    tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)


-@triton.jit(do_not_specialize=["token_offset"])
+@triton.jit
 def _compute_swa_indices_and_lens_kernel(
    swa_indices_ptr,
    swa_indices_stride,
@@ -621,14 +564,12 @@ def _compute_swa_indices_and_lens_kernel(
    block_table_ptr,
    block_table_stride,
    block_size,
-    token_offset,
    TRITON_BLOCK_SIZE: tl.constexpr,
 ):
-    pid = tl.program_id(0)
-    token_idx = pid + token_offset
+    token_idx = tl.program_id(0)
    is_valid = tl.load(is_valid_token_ptr + token_idx)
    if not is_valid:
-        tl.store(swa_lens_ptr + pid, 0)
+        tl.store(swa_lens_ptr + token_idx, 0)
        return

    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
@@ -645,7 +586,7 @@ def _compute_swa_indices_and_lens_kernel(
    end_pos = pos + 1

    swa_len = end_pos - start_pos
-    tl.store(swa_lens_ptr + pid, swa_len)
+    tl.store(swa_lens_ptr + token_idx, swa_len)

    for i in range(0, window_size, TRITON_BLOCK_SIZE):
        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
@@ -661,7 +602,7 @@ def _compute_swa_indices_and_lens_kernel(

        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
        tl.store(
-            swa_indices_ptr + pid * swa_indices_stride + offset,
+            swa_indices_ptr + token_idx * swa_indices_stride + offset,
            slot_ids,
            mask=offset < window_size,
        )
@@ -71,11 +71,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
    )
    FLASHINFER_MLA_SPARSE = (
        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
-        "FlashInferMLASparseTRTLLMBackend"
-    )
-    FLASHINFER_MLA_SPARSE_SM120 = (
-        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
-        "FlashInferMLASparseSM120Backend"
+        "FlashInferMLASparseBackend"
    )
    TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
    CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"
@@ -73,7 +73,7 @@ def is_flashmla_sparse_supported() -> tuple[bool, str | None]:
    ):
        return (
            False,
-            "FlashMLA Sparse is only supported on Hopper and Blackwell DC devices.",
+            "FlashMLA Sparse is only supported on Hopper and Blackwell devices.",
        )
    return True, None

@@ -1872,11 +1872,6 @@ class Scheduler(SchedulerInterface):
        if not cached_encoder_input_ids:
            return

-        # Defer the free by the drafter's look-ahead so an entry stays
-        # referenced until the drafter's +1 read has also passed it, mirroring
-        # the shift the encoder scheduling path applies.
-        spec_lookahead = 1 if self.use_eagle else 0
-
        # Here, we use list(set) to avoid modifying the set while iterating
        # over it.
        for input_id in list(cached_encoder_input_ids):
@@ -1889,12 +1884,13 @@ class Scheduler(SchedulerInterface):
                # KVs have been calculated and cached already.
                self.encoder_cache_manager.free_encoder_input(request, input_id)
            elif (
-                start_pos + num_tokens + spec_lookahead
+                start_pos + num_tokens
                <= request.num_computed_tokens - request.num_output_placeholders
            ):
-                # Processed, stored in the decoder KV cache, and far enough past
-                # the placeholder range (plus the drafter's look-ahead) that no
-                # rejection or drafter gather can reference it.
+                # The encoder output is already processed and stored in the
+                # decoder's KV cache, and progress is far enough past the
+                # placeholder range that no pending draft-token rejection can
+                # roll num_computed_tokens back into it.
                self.encoder_cache_manager.free_encoder_input(request, input_id)

    def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
@@ -12,8 +12,6 @@ import torch
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
-    CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
-    CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
    BatchMemcpyParams,
    build_params,
    copy_blocks,
@@ -45,20 +43,8 @@ class DmaCopyBackend:
        self._load_stream = load_stream
        self._store_stream = store_stream

-        # Stores read the live KV cache -> STREAM (paired with the compute-done
-        # wait in get_finished); loads read stable pinned host memory -> ANY.
-        self._store_params = build_params(
-            gpu_caches,
-            cpu_caches,
-            store_stream,
-            src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
-        )
-        self._load_params = build_params(
-            cpu_caches,
-            gpu_caches,
-            load_stream,
-            src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
-        )
+        self._store_params = build_params(gpu_caches, cpu_caches, store_stream)
+        self._load_params = build_params(cpu_caches, gpu_caches, load_stream)

        self._queue = queue.SimpleQueue()
        self._thread = threading.Thread(
@@ -75,20 +61,11 @@ class DmaCopyBackend:
        is_store: bool,
        event_idx: int,
        events_list: list[tuple[int, torch.Event]],
-        wait_event: torch.Event | None = None,
    ) -> None:
        params = self._store_params if is_store else self._load_params
        assert params is not None and self._queue is not None
        self._queue.put(
-            (
-                src_blocks,
-                dst_blocks,
-                params,
-                is_store,
-                event_idx,
-                events_list,
-                wait_event,
-            )
+            (src_blocks, dst_blocks, params, is_store, event_idx, events_list)
        )

    def shutdown(self) -> None:
@@ -112,19 +89,9 @@ class DmaCopyBackend:
            item = q.get()
            if item is None:
                return
-            (
-                src_blocks,
-                dst_blocks,
-                params,
-                is_store,
-                event_idx,
-                events_list,
-                wait_event,
-            ) = item
-            stream = store_stream if is_store else load_stream
-            if wait_event is not None:
-                stream.wait_event(wait_event)
+            src_blocks, dst_blocks, params, is_store, event_idx, events_list = item
            copy_blocks(src_blocks, dst_blocks, params)
+            stream = store_stream if is_store else load_stream
            event = torch.Event()
            event.record(stream)
            events_list.append((event_idx, event))
@@ -13,12 +13,6 @@ from vllm.platforms import current_platform

 logger = init_logger(__name__)

-# CUmemcpySrcAccessOrder values (CUDA driver API). STREAM(1): source read in
-# stream order, safe when the source may still be written. ANY(3): source may
-# be read early, only safe for a stable source (e.g. pinned host memory).
-CU_MEMCPY_SRC_ACCESS_ORDER_STREAM = 1
-CU_MEMCPY_SRC_ACCESS_ORDER_ANY = 3
-

 def pin_tensor(tensor: torch.Tensor) -> None:
    """Pin a CPU tensor via cudaHostRegister.
@@ -112,8 +106,8 @@ class BatchMemcpyParams(NamedTuple):
    dst_bases: np.ndarray  # [num_layers] uint64
    bpb: np.ndarray  # [num_layers] uint64 — bytes per block
    num_layers: int
-    # CUDA only: one attributes entry carrying srcAccessOrder. Unused on ROCm
-    # (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
+    # CUDA only: one attributes entry with srcAccessOrder=ANY. Unused on
+    # ROCm (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
    attrs: _CUmemcpyAttributes
    attrs_idx: ctypes.c_size_t
    # NOTE: cuMemcpyBatchAsync_v2() removed fail_idx field, but we use
@@ -126,7 +120,6 @@ def build_params(
    src_caches: dict[str, torch.Tensor],
    dst_caches: dict[str, torch.Tensor],
    stream: torch.cuda.Stream,
-    src_access_order: int = CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
 ) -> BatchMemcpyParams:
    global _batch_memcpy_fn
    if _batch_memcpy_fn is None:
@@ -144,7 +137,10 @@ def build_params(
        dst_bases.append(d.data_ptr())
        bpb.append(s_bpb)

-    attrs = _CUmemcpyAttributes(srcAccessOrder=src_access_order)
+    # ``srcAccessOrder=3`` == CU_MEMCPY_SRC_ACCESS_ORDER_ANY /
+    # hipMemcpySrcAccessOrderAny. See
+    # https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1g6f1ff58e3065df3eb4b573dba77ad31f  # noqa: E501
+    attrs = _CUmemcpyAttributes(srcAccessOrder=3)

    return BatchMemcpyParams(
        src_bases=np.array(src_bases, dtype=np.uint64),
@@ -57,10 +57,6 @@ class SimpleCPUOffloadWorker:
        # Metadata for the current step
        self._connector_metadata: SimpleCPUOffloadMetadata | None = None

-        # Compute-done event recorded before each store; reused across steps
-        # (get_finished runs once per step, copy queue is FIFO).
-        self._store_compute_done: torch.Event | None = None
-
        # Pending event index sets, populated in bind_connector_metadata
        self._pending_load_event_indices: set[int] = set()
        self._pending_store_event_indices: set[int] = set()
@@ -210,11 +206,9 @@ class SimpleCPUOffloadWorker:
    ) -> tuple[set[str] | None, set[str] | None]:
        """Submit transfers and report completed events to the scheduler.

-        Stores (GPU->CPU) read the live KV cache, which the compute stream may
-        still be writing under v1 overlapped execution, so they are ordered
-        after a compute-done event recorded on the current stream. Loads
-        (CPU->GPU) read stable pinned host memory and launch immediately. See
-        #45704 for the bug and #39306 for the srcAccessOrder rationale.
+        Called after model execution. The manager only schedules stores for
+        blocks whose KV data is confirmed computed, so we launch both loads
+        and stores immediately — no deferral or cross-stream sync needed.

        Returns:
            tuple of (finished_sending, finished_recving).
@@ -224,6 +218,7 @@ class SimpleCPUOffloadWorker:
        # (1) Submit transfers
        metadata = self._connector_metadata
        if metadata is not None:
+            # Launch loads (CPU->GPU).
            if metadata.load_cpu_blocks:
                self._backend.launch_copy(
                    metadata.load_cpu_blocks,
@@ -232,17 +227,14 @@ class SimpleCPUOffloadWorker:
                    event_idx=metadata.load_event,
                    events_list=self._load_events,
                )
+            # Launch stores (GPU->CPU).
            if metadata.store_gpu_blocks:
-                if self._store_compute_done is None:
-                    self._store_compute_done = torch.Event()
-                self._store_compute_done.record(torch.cuda.current_stream())
                self._backend.launch_copy(
                    metadata.store_gpu_blocks,
                    metadata.store_cpu_blocks,
                    is_store=True,
                    event_idx=metadata.store_event,
                    events_list=self._store_events,
-                    wait_event=self._store_compute_done,
                )

        # (2) Track completed transfer events
@@ -68,19 +68,15 @@ class EncoderRunner:
        query_start_loc: np.ndarray,
        prefill_lens: np.ndarray,
        computed_prefill_lens: np.ndarray,
-        draft_lookahead: int = 0,
    ) -> tuple[list[torch.Tensor], torch.Tensor]:
-        if draft_lookahead:
-            computed_prefill_lens = computed_prefill_lens + draft_lookahead
-
-        is_prefilling_np = computed_prefill_lens < prefill_lens
-        if not is_prefilling_np.any():
+        is_prefilling = (computed_prefill_lens < prefill_lens).tolist()
+        all_decode = not any(is_prefilling)
+        if all_decode:
            # All decode requests, so no need to gather any embeddings.
            return [], torch.zeros(
                total_num_scheduled_tokens, dtype=torch.bool, device=self.device
            )

-        is_prefilling = is_prefilling_np.tolist()
        query_start = computed_prefill_lens.tolist()
        query_end = (computed_prefill_lens + num_scheduled_tokens).tolist()

@@ -93,12 +89,11 @@ class EncoderRunner:
                # OPTIMIZATION: Skip decode requests.
                continue

-            cur_query_start = query_start[i]
-            cur_query_end = query_end[i]
-
            mm_features = self.encoder_cache.mm_features[req_id]
            lo, hi = get_mm_features_in_window(
-                mm_features, start=cur_query_start, end=cur_query_end
+                mm_features,
+                start=query_start[i],
+                end=query_end[i],
            )
            for idx in range(lo, hi):
                mm_feature = mm_features[idx]
@@ -106,8 +101,8 @@ class EncoderRunner:
                start_pos = pos_info.offset
                num_encoder_tokens = pos_info.length

-                start_idx = max(cur_query_start - start_pos, 0)
-                end_idx = min(cur_query_end - start_pos, num_encoder_tokens)
+                start_idx = max(query_start[i] - start_pos, 0)
+                end_idx = min(query_end[i] - start_pos, num_encoder_tokens)
                assert start_idx < end_idx
                curr_embeds_start, curr_embeds_end = (
                    pos_info.get_embeds_indices_in_range(start_idx, end_idx)
@@ -119,13 +114,7 @@ class EncoderRunner:

                mm_hash = mm_feature.identifier
                encoder_output = self.encoder_cache.encoder_outputs.get(mm_hash, None)
-                if encoder_output is None:
-                    # A feature starting at/after the processed boundary is only
-                    # reached via the drafter's +1 look-ahead and might not be
-                    # encoded yet; fall back to the token embedding for drafting.
-                    if start_pos + draft_lookahead >= cur_query_end:
-                        continue
-                    raise RuntimeError(f"Encoder cache miss for {mm_hash}.")
+                assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."

                if (is_embed := pos_info.is_embed) is not None:
                    is_embed = is_embed[start_idx:end_idx]
@@ -133,7 +122,7 @@ class EncoderRunner:
                else:
                    mm_embeds_item = encoder_output[start_idx:end_idx]

-                req_start_pos = query_start_loc[i] + start_pos - cur_query_start
+                req_start_pos = query_start_loc[i] + start_pos - query_start[i]
                is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= (
                    True if is_embed is None else is_embed
                )
@@ -314,7 +314,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
        self.model_state = init_model_state(
            self.vllm_config, self.model, self.encoder_cache, self.device
        )
-        self.model_state.req_states = self.req_states

        self.decode_query_len = (
            self.num_speculative_steps
@@ -1410,9 +1409,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                input_batch.num_scheduled_tokens,
                input_batch.query_start_loc_np,
                input_batch.prefill_len_np,
-                input_batch.num_computed_prefill_tokens_np,
-                # The EAGLE/MTP drafter reads one position ahead of the target.
-                draft_lookahead=1,
+                # +1 to consider the skew in eagle
+                input_batch.num_computed_prefill_tokens_np + 1,
            )

        # Postprocess results and update request states.
@@ -4,7 +4,6 @@ import torch
 import torch.nn as nn

 from vllm.config import VllmConfig
-from vllm.model_executor.layers.attention import CrossAttention
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache


@@ -19,13 +18,13 @@ def init_model_state(
        cls = model.get_model_state_cls()
        return cls(vllm_config, model, encoder_cache, device)

-    # Cross-attention encoder-decoder models (Whisper, CohereASR, NemotronParse, ...)
-    if any(isinstance(m, CrossAttention) for m in model.modules()):
-        from vllm.v1.worker.gpu.model_states.encoder_decoder import (
-            EncoderDecoderModelState,
-        )
+    if (
+        "WhisperForConditionalGeneration" in vllm_config.model_config.architectures
+        or "CohereAsrForConditionalGeneration" in vllm_config.model_config.architectures
+    ):
+        from vllm.v1.worker.gpu.model_states.whisper import WhisperModelState

-        return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
+        return WhisperModelState(vllm_config, model, encoder_cache, device)

    if vllm_config.model_config.is_hybrid:
        from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
@@ -7,8 +7,7 @@ import torch.nn as nn

 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.model_executor.models.interfaces import supports_multimodal_pruning
-from vllm.multimodal.utils import get_mm_features_in_window
+from vllm.tasks import GenerationTask
 from vllm.v1.core.sched.output import NewRequestData
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
@@ -34,7 +33,6 @@ class DefaultModelState(ModelState):
        self.scheduler_config = vllm_config.scheduler_config
        self.model = model
        self.device = device
-        self.req_states: RequestState | None = None

        self.supports_mm_inputs = encoder_cache is not None
        self.max_model_len = self.model_config.max_model_len
@@ -42,11 +40,6 @@ class DefaultModelState(ModelState):
        self.max_num_tokens = self.scheduler_config.max_num_batched_tokens
        self.inputs_embeds_size = self.model_config.get_inputs_embeds_size()
        self.dtype = self.model_config.dtype
-        self.is_multimodal_pruning_enabled = (
-            supports_multimodal_pruning(model)
-            and self.model_config.multimodal_config is not None
-            and self.model_config.multimodal_config.is_multimodal_pruning_enabled()
-        )

        if self.supports_mm_inputs:
            assert encoder_cache is not None
@@ -69,6 +62,28 @@ class DefaultModelState(ModelState):
            device=self.device,
        )

+    def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
+        from vllm.model_executor.models.interfaces import (
+            supports_realtime,
+            supports_transcription,
+        )
+        from vllm.model_executor.models.interfaces_base import is_text_generation_model
+
+        supported_tasks = list[GenerationTask]()
+
+        if is_text_generation_model(self.model):
+            supported_tasks.append("generate")
+
+        if supports_transcription(self.model):
+            if self.model.supports_transcription_only:
+                return ("transcription",)
+            supported_tasks.append("transcription")
+
+        if supports_realtime(self.model):
+            supported_tasks.append("realtime")
+
+        return tuple(supported_tasks)
+
    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
        if self.rope_state is not None:
            assert new_req_data.prefill_token_ids is not None
@@ -83,82 +98,6 @@ class DefaultModelState(ModelState):
        if self.rope_state is not None:
            self.rope_state.apply_staged_writes()

-    def _recompute_mrope_positions(
-        self,
-        mm_embeds: list[torch.Tensor],
-        input_batch: InputBatch,
-    ) -> list[torch.Tensor]:
-        assert self.rope_state is not None
-        assert self.req_states is not None
-        req_states = self.req_states
-
-        mm_embeds_out: list[torch.Tensor] = []
-        mm_embed_idx = 0
-        for batch_idx, req_id in enumerate(input_batch.req_ids):
-            req_idx = req_states.req_id_to_index[req_id]
-            num_computed_tokens = int(req_states.num_computed_tokens_np[req_idx])
-
-            mm_features = self.encoder_cache.mm_features[req_id]
-            query_start = num_computed_tokens
-            query_end = query_start + int(input_batch.num_scheduled_tokens[batch_idx])
-
-            num_req_mm_embeds = 0
-            lo, hi = get_mm_features_in_window(
-                mm_features,
-                start=query_start,
-                end=query_end,
-            )
-            # iterate and get the mm_embeds num in current window
-            for mm_feature in mm_features[lo:hi]:
-                start_pos = mm_feature.mm_position.offset
-                num_encoder_tokens = mm_feature.mm_position.length
-                start_idx = max(query_start - start_pos, 0)
-                end_idx = min(query_end - start_pos, num_encoder_tokens)
-                curr_embeds_start, curr_embeds_end = (
-                    mm_feature.mm_position.get_embeds_indices_in_range(
-                        start_idx, end_idx
-                    )
-                )
-                if curr_embeds_start != curr_embeds_end:
-                    num_req_mm_embeds += 1
-
-            if num_req_mm_embeds == 0:
-                continue
-
-            req_mm_embeds = mm_embeds[mm_embed_idx : mm_embed_idx + num_req_mm_embeds]
-            mm_embed_idx += num_req_mm_embeds
-
-            # get prompttoken ids
-            prompt_len = int(req_states.prompt_len.np[req_idx])
-            prompt_token_ids = req_states.all_token_ids._uva_buf.np[
-                req_idx, :prompt_len
-            ].tolist()
-            # get mrope positions
-            start = req_idx * self.rope_state.num_dims
-            end = start + self.rope_state.num_dims
-            mrope_positions = torch.tensor(
-                self.rope_state.prefill_positions._uva_buf.np[start:end, :prompt_len],
-                dtype=torch.long,
-            )
-            req_mm_embeds, new_positions, new_delta = (
-                self.model.recompute_mrope_positions(
-                    input_ids=prompt_token_ids,
-                    multimodal_embeddings=tuple(req_mm_embeds),
-                    mrope_positions=mrope_positions,
-                    num_computed_tokens=num_computed_tokens,
-                )
-            )
-            new_positions_cpu = new_positions.to(device="cpu", dtype=torch.int32)
-            self.rope_state.prefill_positions._uva_buf.cpu[
-                start:end, : new_positions_cpu.shape[1]
-            ].copy_(new_positions_cpu)
-            self.rope_state.prefill_delta.np[req_idx] = new_delta
-            self.rope_state.prefill_delta.copy_to_uva()
-            mm_embeds_out.extend(req_mm_embeds)
-
-        assert mm_embed_idx == len(mm_embeds)
-        return mm_embeds_out
-
    def get_mm_embeddings(
        self,
        scheduled_encoder_inputs: dict[str, list[int]],
@@ -181,13 +120,6 @@ class DefaultModelState(ModelState):
            input_batch.prefill_len_np,
            input_batch.num_computed_prefill_tokens_np,
        )
-        if (
-            mm_embeds
-            and self.is_multimodal_pruning_enabled
-            and self.rope_state is not None
-            and self.rope_state.has_delta
-        ):
-            mm_embeds = self._recompute_mrope_positions(mm_embeds, input_batch)
        # Use unpadded input_ids to match is_mm_embed size (num_tokens).
        # input_batch.input_ids may be padded for CUDA graphs.
        input_ids_unpadded = input_batch.input_ids[: input_batch.num_tokens]
@@ -46,25 +46,9 @@ class ModelState(ABC):
    ) -> None:
        raise NotImplementedError

-    model: nn.Module
-
+    @abstractmethod
    def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
-        from vllm.model_executor.models.interfaces import (
-            supports_realtime,
-            supports_transcription,
-        )
-        from vllm.model_executor.models.interfaces_base import is_text_generation_model
-
-        supported_tasks = list[GenerationTask]()
-        if is_text_generation_model(self.model):
-            supported_tasks.append("generate")
-        if supports_transcription(self.model):
-            if self.model.supports_transcription_only:
-                return ("transcription",)
-            supported_tasks.append("transcription")
-        if supports_realtime(self.model):
-            supported_tasks.append("realtime")
-        return tuple(supported_tasks)
+        raise NotImplementedError

    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
        return None
@@ -23,7 +23,7 @@ from vllm.v1.worker.utils import AttentionGroup


@dataclass
-class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
+class WhisperAttnMetadata(ModelSpecificAttnMetadata):
    encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]

    def get_extra_common_attn_kwargs(
@@ -41,11 +41,7 @@ class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
        }


-class EncoderDecoderModelState(ModelState):
-    """ModelState for cross-attention encoder-decoder models
-    (Whisper, CohereASR, NemotronParse, FireRedLID, ...)
-    """
-
+class WhisperModelState(ModelState):
    def __init__(
        self,
        vllm_config: VllmConfig,
@@ -84,6 +80,9 @@ class EncoderDecoderModelState(ModelState):

        self.encoder_outputs: list[torch.Tensor] = []

+    def get_supported_generation_tasks(self):
+        return ("transcription",)
+
    def get_mm_embeddings(
        self, scheduled_encoder_inputs: dict[str, list[int]], input_batch: InputBatch
    ) -> None:
@@ -95,11 +94,11 @@ class EncoderDecoderModelState(ModelState):
                encoder_inputs[req_id] = req_encoder_inputs
        _, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs)
        if mm_kwargs:
-            # Encoder-decoder models consume encoder outputs through the
-            # `encoder_outputs` forward kwarg, not `inputs_embeds`. Single modality
-            # so execute_mm_encoder preserves request order; use its return value
-            # directly. No need to store in encoder_cache: cross-attention K/V are
-            # written to the KV cache on the first step; decode steps use the cache.
+            # Whisper consumes encoder outputs through `encoder_outputs`, not
+            # `inputs_embeds`. Single modality (audio) so execute_mm_encoder
+            # preserves request order; use its return value directly.
+            # No need to store in encoder_cache: cross-attention K/V are written
+            # to the KV cache on the first step; decode steps use the cache.
            self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
        else:
            # Decode steps: encoder K/V are in cross-attention KV cache.
@@ -132,7 +131,7 @@ class EncoderDecoderModelState(ModelState):
        else:
            num_reqs = input_batch.num_reqs
            num_tokens = input_batch.num_tokens
-        enc_dec_attn_metadata = EncoderDecoderAttnMetadata(
+        whisper_attn_metadata = WhisperAttnMetadata(
            self._get_encoder_seq_lens(
                input_batch.req_ids, attn_groups, for_capture, num_reqs
            )
@@ -159,7 +158,7 @@ class EncoderDecoderModelState(ModelState):
            kv_cache_config=kv_cache_config,
            seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
            dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
-            model_specific_attn_metadata=enc_dec_attn_metadata,
+            model_specific_attn_metadata=whisper_attn_metadata,
            for_cudagraph_capture=for_capture,
        )
        return attn_metadata
@@ -2,14 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 from collections.abc import Callable
-from contextlib import AbstractContextManager, nullcontext
 from typing import Any

 import numpy as np
 import torch

 from vllm import PoolingParams, SamplingParams
-from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import (
    CachedRequestData,
@@ -20,134 +18,6 @@ from vllm.v1.core.sched.output import (
 from vllm.v1.request import Request
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner

-logger = init_logger(__name__)
-
-
-def run_mixed_prefill_decode_warmup(
-    model_runner: GPUModelRunner,
-    worker_execute_model: Callable[[SchedulerOutput], Any],
-    worker_sample_tokens: Callable[[GrammarOutput | None], Any],
-    num_tokens: int,
-    *,
-    mixed_step_context: AbstractContextManager[object] | None = None,
-    req_id_prefix: str = "_v2_mixed_warmup",
-) -> bool:
-    """Run a V2 mixed prefill+decode step through normal scheduler inputs."""
-    if model_runner.is_pooling_model or num_tokens < 3:
-        return False
-
-    decode_req_id = f"{req_id_prefix}_decode_"
-    prefill_req_id = f"{req_id_prefix}_prefill_"
-    decode_prompt_len = 2
-    decode_scheduled_tokens = 1
-    prefill_len = num_tokens - decode_scheduled_tokens
-    decode_token_ids = list(range(decode_prompt_len))
-    prefill_token_ids = list(range(prefill_len))
-
-    kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
-    num_kv_cache_groups = len(kv_cache_groups)
-    group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups]
-    decode_prefill_block_counts = [
-        cdiv(decode_prompt_len, block_size) for block_size in group_block_sizes
-    ]
-    decode_block_counts = [
-        cdiv(decode_prompt_len + decode_scheduled_tokens, block_size)
-        for block_size in group_block_sizes
-    ]
-    decode_block_deltas = [
-        decode - prefill
-        for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts)
-    ]
-    prefill_block_counts = [
-        cdiv(prefill_len, block_size) for block_size in group_block_sizes
-    ]
-    required_blocks = sum(decode_block_counts) + sum(prefill_block_counts)
-    if model_runner.kv_cache_config.num_blocks <= required_blocks:
-        logger.warning(
-            "Skipping V2 mixed prefill+decode warmup because only %d KV blocks "
-            "are available for %d required warmup blocks.",
-            model_runner.kv_cache_config.num_blocks,
-            required_blocks,
-        )
-        return False
-
-    next_block_id = 1
-
-    def _alloc_blocks(num_blocks: int) -> list[int]:
-        nonlocal next_block_id
-        block_ids = list(range(next_block_id, next_block_id + num_blocks))
-        next_block_id += num_blocks
-        return block_ids
-
-    sampling_params = SamplingParams(max_tokens=2, temperature=0.0)
-
-    decode_prefill_output = SchedulerOutput.make_empty()
-    decode_prefill_output.scheduled_new_reqs = [
-        NewRequestData(
-            req_id=decode_req_id,
-            prompt_token_ids=decode_token_ids,
-            mm_features=[],
-            sampling_params=sampling_params,
-            pooling_params=None,
-            block_ids=tuple(_alloc_blocks(n) for n in decode_prefill_block_counts),
-            num_computed_tokens=0,
-            lora_request=None,
-            prefill_token_ids=decode_token_ids,
-        ),
-    ]
-    decode_prefill_output.num_scheduled_tokens = {
-        decode_req_id: decode_prompt_len,
-    }
-    decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len
-    decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
-
-    decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas)
-    cached_decode_req = CachedRequestData.make_empty()
-    cached_decode_req.req_ids = [decode_req_id]
-    cached_decode_req.num_computed_tokens = [decode_prompt_len]
-    cached_decode_req.num_output_tokens = [1]
-    cached_decode_req.new_block_ids = [
-        decode_new_blocks if any(decode_block_deltas) else None
-    ]
-
-    mixed_output = SchedulerOutput.make_empty()
-    mixed_output.scheduled_cached_reqs = cached_decode_req
-    mixed_output.scheduled_new_reqs = [
-        NewRequestData(
-            req_id=prefill_req_id,
-            prompt_token_ids=prefill_token_ids,
-            mm_features=[],
-            sampling_params=sampling_params,
-            pooling_params=None,
-            block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
-            num_computed_tokens=0,
-            lora_request=None,
-            prefill_token_ids=prefill_token_ids,
-        ),
-    ]
-    mixed_output.num_scheduled_tokens = {
-        decode_req_id: decode_scheduled_tokens,
-        prefill_req_id: prefill_len,
-    }
-    mixed_output.total_num_scheduled_tokens = num_tokens
-    mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
-
-    cleanup_output = SchedulerOutput.make_empty()
-    cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id}
-
-    context = mixed_step_context or nullcontext()
-    model_runner.kv_connector.set_disabled(True)
-    try:
-        worker_execute_model(decode_prefill_output)
-        worker_sample_tokens(None)
-        with context:
-            worker_execute_model(mixed_output)
-            worker_sample_tokens(None)
-        worker_execute_model(cleanup_output)
-    finally:
-        model_runner.kv_connector.set_disabled(False)
-    return True
-

@torch.inference_mode()
 def warmup_kernels(
@@ -3152,16 +3152,7 @@ class GPUModelRunner(

                mm_hash = mm_feature.identifier
                encoder_output = self.encoder_cache.get(mm_hash, None)
-                if encoder_output is None:
-                    # A feature starting at/after the processed boundary is only
-                    # reached via the drafter's +1 look-ahead and might not be
-                    # encoded yet; fall back to the token embedding for drafting.
-                    if (
-                        start_pos
-                        >= req_state.num_computed_tokens + num_scheduled_tokens
-                    ):
-                        continue
-                    raise RuntimeError(f"Encoder cache miss for {mm_hash}.")
+                assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."

                if (is_embed := pos_info.is_embed) is not None:
                    is_embed = is_embed[start_idx:end_idx]