Enable DeepSeek V4 and GLM-5.1 on SM120 (#43477)

Signed-off-by: Zihua Wu <[email protected]> Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]> Co-authored-by: Yongye Zhu <[email protected]>
2026-08-23 22:20:15 +00:00 · 2026-06-22 11:54:14 -07:00
parent 3ce15fd574
commit 44d95069e9
37 changed files with 2347 additions and 476 deletions
@@ -1,5 +1,5 @@
 group: Expert Parallelism
-depends_on: 
+depends_on:
  - image-build-xpu
 steps:
 - label: EPLB Algorithm
@@ -1,5 +1,5 @@
 group: Models - Multimodal
-depends_on: 
+depends_on:
  - image-build-xpu
 steps:
 - label: "Multi-Modal Models (Standard) 1: qwen2"
@@ -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
@@ -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.")
@@ -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 |
@@ -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"
@@ -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 == []
@@ -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"
@@ -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
    ]


@@ -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.",
                "",
            ]
        )
@@ -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(
@@ -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(
@@ -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:
@@ -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)
@@ -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,
        )
@@ -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:
@@ -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())
@@ -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,
+    )
@@ -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
@@ -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,
+            )
@@ -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

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

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


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

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