Merge branch 'main' into wentao-fix-v2-mrope

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/hardware_tests/cpu.yaml

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

.buildkite/image_build/image_build_arm64.sh

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

.buildkite/intel_jobs/expert_parallelism_intel.yaml

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

.buildkite/intel_jobs/models_multimodal_intel.yaml

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

.buildkite/test_areas/lm_eval.yaml

@@ -136,6 +136,49 @@ steps:
   commands:
     - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
 
+- label: LM Eval Humming (A100 - TEMPORARY)
+  key: lm-eval-humming-a100
+  timeout_in_minutes: 30
+  device: a100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+
+- label: LM Eval Humming (H100 - TEMPORARY)
+  key: lm-eval-humming-h100
+  timeout_in_minutes: 30
+  device: h100
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
+
+- label: LM Eval Humming (B200 - TEMPORARY)
+  key: lm-eval-humming-b200
+  timeout_in_minutes: 30
+  device: b200-k8s
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - vllm/model_executor/layers/quantization/humming.py
+  - vllm/model_executor/layers/quantization/utils/humming_utils.py
+  - vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+  - vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
 
 - label: LM Eval TurboQuant KV Cache
   key: lm-eval-turboquant-kv-cache
@@ -177,6 +220,20 @@ steps:
     - uv pip install --system 'gpt-oss[eval]==0.0.5'
     - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
 
+- label: GPQA Eval (GPT-OSS) (DGX Spark)
+  key: gpqa-eval-gpt-oss-spark
+  timeout_in_minutes: 120
+  device: dgx-spark
+  optional: true
+  num_devices: 1
+  source_file_dependencies:
+  - csrc/
+  - vllm/model_executor/layers/quantization
+  - tests/evals/gpt_oss/
+  commands:
+    - uv pip install --system 'gpt-oss[eval]==0.0.5'
+    - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-spark.txt
+
 - label: MRCR Eval Small Models
   device: h200_35gb
   timeout_in_minutes: 30

cmake/external_projects/deepgemm.cmake

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

cmake/external_projects/qutlass.cmake

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

docs/design/attention_backends.md

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

requirements/cuda.txt

@@ -26,4 +26,4 @@ quack-kernels>=0.3.3
 tokenspeed-mla==0.1.2
 
 # Humming kernels for quantization gemm
-humming-kernels[cu13]==0.1.4
+humming-kernels[cu13]==0.1.6

tests/distributed/test_elastic_ep.py

@@ -59,9 +59,8 @@ def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
     return accuracy
 
 
-@multi_gpu_test(num_gpus=4)
-def test_elastic_ep_scaling():
-    vllm_serve_args = [
+def _base_serve_args(use_async_eplb: bool = False) -> list[str]:
+    args = [
         "--trust-remote-code",
         "--tensor-parallel-size",
         "1",
@@ -79,7 +78,11 @@ def test_elastic_ep_scaling():
         "--eplb-config.num_redundant_experts",
         "0",
         "--eplb-config.use_async",
-        "false",
+        "true" if use_async_eplb else "false",
+        "--eplb-config.step_interval",
+        "10",
+        "--eplb-config.window_size",
+        "5",
         "--data-parallel-backend",
         "ray",
         "--data-parallel-size",
@@ -90,7 +93,23 @@ def test_elastic_ep_scaling():
 
     leader_address = os.environ.get("LEADER_ADDRESS")
     if leader_address:
-        vllm_serve_args.extend(["--data-parallel-address", leader_address])
+        args.extend(["--data-parallel-address", leader_address])
+
+    return args
+
+
+@pytest.mark.parametrize(
+    "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
+)
+@multi_gpu_test(num_gpus=4)
+def test_elastic_ep_scaling(use_async_eplb: bool):
+    if use_async_eplb:
+        from vllm.distributed.eplb.eplb_communicator import has_nixl
+
+        if not has_nixl():
+            pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
+
+    vllm_serve_args = _base_serve_args(use_async_eplb)
 
     with RemoteOpenAIServer(
         MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
@@ -128,44 +147,24 @@ def test_elastic_ep_scaling():
         print(f"  Tolerance:  {ACCURACY_TOL:.3f}")
 
 
+@pytest.mark.parametrize(
+    "use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
+)
 @multi_gpu_test(num_gpus=4)
-def test_elastic_ep_scaling_uneven():
+def test_elastic_ep_scaling_uneven(use_async_eplb: bool):
     """Test scale up with uneven worker distribution.
 
     This tests the case where num_new_workers % old_dp_size != 0,
     specifically 2 -> 3 where remainder = 1 % 2 = 1.
     This exercises the remainder handling in sender-receiver pairing.
     """
-    vllm_serve_args = [
-        "--trust-remote-code",
-        "--tensor-parallel-size",
-        "1",
-        "--gpu-memory-utilization",
-        "0.8",
-        "--max-model-len",
-        "4096",
-        "--max-num-seqs",
-        str(MAX_NUM_SEQS),
-        "--enable-expert-parallel",
-        "--all2all-backend",
-        "allgather_reducescatter",
-        "--enable-elastic-ep",
-        "--enable-eplb",
-        "--eplb-config.num_redundant_experts",
-        "0",
-        "--eplb-config.use_async",
-        "false",
-        "--data-parallel-backend",
-        "ray",
-        "--data-parallel-size",
-        "2",
-        "--api-server-count",
-        "1",
-    ]
+    if use_async_eplb:
+        from vllm.distributed.eplb.eplb_communicator import has_nixl
 
-    leader_address = os.environ.get("LEADER_ADDRESS")
-    if leader_address:
-        vllm_serve_args.extend(["--data-parallel-address", leader_address])
+        if not has_nixl():
+            pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
+
+    vllm_serve_args = _base_serve_args(use_async_eplb)
 
     with RemoteOpenAIServer(
         MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200

tests/distributed/test_eplb_execute.py

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

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

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+model_name: "openai/gpt-oss-20b"
+metric_threshold: 0.568
+reasoning_effort: "low"

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

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

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

@@ -0,0 +1,12 @@
+model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
+accuracy_threshold: 0.86
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --quantization humming
+  --kernel-config.enable_flashinfer_autotune=False
+env:
+  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'

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

@@ -0,0 +1,10 @@
+model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
+accuracy_threshold: 0.86
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --quantization humming
+  --kernel-config.enable_flashinfer_autotune=False

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

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

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

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

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

@@ -0,0 +1,11 @@
+model_name: "openai/gpt-oss-20b"
+accuracy_threshold: 0.30
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --moe-backend humming
+env:
+  VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3"}'

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

@@ -0,0 +1,9 @@
+model_name: "openai/gpt-oss-20b"
+accuracy_threshold: 0.30
+num_questions: 1319
+num_fewshot: 5
+server_args: >-
+  --enforce-eager
+  --max-model-len 8192
+  --tensor-parallel-size 1
+  --moe-backend humming

tests/kernels/moe/test_moe_weight_loading_padded.py

@@ -325,3 +325,24 @@ class TestWeightLoadingWithPaddedHiddenSize:
                 shard_id="w2",
                 expert_id=0,
             )
+
+
+class TestPerTensorScaleCoercion:
+    """Regression test for shape-(1,) per-tensor scales (issue #43297).
+
+    llm-compressor NVFP4 emits per-tensor weight and input scales as
+    shape-(1,) tensors. `_to_scalar` collapses them to a 0-D scalar so the
+    scalar-slot assignments in the weight loader neither broadcast nor raise.
+    """
+
+    def test_collapses_to_scalar(self):
+        # shape-(1,) and 0-D both reduce to a 0-D scalar.
+        for loaded_weight in (torch.tensor([0.5]), torch.tensor(0.5)):
+            scalar = RoutedExperts._to_scalar(loaded_weight)
+            assert scalar.shape == ()
+            assert scalar.item() == pytest.approx(0.5)
+
+    def test_rejects_non_scalar(self):
+        # numel > 1 must fail loudly instead of silently picking an element.
+        with pytest.raises(RuntimeError):
+            RoutedExperts._to_scalar(torch.tensor([0.1, 0.2]))

tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py

@@ -75,7 +75,7 @@ def get_ref_results(
 @pytest.mark.parametrize("shape", SHAPES)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
-@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm", "b12x"])
+@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass", "cudnn", "trtllm", "b12x"])
 @pytest.mark.parametrize("autotune", [False, True])
 @torch.inference_mode()
 def test_flashinfer_nvfp4_gemm(
@@ -88,6 +88,8 @@ def test_flashinfer_nvfp4_gemm(
 ) -> None:
     if "trtllm" in backend and dtype == torch.float16:
         pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
+    if backend == "cute-dsl" and not current_platform.is_device_capability_family(100):
+        pytest.skip("FlashInfer cutedsl backend is only supported on SM10x")
     if backend == "b12x" and not current_platform.has_device_capability(120):
         pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
     if backend == "b12x" and not has_flashinfer_b12x_gemm():

tests/model_executor/test_flashinfer_autotune_cache.py

@@ -1,60 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import sys
-from hashlib import sha256
-from pathlib import Path
-from types import SimpleNamespace
-
-from vllm.model_executor.warmup import kernel_warmup
-
-
-def test_resolve_flashinfer_autotune_file_default_layout(
-    monkeypatch, tmp_path: Path
-) -> None:
-    fake_jit = SimpleNamespace(
-        env=SimpleNamespace(
-            FLASHINFER_WORKSPACE_DIR=Path("/flashinfer-cache/0.6.11.post2/103a")
-        )
-    )
-    fake_flashinfer = SimpleNamespace(jit=fake_jit)
-    monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
-    monkeypatch.setitem(sys.modules, "flashinfer.jit", fake_jit)
-    monkeypatch.setattr(
-        kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
-    )
-    monkeypatch.setattr(kernel_warmup.envs, "VLLM_CACHE_ROOT", str(tmp_path))
-    monkeypatch.setattr(kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None)
-
-    runner = SimpleNamespace(vllm_config=SimpleNamespace())
-    cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
-
-    path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
-
-    assert path == (
-        tmp_path
-        / "flashinfer_autotune_cache"
-        / "0.6.11.post2"
-        / "103a"
-        / cache_hash
-        / "autotune_configs.json"
-    )
-    assert path.parent.is_dir()
-
-
-def test_resolve_flashinfer_autotune_file_uses_override_dir(
-    monkeypatch, tmp_path: Path
-) -> None:
-    monkeypatch.setattr(
-        kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)
-    )
-    monkeypatch.setattr(
-        kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
-    )
-
-    runner = SimpleNamespace(vllm_config=SimpleNamespace())
-    cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
-
-    path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
-
-    assert path == tmp_path / cache_hash / "autotune_configs.json"

tests/models/quantization/test_nvfp4.py

@@ -90,6 +90,7 @@ def test_models(example_prompts, model_name) -> None:
 EAGER = [True, False]
 
 SM_100_NVFP4_BACKENDS = [
+    "flashinfer_cutedsl",
     "flashinfer_cudnn",
     "flashinfer_trtllm",
     "flashinfer_cutlass",
@@ -102,12 +103,18 @@ SM_100_NVFP4_BACKENDS = [
     "backend",
     [
         "emulation",
+        "flashinfer_cutedsl",
         "flashinfer_cudnn",
         "flashinfer_trtllm",  # the small seq_len ensures trtllm_8x4_layout backend is used
         "flashinfer_cutlass",
     ],
 )
 def test_nvfp4(vllm_runner, model, eager, backend):
+    if backend == "flashinfer_cutedsl" and not (
+        current_platform.is_device_capability_family(100)
+    ):
+        pytest.skip("The flashinfer_cutedsl backend is only supported on SM10x")
+
     if (
         not current_platform.has_device_capability(100)
         and backend in SM_100_NVFP4_BACKENDS

tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py

@@ -0,0 +1,54 @@
+# 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 == []

tests/v1/attention/test_sparse_mla_backends.py

@@ -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 (
-    FlashInferMLASparseBackend,
+    FlashInferMLASparseTRTLLMBackend,
 )
 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, FlashInferMLASparseBackend],
-    ids=["FlashMLA", "FlashInfer"],
+    [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
+    ids=["FlashMLA", "FlashInferTRTLLM"],
 )
 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])
@@ -217,9 +217,12 @@ def test_sparse_backend_decode_correctness(
         ok, reason = flashmla.is_flashmla_sparse_supported()
         if not ok:
             pytest.skip(reason)
-    elif backend_cls == FlashInferMLASparseBackend:
-        if not current_platform.has_device_capability(100):
-            pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher")
+    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")
 
     batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
     use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"

tests/v1/core/test_scheduler.py

@@ -4708,6 +4708,38 @@ 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."""

tests/v1/kv_connector/unit/test_kv_cache_layout.py

@@ -3,6 +3,8 @@
 
 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
@@ -24,8 +26,13 @@ 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",

tests/v1/simple_kv_offload/test_worker.py

@@ -0,0 +1,183 @@
+# 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

tests/v1/spec_decode/test_acceptance_length.py

@@ -132,9 +132,9 @@ def get_available_attention_backends() -> list[str]:
     )
 
     return [
-        backend.name
-        for backend, _ in valid_backends
-        if backend not in EXCLUDED_BACKENDS
+        candidate.backend.name
+        for candidate in valid_backends
+        if candidate.backend not in EXCLUDED_BACKENDS
     ]
 
 

tests/v1/worker/test_encoder_runner.py

@@ -0,0 +1,157 @@
+# 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)

tests/v1/worker/test_gpu_model_runner_mm_gather.py

@@ -0,0 +1,99 @@
+# 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)

tools/pre_commit/generate_attention_backend_docs.py

@@ -690,7 +690,9 @@ def parse_compute_capability(node: ast.ClassDef) -> str:
         major_list.sort()
         if len(major_list) == 1:
             return f"{major_list[0]}.x"
-        return f"{major_list[0]}.x-{major_list[-1]}.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)
 
     if min_cap:
         if max_cap:
@@ -1668,7 +1670,8 @@ 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 `FLASHMLA_SPARSE_DSV4`.",
+                "default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and",
+                "`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.",
                 "",
             ]
         )

vllm/compilation/passes/fusion/collective_fusion.py

@@ -956,6 +956,15 @@ 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

vllm/config/kernel.py

@@ -141,6 +141,7 @@ LinearBackend = Literal[
     "auto",
     "cutlass",
     "flashinfer_cutlass",
+    "flashinfer_cutedsl",
     "flashinfer_trtllm",
     "flashinfer_cudnn",
     "flashinfer_b12x",
@@ -198,6 +199,7 @@ 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+)

vllm/config/parallel.py

@@ -803,13 +803,6 @@ 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 "
@@ -821,6 +814,15 @@ 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.
@@ -929,23 +931,21 @@ class ParallelConfig:
             )
 
         if self.enable_eplb and self.eplb_config.communicator is None:
-            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
+            # 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:
                 self.eplb_config.communicator = "pynccl"
             else:
-                # 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"
+                self.eplb_config.communicator = "torch_gloo"
 
     @property
     def use_ray(self) -> bool:

vllm/distributed/elastic_ep/elastic_execute.py

@@ -207,6 +207,9 @@ 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()
 
@@ -540,6 +543,11 @@ 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")
 
@@ -549,6 +557,9 @@ 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

vllm/distributed/eplb/async_worker.py

@@ -8,7 +8,6 @@ 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
@@ -26,8 +25,7 @@ def start_async_worker(
     state: "EplbState",
     is_profile: bool = False,
 ) -> threading.Thread:
-    eplb_group = get_eplb_group().device_group
-    rank = eplb_group.rank()
+    rank = get_eplb_group().device_group.rank()
     device_index = state.cuda_device_index
     assert state.is_async
 
@@ -38,7 +36,6 @@ def start_async_worker(
         try:
             transfer_run_periodically(
                 state=state,
-                eplb_group=eplb_group,
                 cuda_stream=cuda_stream,
                 is_profile=is_profile,
             )
@@ -78,13 +75,15 @@ 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)
-        logger.info("async worker woke up for EPLB transfer")
+
+        eplb_group = get_eplb_group().device_group
+        eplb_cpu_group = get_eplb_group().cpu_group
+        ep_rank = eplb_group.rank()
 
         assert state.is_async
         for model_state in state.model_states.values():
@@ -101,16 +100,32 @@ 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
-            while model_state.rebalanced and layer_idx < num_layers:
+            # 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
+
                 transfer_metadata = transfer_layer(
                     old_layer_indices=physical_to_logical_map_cpu[layer_idx],
                     new_layer_indices=new_physical_to_logical_map[layer_idx],
@@ -143,6 +158,5 @@ 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

vllm/distributed/eplb/eplb_communicator.py

@@ -246,7 +246,20 @@ 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."
         )
@@ -302,10 +315,29 @@ 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._log_initialized()
+        self._remote_state_initialized = True
+
+    def _ensure_remote_state(self) -> None:
+        if not self._remote_state_initialized:
+            self._init_remote_state()
 
     @property
     def needs_profile_buffer_reservation(self) -> bool:
@@ -339,8 +371,7 @@ class NixlEplbCommunicator(EplbCommunicator):
         pass
 
     def set_transfer_context(self, old_indices: np.ndarray, layer_idx: int) -> None:
-        # Pre-compute expert_id -> src_row mapping for every rank so that
-        # add_recv can immediately issue NIXL READs.
+        self._ensure_remote_state()
         assert not self._xfer_entries, (
             f"set_transfer_context() called with {len(self._xfer_entries)} "
             f"pending transfers from layer {self._layer_idx}; "
@@ -523,6 +554,21 @@ 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() "
@@ -531,13 +577,7 @@ class NixlEplbCommunicator(EplbCommunicator):
         try:
             self._wait_for_all_transfers([x[2] for x in self._xfer_entries])
 
-            # 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),
-            )
+            self._post_read_barrier()
         finally:
             for local_h, remote_h, xfer_h in self._xfer_entries:
                 with contextlib.suppress(Exception):
@@ -628,11 +668,13 @@ def create_eplb_communicator(
             device and CPU communication groups.
         backend: Communicator backend name (``"torch_nccl"``,
             ``"torch_gloo"``, ``"pynccl"``, or ``"nixl"``).
-            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.
+            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.
         expert_weights: Expert weight tensors for *all* MoE layers.
             Shape ``(num_layers)(num_tensors_per_layer)``.
             NixlEplbCommunicator registers all layers with NIXL for
@@ -686,18 +728,21 @@ def create_eplb_communicator(
 
     is_stateless = isinstance(group_coordinator, StatelessGroupCoordinator)
     if is_stateless:
-        if backend not in ("torch_nccl", "pynccl"):
+        if backend == "nixl":
+            pass  # handled below with defer_remote_setup=True
+        elif backend not in ("torch_nccl", "pynccl"):
             raise ValueError(
-                f"Elastic EP requires 'torch_nccl' or 'pynccl' EPLB communicator "
-                f"(got '{backend}')."
+                f"Elastic EP requires 'torch_nccl', 'pynccl', or 'nixl' "
+                f"EPLB communicator (got '{backend}')."
             )
-        if backend == "torch_nccl":
-            logger.warning(
-                "Stateless elastic EP requires PyNCCL backend. "
-                "Forcing EPLB communicator to 'pynccl'."
-            )
-            backend = "pynccl"
-        return _create_pynccl()
+        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 == "nixl":
         if not has_nixl():
@@ -714,6 +759,7 @@ 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(

vllm/distributed/eplb/eplb_state.py

@@ -27,6 +27,7 @@ physical experts.
 """
 
 import threading
+import time
 from collections.abc import Sequence
 from dataclasses import dataclass
 
@@ -825,6 +826,45 @@ 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)
@@ -850,8 +890,9 @@ 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=get_ep_group().device_group)
+            all_reduce(tensor_list[0], group=ep_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), (
@@ -863,7 +904,6 @@ 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 = []

vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py

@@ -387,40 +387,28 @@ class DecodeBenchConnectorWorker:
 
             kv_cache = self.kv_caches[layer_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,
-                )
+            # 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:
-                # Constant fill value
-                fill_values = torch.full(
-                    (len(valid_block_ids),) + block_shape,
-                    self.fill_mean,
-                    dtype=kv_cache.dtype,
-                    device=kv_cache.device,
+                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__,
                 )
-
-            # Batch fill operation
-            kv_cache[valid_block_ids] = fill_values
+                continue
 
         logger.debug(
             "DecodeBenchConnector: Filled %d blocks in group %d with %s values "
@@ -431,3 +419,62 @@ 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)

vllm/model_executor/kernels/linear/__init__.py

@@ -115,6 +115,7 @@ from vllm.model_executor.kernels.linear.nvfp4.fbgemm import (
 from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
     FlashInferB12xNvFp4LinearKernel,
     FlashInferCudnnNvFp4LinearKernel,
+    FlashInferCuteDslNvFp4LinearKernel,
     FlashInferCutlassNvFp4LinearKernel,
     FlashInferTrtllmNvFp4LinearKernel,
 )
@@ -209,6 +210,9 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
         FlashInferCutlassNvFp4LinearKernel,
         FlashInferMxFp4LinearKernel,
     },
+    "flashinfer_cutedsl": {
+        FlashInferCuteDslNvFp4LinearKernel,
+    },
     "flashinfer_trtllm": {
         FlashInferTrtllmNvFp4LinearKernel,
     },
@@ -399,6 +403,7 @@ _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.
@@ -1038,6 +1043,7 @@ __all__ = [
     "CutlassNvFp4LinearKernel",
     "EmulationNvFp4LinearKernel",
     "FbgemmNvFp4LinearKernel",
+    "FlashInferCuteDslNvFp4LinearKernel",
     "FlashInferB12xNvFp4LinearKernel",
     "FlashInferCutlassNvFp4LinearKernel",
     "FlashInferTrtllmNvFp4LinearKernel",

vllm/model_executor/kernels/linear/nvfp4/flashinfer.py

@@ -28,6 +28,72 @@ 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."""
 

vllm/model_executor/layers/attention/mla_attention.py

@@ -208,6 +208,7 @@ 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,
@@ -319,6 +320,22 @@ 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.
 
@@ -369,7 +386,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 = cache_config.cache_dtype
+            kv_cache_dtype: CacheDType = cache_config.cache_dtype
             calculate_kv_scales = cache_config.calculate_kv_scales
         else:
             kv_cache_dtype = "auto"
@@ -393,24 +410,22 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 num_heads=self.num_heads,
             )
 
-        # 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"
+        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
             logger.info_once(
-                "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "
-                "fp8 kv-cache format, please set `--attention-backend "
-                "FLASHINFER_MLA_SPARSE`"
+                "Using %s KV cache format for %s backend.",
+                kv_cache_dtype,
+                self.attn_backend.get_name(),
             )
 
         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(

vllm/model_executor/layers/fused_moe/deep_gemm_utils.py

@@ -23,24 +23,82 @@ 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,
-):
-    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
-        )
+) -> int:
+    """Return ``M_sum`` only (backward-compat wrapper).
 
-    # 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)
+    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
+    )
     return M_sum
 
 
@@ -51,12 +109,6 @@ 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,
@@ -65,6 +117,7 @@ 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)
 
@@ -74,7 +127,8 @@ def _fwd_kernel_ep_scatter_1(
         mask=offset_cumsum < num_experts,
         other=0,
     )
-    tokens_per_expert = round_up_128(tokens_per_expert)
+    # 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
     cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
 
     # Extract this block's offset from the register vector (warp shuffle,
@@ -227,10 +281,12 @@ 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 = 128  # token num of per expert is aligned to 128
+    # BLOCK_E is the m_indices fill-loop tile (masked), independent of align_m.
+    BLOCK_E = 128
     BLOCK_D = block_size  # block size of activation-scale quantization
     num_warps = 8
     num_experts = num_recv_tokens_per_expert.shape[0]
@@ -238,7 +294,7 @@ def ep_scatter(
     # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
     grid = num_experts
 
-    assert m_indices.shape[0] % BLOCK_E == 0
+    assert m_indices.shape[0] % align_m == 0
     assert expert_start_loc.shape[0] == num_experts
 
     # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
@@ -253,6 +309,7 @@ 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)
@@ -418,7 +475,7 @@ def deepgemm_moe_permute(
     if block_size is not None:
         block_k = block_size
 
-    M_sum = compute_aligned_M(
+    M_sum, align_used = compute_aligned_M_and_alignment(
         M=topk_ids.size(0),
         num_topk=topk_ids.size(1),
         local_num_experts=local_num_experts,
@@ -482,11 +539,12 @@ 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
+    return aq_out, aq_scale_out, expert_ids, inv_perm, align_used
 
 
 def deepgemm_unpermute_and_reduce(

vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py

@@ -318,11 +318,12 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):
 
     def supports_packed_ue8m0_act_scales(self) -> bool:
         """
-        DeepGemm supports packed ue8m0 activation scales format in devices == sm100
+        DeepGemm supports packed ue8m0 activation scales on Blackwell-family
+        GPUs (SM100 datacenter and SM120 consumer).
         """
-        return (
-            is_deep_gemm_e8m0_used()
-            and current_platform.is_device_capability_family(100)
+        return is_deep_gemm_e8m0_used() and (
+            current_platform.is_device_capability_family(100)
+            or current_platform.is_device_capability_family(120)
         )
 
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py

@@ -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,
+    compute_aligned_M_and_alignment,
     deepgemm_moe_permute,
     deepgemm_unpermute_and_reduce,
 )
@@ -43,6 +43,7 @@ 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
 
@@ -210,10 +211,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 = compute_aligned_M(
+        M_sum, align_used = compute_aligned_M_and_alignment(
             M, topk, local_num_experts, block_m, expert_tokens_meta
         )
-        assert M_sum % block_m == 0
+        assert M_sum % align_used == 0
 
         activation_out_dim = self.adjust_N_for_activation(N, activation)
         workspace1 = (M_sum, max(activation_out_dim, K))
@@ -316,7 +317,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
 
         assert w2.size(1) == K
 
-        M_sum = compute_aligned_M(
+        M_sum, _ = compute_aligned_M_and_alignment(
             M=topk_ids.size(0),
             num_topk=topk_ids.size(1),
             local_num_experts=local_num_experts,
@@ -327,7 +328,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 = deepgemm_moe_permute(
+        a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
             aq=a1q,
             aq_scale=a1q_scale,
             topk_ids=topk_ids,
@@ -349,23 +350,35 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
             else {}
         )
 
-        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
-        )
+        # 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,
+            )
 
-        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)
@@ -384,7 +397,8 @@ 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 SM100+ (Blackwell).
+    MXFP4 (FP4 E2M1 packed as uint8) weights. Requires Blackwell-family
+    GPUs (SM100 datacenter or SM120 consumer).
     """
 
     # FP8 activation block size (hardcoded since mxfp4_w4a8 quant config
@@ -409,9 +423,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)
+        return is_deep_gemm_supported() and (
+            current_platform.is_device_capability_family(100)
+            or current_platform.is_device_capability_family(120)
         )
 
     @staticmethod
@@ -454,10 +468,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 = compute_aligned_M(
+        M_sum, align_used = compute_aligned_M_and_alignment(
             M, topk, local_num_experts, block_m, expert_tokens_meta
         )
-        assert M_sum % block_m == 0
+        assert M_sum % align_used == 0
 
         activation_out_dim = self.adjust_N_for_activation(N, activation)
         workspace1 = (M_sum, max(activation_out_dim, K))
@@ -533,7 +547,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
         if global_num_experts == -1:
             global_num_experts = local_num_experts
 
-        M_sum = compute_aligned_M(
+        M_sum, _ = compute_aligned_M_and_alignment(
             M=topk_ids.size(0),
             num_topk=topk_ids.size(1),
             local_num_experts=local_num_experts,
@@ -544,7 +558,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 = deepgemm_moe_permute(
+        a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
             aq=a1q,
             aq_scale=a1q_scale,
             topk_ids=topk_ids,
@@ -555,37 +569,40 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
         )
         assert a1q.size(0) == M_sum
 
-        # 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),
-        )
+        # 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),
+            )
 
-        # 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)

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

@@ -57,6 +57,40 @@ 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+)
@@ -652,15 +686,18 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
     """Convert loaded weights into backend-specific kernel format."""
 
     if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
-        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-            _upcast_e8m0_to_fp32,
+        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
+            w13_weight,
+            w2_weight,
+            w13_weight_scale,
+            w2_weight_scale,
         )
 
         return (
             w13_weight.data,
             w2_weight.data,
-            _upcast_e8m0_to_fp32(w13_weight_scale.data),
-            _upcast_e8m0_to_fp32(w2_weight_scale.data),
+            w13_weight_scale,
+            w2_weight_scale,
             w13_bias,
             w2_bias,
         )
@@ -1195,17 +1232,18 @@ def convert_weight_to_mxfp4_moe_kernel_format(
     """
 
     if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
-        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-            _upcast_e8m0_to_fp32,
+        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
+            w13_weight,
+            w2_weight,
+            w13_weight_scale,
+            w2_weight_scale,
         )
 
-        # Weights stay as uint8 packed FP4 — no layout change needed.
-        # Convert E8M0 uint8 scales to float32.
         return (
             w13_weight.data,
             w2_weight.data,
-            _upcast_e8m0_to_fp32(w13_weight_scale.data),
-            _upcast_e8m0_to_fp32(w2_weight_scale.data),
+            w13_weight_scale,
+            w2_weight_scale,
             w13_bias,
             w2_bias,
         )

vllm/model_executor/layers/fused_moe/routed_experts.py

@@ -275,6 +275,12 @@ 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,
@@ -288,10 +294,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] = loaded_weight
+            param_data[expert_id][idx] = self._to_scalar(loaded_weight)
         # If we are in the row parallel case (down_proj)
         elif shard_id == "w2":
-            param_data[expert_id] = loaded_weight
+            param_data[expert_id] = self._to_scalar(loaded_weight)
 
     def _load_combined_w13_weight_scale(
         self,
@@ -525,7 +531,7 @@ class RoutedExperts(PluggableLayer):
         param_data = param.data
 
         # Input scales can be loaded directly and should be equal.
-        param_data[expert_id] = loaded_weight
+        param_data[expert_id] = self._to_scalar(loaded_weight)
 
     def _load_g_idx(
         self,
@@ -692,7 +698,9 @@ 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] = loaded_weight.reshape(())
+                param.data[scale_expert_id][scale_shard_id] = self._to_scalar(
+                    loaded_weight
+                )
                 return True if return_success else None
 
             if (

vllm/model_executor/layers/quantization/humming.py

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

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -1058,6 +1058,34 @@ 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,
@@ -1073,13 +1101,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, got {ws.dtype} instead"
+            f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
         )
         if use_e8m0:
             requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)
@@ -1094,16 +1122,12 @@ 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])
-        # 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,
+        dg_ws = deepgemm_post_process_weight_scale_block(
+            ws=ws,
             mn=r,
             k=d,
-            recipe=(1, quant_block_shape[0], quant_block_shape[1]),
+            quant_block_shape=quant_block_shape,
             num_groups=g,
-            is_sfa=False,
         )
         return wq, dg_ws
 
@@ -1113,22 +1137,12 @@ def deepgemm_post_process_fp8_weight_block(
         wq = wq.unsqueeze(0)
         ws = ws.unsqueeze(0)
 
-    # 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,
+    dg_ws = deepgemm_post_process_weight_scale_block(
+        ws=ws,
         mn=wq.size(1),
         k=wq.size(2),
-        recipe=recipe,
+        quant_block_shape=quant_block_shape,
         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:

vllm/model_executor/warmup/deep_gemm_warmup.py

@@ -12,7 +12,9 @@ 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
+from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
+    compute_aligned_M_and_alignment,
+)
 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,
@@ -25,6 +27,7 @@ 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
@@ -238,7 +241,7 @@ def _get_grouped_gemm_params(
     w2: torch.Tensor,
     num_topk: int,
     max_tokens: int,
-) -> tuple[int, int, torch.Tensor]:
+) -> tuple[int, int, list[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]
@@ -248,19 +251,46 @@ def _get_grouped_gemm_params(
     # Assumes all ranks have the same max_num_batched_tokens
     max_tokens = get_dp_group().world_size * max_tokens
 
-    # 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 = _generate_optimal_warmup_m_values(
+        max_tokens,
+        max(w1.size(1), w2.size(1)),
+        device,
     )
-    # 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)
+    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, []
 
-    return MAX_M, block_m, expert_ids
+    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
 
 
 def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
@@ -278,7 +308,11 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
     ):
         return
 
-    MAX_M, block_m, expert_ids = _get_grouped_gemm_params(w1, w2, num_topk, max_tokens)
+    MAX_M, block_m, warmup_cases = _get_grouped_gemm_params(
+        w1, w2, num_topk, max_tokens
+    )
+    if not warmup_cases:
+        return
     device = w1.device
 
     def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
@@ -289,15 +323,14 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
         )
         out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
 
-        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],
-            )
+        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,
+                )
             if pbar is not None:
                 pbar.update(1)
 
@@ -350,8 +383,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
-            MAX_M, block_m, _ = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
-            n_values = (MAX_M - block_m) // block_m + 1
+            _, _, warmup_cases = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
+            n_values = len(warmup_cases)
             if w13.size() not in seen_grouped_sizes:
                 total += n_values
                 seen_grouped_sizes.add(w13.size())

vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py

@@ -0,0 +1,226 @@
+# 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,
+    )

vllm/model_executor/warmup/flashinfer_autotune_cache.py

@@ -0,0 +1,56 @@
+# 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

vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py

@@ -0,0 +1,255 @@
+# 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,
+            )

vllm/model_executor/warmup/kernel_warmup.py

@@ -6,16 +6,24 @@ 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
@@ -27,36 +35,26 @@ 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
@@ -147,7 +145,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)
 
@@ -183,9 +181,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
             "Falling back to default tactics."
         )
     else:
-        if not is_leader and world.local_rank == 0:
-            with open(cache_path, "wb") as f:
-                f.write(tune_results)
+        write_flashinfer_autotune_cache(cache_path, tune_results)
         world.barrier()
         from flashinfer.autotuner import AutoTuner
 

vllm/models/deepseek_v4/attention.py

@@ -62,23 +62,22 @@ logger = init_logger(__name__)
 
 
 def _resolve_dsv4_kv_cache_dtype(
-    use_flashmla_fp8_layout: bool,
+    use_fp8_ds_mla_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
-    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.
+    ``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.
     """
-    if use_flashmla_fp8_layout:
+    if use_fp8_ds_mla_layout:
         # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
         assert kv_cache_dtype.startswith("fp8"), (
-            f"DeepseekV4 FlashMLA fp8 layout only supports fp8 kv-cache, "
+            f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, "
             f"got {kv_cache_dtype}"
         )
         if kv_cache_dtype != "fp8_ds_mla":
@@ -100,18 +99,20 @@ 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`` / ``DeepseekV4FlashInferMLAAttention``
-    (CUDA) or ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the
-    platform-specific deepseek_v4 model module. The base is never instantiated
-    directly.
+    subclass — ``DeepseekV4FlashMLAAttention`` /
+    ``DeepseekV4FlashInferSM120Attention`` /
+    ``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)
-    # = 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
+    # = 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
     # 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.
@@ -145,6 +146,10 @@ 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,
@@ -276,13 +281,10 @@ 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 (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.
+        # Resolve the kv-cache dtype from this backend's block format. 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.use_flashmla_fp8_layout, cache_config.cache_dtype, cache_config
+            self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
         )
 
         self.swa_cache_layer = DeepseekV4SWACache(
@@ -539,7 +541,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
 
         # kv is unchanged; attention reads kv solely via swa_kv_cache.
         if cache_dtype == torch.uint8:
-            # Legacy FlashMLA UE8M0 paged path. Horizontally fused:
+            # fp8_ds_mla 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.
@@ -557,10 +559,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
                 swa_metadata.block_size,
             )
 
-        # 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.
+        # 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.
         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:
@@ -601,18 +603,18 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             self.compress_ratio <= 1
         ):  # SWA part. Allocated separately as DeepseekV4SWACache.
             return None
-        # 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"
+        # 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"
         return MLAAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=1,
             head_size=self.head_dim,
-            dtype=torch.uint8 if is_flashmla else self.kv_cache_torch_dtype,
+            dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
             compress_ratio=self.compress_ratio,
             cache_dtype_str=self.kv_cache_dtype,
-            alignment=576 if is_flashmla else None,  # FlashMLA needs 576B
+            alignment=576 if uses_fp8_ds_mla_layout else None,
             model_version="deepseek_v4",
         )
 

vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py

@@ -40,9 +40,18 @@ def _fused_inv_rope_fp8_quant_per_head(
     USE_GDC: tl.constexpr,
     launch_pdl: tl.constexpr,  # triton metadata
 ):
-    # int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
+    # 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.
     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

vllm/models/deepseek_v4/compressor.py

@@ -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:
-        # 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"
+        # 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"
         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 is_flashmla else None,
+            alignment=576 if uses_fp8_ds_mla_layout 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
 
-        # FlashInfer V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
-        # legacy FlashMLA path uses the UE8M0 paged uint8 layout.
+        # Plain-row V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
+        # fp8_ds_mla 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 legacy UE8M0
-            # layout and the FlashInfer full-cache layout. The full-cache flags
+            # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla
+            # layout and the plain 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(

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -1,13 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""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.
-"""
+"""DeepSeek V4 FlashInfer sparse MLA backend."""
 
 from typing import TYPE_CHECKING, ClassVar, cast
 
@@ -18,6 +11,7 @@ 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,
@@ -27,13 +21,14 @@ 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] = {}
 
@@ -51,34 +46,113 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
 
 
 class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
-    """Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl.
+    """FlashInfer backend using the DSv4 sparse metadata/cache layout.
 
-    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"]
+    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]
 
     @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 DeepSeek V4."""
+    """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
 
     backend_cls = DeepseekV4FlashInferMLASparseBackend
-    # 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
+    use_fp8_ds_mla_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
@@ -106,8 +180,6 @@ 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(
@@ -125,9 +197,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
             torch.tensor([fp8_kv_scale], dtype=torch.float32),
             persistent=False,
         )
-        # 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.
+        # TRTLLM-gen takes scalar scale args on a distinct C++ path vs
+        # one-element tensors, so these are Python floats.
         self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
         self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
 
@@ -387,9 +458,8 @@ 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 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.
+        # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
+        # uniform-q batches, and this avoids flattening mixed batches into one call.
         if num_decode_tokens > 0:
             decode_cu = query_start_loc[: num_decodes + 1]
             decode_cu_cpu = query_start_loc_cpu[: num_decodes + 1]
@@ -434,3 +504,379 @@ 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,
+            )

vllm/models/deepseek_v4/nvidia/model.py

@@ -60,9 +60,11 @@ 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
@@ -736,14 +738,34 @@ 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.
 
-    An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the
-    FlashInfer TRTLLM-gen path; otherwise the FlashMLA path is used.
+    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.
     """
-    if (
-        vllm_config.attention_config.backend
-        == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4
+    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,
     ):
+        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
 
 

vllm/models/deepseek_v4/sparse_mla.py

@@ -86,6 +86,10 @@ 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]

vllm/platforms/cuda.py

@@ -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, TypeVar
+from typing import TYPE_CHECKING, NamedTuple, TypeVar
 
 import torch
 from torch.distributed import PrefixStore, ProcessGroup
@@ -31,6 +31,7 @@ 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
@@ -126,6 +127,11 @@ 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,
@@ -153,6 +159,21 @@ 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:
@@ -340,7 +361,7 @@ class CudaPlatformBase(Platform):
         attn_selector_config: AttentionSelectorConfig,
         num_heads: int | None = None,
     ) -> tuple[
-        list[tuple[AttentionBackendEnum, int]],
+        list[_BackendCandidate],
         dict[AttentionBackendEnum, tuple[int, list[str]]],
     ]:
         valid_backends_priorities = []
@@ -354,7 +375,7 @@ class CudaPlatformBase(Platform):
         )
         for priority, backend in enumerate(backend_priorities):
             try:
-                backend_class = backend.get_class()
+                backend_class = _get_attn_backend_class(backend)
                 invalid_reasons_i = backend_class.validate_configuration(
                     device_capability=device_capability,
                     **attn_selector_config._asdict(),
@@ -364,7 +385,9 @@ class CudaPlatformBase(Platform):
             if invalid_reasons_i:
                 invalid_reasons[backend] = (priority, invalid_reasons_i)
             else:
-                valid_backends_priorities.append((backend, priority))
+                valid_backends_priorities.append(
+                    _BackendCandidate(backend_class, backend, priority)
+                )
 
         return valid_backends_priorities, invalid_reasons
 
@@ -381,7 +404,7 @@ class CudaPlatformBase(Platform):
         # First try checking just the selected backend, if there is one.
         if selected_backend is not None:
             try:
-                backend_class = selected_backend.get_class()
+                backend_class = _get_attn_backend_class(selected_backend)
                 invalid_reasons = backend_class.validate_configuration(
                     device_capability=device_capability,
                     **attn_selector_config._asdict(),
@@ -395,7 +418,7 @@ class CudaPlatformBase(Platform):
                 )
             else:
                 logger.info("Using %s backend.", selected_backend)
-                return selected_backend.get_path()
+                return _backend_cls_path(backend_class)
 
         # No selected backend or the selected backend is invalid,
         # so we try finding a valid backend.
@@ -425,13 +448,13 @@ class CudaPlatformBase(Platform):
 
         # We have found some valid backends. Select the one with the
         # highest priority.
-        sorted_indices = sorted(
-            range(len(valid_backends_priorities)),
-            key=lambda i: valid_backends_priorities[i][1],
+        selected_candidate = min(
+            valid_backends_priorities,
+            key=lambda candidate: candidate.priority,
         )
-        selected_index = sorted_indices[0]
-        selected_backend = valid_backends_priorities[selected_index][0]
-        selected_priority = valid_backends_priorities[selected_index][1]
+        selected_backend_class = selected_candidate.backend_class
+        selected_backend = selected_candidate.backend
+        selected_priority = selected_candidate.priority
 
         # If the user specified --block-size (but not --attention-backend),
         # check whether that constraint precluded any higher-priority backends.
@@ -457,10 +480,14 @@ class CudaPlatformBase(Platform):
         logger.info_once(
             "Using %s attention backend out of potential backends: %s.",
             selected_backend.name,
-            "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
+            "["
+            + ", ".join(
+                f"'{candidate.backend.name}'" for candidate in valid_backends_priorities
+            )
+            + "]",
         )
 
-        return selected_backend.get_path()
+        return _backend_cls_path(selected_backend_class)
 
     @classmethod
     def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]:
@@ -635,7 +662,11 @@ 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)
+        return (
+            cls.is_device_capability(90)
+            or cls.is_device_capability_family(100)
+            or cls.is_device_capability_family(120)
+        )
 
     @classmethod
     def is_integrated_gpu(cls, device_id: int = 0) -> bool:

vllm/platforms/rocm.py

@@ -27,8 +27,10 @@ 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,
@@ -167,6 +169,14 @@ 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.
@@ -726,8 +736,22 @@ class RocmPlatform(Platform):
 
     @classmethod
     def get_device_total_memory(cls, device_id: int = 0) -> int:
-        device_props = torch.cuda.get_device_properties(device_id)
-        return device_props.total_memory
+        # 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
 
     @classmethod
     def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:

vllm/utils/deep_gemm.py

@@ -5,6 +5,7 @@
 Users of vLLM should always import **only** these wrappers.
 """
 
+import contextlib
 import functools
 import importlib
 import os
@@ -37,7 +38,10 @@ 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):
+    if not (
+        current_platform.is_device_capability_family(100)
+        or current_platform.is_device_capability_family(120)
+    ):
         return False
     return model_type in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES
 
@@ -71,7 +75,10 @@ class DeepGemmQuantScaleFMT(Enum):
 
         cls._oracle_cache = (  # type: ignore
             cls.UE8M0
-            if current_platform.is_device_capability_family(100)
+            if (
+                current_platform.is_device_capability_family(100)
+                or current_platform.is_device_capability_family(120)
+            )
             else cls.FLOAT32_CEIL_UE8M0
         )
 
@@ -138,7 +145,15 @@ _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
@@ -203,7 +218,11 @@ 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
@@ -218,6 +237,9 @@ 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
 
@@ -258,9 +280,19 @@ 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()
 
 
@@ -280,7 +312,6 @@ 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:
@@ -289,6 +320,70 @@ 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()
@@ -297,6 +392,48 @@ 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:
@@ -601,4 +738,8 @@ __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",
 ]

vllm/utils/flashinfer.py

@@ -72,11 +72,10 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
     )
 
 
-def _missing_dsv4_sparse_mla(*_: Any, **__: Any) -> NoReturn:
+def _missing_sparse_mla(*_: Any, **__: Any) -> NoReturn:
     raise RuntimeError(
-        "flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4 is not available. "
-        "Install a FlashInfer build that includes DeepSeek V4 sparse MLA "
-        "TRTLLM-GEN support."
+        "FlashInfer sparse MLA decode APIs are not available. "
+        "Install a FlashInfer build that includes sparse MLA decode support."
     )
 
 
@@ -149,14 +148,18 @@ flashinfer_b12x_fused_moe = _lazy_import_wrapper(
 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
     "flashinfer", "trtllm_fp4_block_scale_moe"
 )
-# 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.mla",
-    "trtllm_batch_decode_sparse_mla_dsv4",
-    fallback_fn=_missing_dsv4_sparse_mla,
+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,
 )
+flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper(
+    "flashinfer.decode",
+    "trtllm_batch_decode_sparse_mla_dsv4",
+    fallback_fn=_missing_sparse_mla,
+)
+
+
 # Special case for autotune since it returns a context manager
 autotune = _lazy_import_wrapper(
     "flashinfer.autotuner",
@@ -209,6 +212,26 @@ 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."""
@@ -988,6 +1011,7 @@ __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",

vllm/v1/attention/backend.py

@@ -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
-    metadata (e.g. C128A topk indices for DeepSeek V4)."""
+    sparse metadata for DeepSeek V4 C128A layers."""
 
     is_prefilling: torch.Tensor | None = None
     """(batch_size,) bool tensor: True if request is still in prefill phase

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -1,23 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""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
-"""
+"""FlashInfer sparse MLA attention backend."""
 
 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
@@ -52,34 +41,13 @@ logger = init_logger(__name__)
 FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
 
 
-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]
+class _FlashInferMLASparseBackendBase(AttentionBackend):
+    """Common metadata for concrete FlashInfer sparse MLA backends."""
 
     @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
@@ -96,9 +64,29 @@ class FlashInferMLASparseBackend(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
@@ -114,10 +102,15 @@ class FlashInferMLASparseBackend(AttentionBackend):
         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)
@@ -146,6 +139,102 @@ class FlashInferMLASparseBackend(AttentionBackend):
         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."""
@@ -353,6 +442,8 @@ 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),

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py

@@ -0,0 +1,155 @@
+# 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

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -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: legacy FlashMLA UE8M0 paged layout. bfloat16 / float8_e4m3fn:
-        # FlashInfer contiguous full-cache layout.
+        # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
+        # 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:
-        # 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"
+        # 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"
         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 is_flashmla else None,
+            alignment=576 if uses_fp8_ds_mla_layout else None,
             model_version="deepseek_v4",
         )
 
@@ -164,6 +164,11 @@ 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
@@ -343,6 +348,20 @@ 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,
@@ -402,6 +421,29 @@ 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,
             )
 
@@ -431,6 +473,16 @@ 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,
@@ -465,6 +517,7 @@ 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:
@@ -489,7 +542,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         Returns a dict of keyword arguments to pass to the
         DeepseekSparseSWAMetadata constructor.
 
-        Note: C128A topk indices are computed by the FlashMLASparse builder
+        Note: C128A sparse metadata is computed by the FlashMLASparse builder
         (which owns the C128A block_table), not here.
         """
         result: dict[str, torch.Tensor | int | None] = {}
@@ -539,10 +592,14 @@ 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 + 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)
+    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)
 
     query_len = qsl_end - qsl_start
     prefix_len = seq_len - query_len
@@ -551,7 +608,7 @@ def _compute_prefill_metadata_kernel(
     tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["token_offset"])
 def _compute_swa_indices_and_lens_kernel(
     swa_indices_ptr,
     swa_indices_stride,
@@ -564,12 +621,14 @@ def _compute_swa_indices_and_lens_kernel(
     block_table_ptr,
     block_table_stride,
     block_size,
+    token_offset,
     TRITON_BLOCK_SIZE: tl.constexpr,
 ):
-    token_idx = tl.program_id(0)
+    pid = tl.program_id(0)
+    token_idx = pid + token_offset
     is_valid = tl.load(is_valid_token_ptr + token_idx)
     if not is_valid:
-        tl.store(swa_lens_ptr + token_idx, 0)
+        tl.store(swa_lens_ptr + pid, 0)
         return
 
     req_idx = tl.load(token_to_req_indices_ptr + token_idx)
@@ -586,7 +645,7 @@ def _compute_swa_indices_and_lens_kernel(
     end_pos = pos + 1
 
     swa_len = end_pos - start_pos
-    tl.store(swa_lens_ptr + token_idx, swa_len)
+    tl.store(swa_lens_ptr + pid, swa_len)
 
     for i in range(0, window_size, TRITON_BLOCK_SIZE):
         offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
@@ -602,7 +661,7 @@ def _compute_swa_indices_and_lens_kernel(
 
         slot_ids = tl.where(offset < swa_len, slot_ids, -1)
         tl.store(
-            swa_indices_ptr + token_idx * swa_indices_stride + offset,
+            swa_indices_ptr + pid * swa_indices_stride + offset,
             slot_ids,
             mask=offset < window_size,
         )

vllm/v1/attention/backends/registry.py

@@ -71,7 +71,11 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     )
     FLASHINFER_MLA_SPARSE = (
         "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
-        "FlashInferMLASparseBackend"
+        "FlashInferMLASparseTRTLLMBackend"
+    )
+    FLASHINFER_MLA_SPARSE_SM120 = (
+        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
+        "FlashInferMLASparseSM120Backend"
     )
     TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
     CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"

vllm/v1/attention/ops/flashmla.py

@@ -73,7 +73,7 @@ def is_flashmla_sparse_supported() -> tuple[bool, str | None]:
     ):
         return (
             False,
-            "FlashMLA Sparse is only supported on Hopper and Blackwell devices.",
+            "FlashMLA Sparse is only supported on Hopper and Blackwell DC devices.",
         )
     return True, None
 

vllm/v1/core/sched/scheduler.py

@@ -1872,6 +1872,11 @@ 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):
@@ -1884,13 +1889,12 @@ 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
+                start_pos + num_tokens + spec_lookahead
                 <= request.num_computed_tokens - request.num_output_placeholders
             ):
-                # 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.
+                # 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.
                 self.encoder_cache_manager.free_encoder_input(request, input_id)
 
     def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:

vllm/v1/simple_kv_offload/copy_backend.py

@@ -12,6 +12,8 @@ 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,
@@ -43,8 +45,20 @@ class DmaCopyBackend:
         self._load_stream = load_stream
         self._store_stream = store_stream
 
-        self._store_params = build_params(gpu_caches, cpu_caches, store_stream)
-        self._load_params = build_params(cpu_caches, gpu_caches, load_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._queue = queue.SimpleQueue()
         self._thread = threading.Thread(
@@ -61,11 +75,20 @@ 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)
+            (
+                src_blocks,
+                dst_blocks,
+                params,
+                is_store,
+                event_idx,
+                events_list,
+                wait_event,
+            )
         )
 
     def shutdown(self) -> None:
@@ -89,9 +112,19 @@ class DmaCopyBackend:
             item = q.get()
             if item is None:
                 return
-            src_blocks, dst_blocks, params, is_store, event_idx, events_list = item
-            copy_blocks(src_blocks, dst_blocks, params)
+            (
+                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)
+            copy_blocks(src_blocks, dst_blocks, params)
             event = torch.Event()
             event.record(stream)
             events_list.append((event_idx, event))

vllm/v1/simple_kv_offload/cuda_mem_ops.py

@@ -13,6 +13,12 @@ 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.
@@ -106,8 +112,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 with srcAccessOrder=ANY. Unused on
-    # ROCm (7.2.1 or 7.2.2) because the current runtime rejects numAttrs > 0.
+    # CUDA only: one attributes entry carrying srcAccessOrder. 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
@@ -120,6 +126,7 @@ 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:
@@ -137,10 +144,7 @@ def build_params(
         dst_bases.append(d.data_ptr())
         bpb.append(s_bpb)
 
-    # ``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)
+    attrs = _CUmemcpyAttributes(srcAccessOrder=src_access_order)
 
     return BatchMemcpyParams(
         src_bases=np.array(src_bases, dtype=np.uint64),

vllm/v1/simple_kv_offload/worker.py

@@ -57,6 +57,10 @@ 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()
@@ -206,9 +210,11 @@ class SimpleCPUOffloadWorker:
     ) -> tuple[set[str] | None, set[str] | None]:
         """Submit transfers and report completed events to the scheduler.
 
-        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.
+        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.
 
         Returns:
             tuple of (finished_sending, finished_recving).
@@ -218,7 +224,6 @@ 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,
@@ -227,14 +232,17 @@ 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

vllm/v1/worker/gpu/mm/encoder_runner.py

@@ -68,15 +68,19 @@ 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]:
-        is_prefilling = (computed_prefill_lens < prefill_lens).tolist()
-        all_decode = not any(is_prefilling)
-        if all_decode:
+        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():
             # 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()
 
@@ -89,11 +93,12 @@ 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=query_start[i],
-                end=query_end[i],
+                mm_features, start=cur_query_start, end=cur_query_end
             )
             for idx in range(lo, hi):
                 mm_feature = mm_features[idx]
@@ -101,8 +106,8 @@ class EncoderRunner:
                 start_pos = pos_info.offset
                 num_encoder_tokens = pos_info.length
 
-                start_idx = max(query_start[i] - start_pos, 0)
-                end_idx = min(query_end[i] - start_pos, num_encoder_tokens)
+                start_idx = max(cur_query_start - start_pos, 0)
+                end_idx = min(cur_query_end - 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)
@@ -114,7 +119,13 @@ class EncoderRunner:
 
                 mm_hash = mm_feature.identifier
                 encoder_output = self.encoder_cache.encoder_outputs.get(mm_hash, None)
-                assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
+                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}.")
 
                 if (is_embed := pos_info.is_embed) is not None:
                     is_embed = is_embed[start_idx:end_idx]
@@ -122,7 +133,7 @@ class EncoderRunner:
                 else:
                     mm_embeds_item = encoder_output[start_idx:end_idx]
 
-                req_start_pos = query_start_loc[i] + start_pos - query_start[i]
+                req_start_pos = query_start_loc[i] + start_pos - cur_query_start
                 is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= (
                     True if is_embed is None else is_embed
                 )

vllm/v1/worker/gpu/model_runner.py

@@ -314,6 +314,7 @@ 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
@@ -1409,8 +1410,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 input_batch.num_scheduled_tokens,
                 input_batch.query_start_loc_np,
                 input_batch.prefill_len_np,
-                # +1 to consider the skew in eagle
-                input_batch.num_computed_prefill_tokens_np + 1,
+                input_batch.num_computed_prefill_tokens_np,
+                # The EAGLE/MTP drafter reads one position ahead of the target.
+                draft_lookahead=1,
             )
 
         # Postprocess results and update request states.

vllm/v1/worker/gpu/model_states/__init__.py

@@ -4,6 +4,7 @@ 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
 
 
@@ -18,13 +19,13 @@ def init_model_state(
         cls = model.get_model_state_cls()
         return cls(vllm_config, model, encoder_cache, device)
 
-    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
+    # 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,
+        )
 
-        return WhisperModelState(vllm_config, model, encoder_cache, device)
+        return EncoderDecoderModelState(vllm_config, model, encoder_cache, device)
 
     if vllm_config.model_config.is_hybrid:
         from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState

vllm/v1/worker/gpu/model_states/default.py

@@ -7,7 +7,8 @@ import torch.nn as nn
 
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.tasks import GenerationTask
+from vllm.model_executor.models.interfaces import supports_multimodal_pruning
+from vllm.multimodal.utils import get_mm_features_in_window
 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
@@ -33,6 +34,7 @@ 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
@@ -40,6 +42,11 @@ 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
@@ -62,28 +69,6 @@ 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
@@ -98,6 +83,82 @@ 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]],
@@ -120,6 +181,13 @@ 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]

vllm/v1/worker/gpu/model_states/encoder_decoder.py

@@ -23,7 +23,7 @@ from vllm.v1.worker.utils import AttentionGroup
 
 
 @dataclass
-class WhisperAttnMetadata(ModelSpecificAttnMetadata):
+class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
     encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
 
     def get_extra_common_attn_kwargs(
@@ -41,7 +41,11 @@ class WhisperAttnMetadata(ModelSpecificAttnMetadata):
         }
 
 
-class WhisperModelState(ModelState):
+class EncoderDecoderModelState(ModelState):
+    """ModelState for cross-attention encoder-decoder models
+    (Whisper, CohereASR, NemotronParse, FireRedLID, ...)
+    """
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -80,9 +84,6 @@ class WhisperModelState(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:
@@ -94,11 +95,11 @@ class WhisperModelState(ModelState):
                 encoder_inputs[req_id] = req_encoder_inputs
         _, mm_kwargs = self.encoder_runner.prepare_mm_inputs(encoder_inputs)
         if mm_kwargs:
-            # 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.
+            # 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.
             self.encoder_outputs = self.encoder_runner.execute_mm_encoder(mm_kwargs)
         else:
             # Decode steps: encoder K/V are in cross-attention KV cache.
@@ -131,7 +132,7 @@ class WhisperModelState(ModelState):
         else:
             num_reqs = input_batch.num_reqs
             num_tokens = input_batch.num_tokens
-        whisper_attn_metadata = WhisperAttnMetadata(
+        enc_dec_attn_metadata = EncoderDecoderAttnMetadata(
             self._get_encoder_seq_lens(
                 input_batch.req_ids, attn_groups, for_capture, num_reqs
             )
@@ -158,7 +159,7 @@ class WhisperModelState(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=whisper_attn_metadata,
+            model_specific_attn_metadata=enc_dec_attn_metadata,
             for_cudagraph_capture=for_capture,
         )
         return attn_metadata

vllm/v1/worker/gpu/model_states/interface.py

@@ -46,9 +46,25 @@ class ModelState(ABC):
     ) -> None:
         raise NotImplementedError
 
-    @abstractmethod
+    model: nn.Module
+
     def get_supported_generation_tasks(self) -> tuple[GenerationTask, ...]:
-        raise NotImplementedError
+        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:
         return None

vllm/v1/worker/gpu/warmup.py

@@ -2,12 +2,14 @@
 # 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,
@@ -18,6 +20,134 @@ 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(

vllm/v1/worker/gpu_model_runner.py

@@ -3152,7 +3152,16 @@ class GPUModelRunner(
 
                 mm_hash = mm_feature.identifier
                 encoder_output = self.encoder_cache.get(mm_hash, None)
-                assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
+                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}.")
 
                 if (is_embed := pos_info.is_embed) is not None:
                     is_embed = is_embed[start_idx:end_idx]