@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 44d95069e9

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

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Gabriel Wu committed Jun 23, 2026 at 02:54 UTC 44d95069e9d6f764bea72f8a9ae6fa7f21187182
37 files changed +2340 -469
.buildkite/intel_jobs/expert_parallelism_intel.yaml
+1 -1
@@ -1,5 +1,5 @@
1 group: Expert Parallelism
2 -depends_on:
2 +depends_on:
3 - image-build-xpu
4 steps:
5 - label: EPLB Algorithm
.buildkite/intel_jobs/models_multimodal_intel.yaml
+1 -1
@@ -1,5 +1,5 @@
1 group: Models - Multimodal
2 -depends_on:
2 +depends_on:
3 - image-build-xpu
4 steps:
5 - label: "Multi-Modal Models (Standard) 1: qwen2"
cmake/external_projects/deepgemm.cmake
+58 -28
@@ -8,43 +8,73 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR})
8 set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR})
9 endif()
10
11 +# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download).
12 +# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW
13 +# disallows one-argument Populate(dep) after Declare; MakeAvailable would run
14 +# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load).
15 if(DEEPGEMM_SRC_DIR)
12 - FetchContent_Declare(
13 - deepgemm
14 - SOURCE_DIR ${DEEPGEMM_SRC_DIR}
15 - CONFIGURE_COMMAND ""
16 - BUILD_COMMAND ""
17 - )
16 + # cmake_path(ABSOLUTE_PATH <var> ...) reads the path from <var>; NORMALIZE is a
17 + # flag (no trailing path argument). Resolve relative paths against vLLM root.
18 + set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}")
19 + cmake_path(ABSOLUTE_PATH _deepgemm_user_src
20 + BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
21 + NORMALIZE)
22 + set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}")
23 + if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}")
24 + message(FATAL_ERROR
25 + "DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'")
26 + endif()
27 + set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}")
28 + message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}")
29 else()
19 - # This ref should be kept in sync with tools/install_deepgemm.sh
20 - FetchContent_Declare(
21 - deepgemm
22 - GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
23 - GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
24 - GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
25 - GIT_PROGRESS TRUE
26 - CONFIGURE_COMMAND ""
27 - BUILD_COMMAND ""
28 - )
29 -endif()
30 + # Keep in sync with tools/install_deepgemm.sh
31 + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git")
32 + set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f")
33
31 -# Use FetchContent_Populate (not MakeAvailable) to avoid processing
32 -# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls.
33 -FetchContent_GetProperties(deepgemm)
34 -if(NOT deepgemm_POPULATED)
35 - FetchContent_Populate(deepgemm)
34 + set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}")
35 + if(NOT _deepgemm_fc_root)
36 + set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps")
37 + endif()
38 + set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src")
39 + set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build")
40 + set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild")
41 +
42 + if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp")
43 + set(deepgemm_SOURCE_DIR "${_deepgemm_src}")
44 + set(deepgemm_BINARY_DIR "${_deepgemm_bin}")
45 + else()
46 + FetchContent_Populate(
47 + deepgemm
48 + SUBBUILD_DIR "${_deepgemm_sub}"
49 + SOURCE_DIR "${_deepgemm_src}"
50 + BINARY_DIR "${_deepgemm_bin}"
51 + GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}"
52 + GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}"
53 + GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
54 + GIT_PROGRESS TRUE
55 + )
56 + endif()
57 + message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
58 endif()
37 -message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}")
59
39 -# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100
60 +# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream),
61 +# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x
62 +# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants.
63 set(DEEPGEMM_SUPPORT_ARCHS)
64 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3)
65 list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a")
66 endif()
44 -if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
45 - list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
46 -elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
47 - list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
67 +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8)
68 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
69 + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f")
70 + else()
71 + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a")
72 + endif()
73 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
74 + list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f")
75 + else()
76 + list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a")
77 + endif()
78 endif()
79
80 cuda_archs_loose_intersection(DEEPGEMM_ARCHS
cmake/external_projects/qutlass.cmake
+38 -16
@@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR})
6 set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
7 endif()
8
9 +# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid.
10 +# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees.
11 if(QUTLASS_SRC_DIR)
10 - FetchContent_Declare(
11 - qutlass
12 - SOURCE_DIR ${QUTLASS_SRC_DIR}
13 - CONFIGURE_COMMAND ""
14 - BUILD_COMMAND ""
15 - )
12 + set(_qutlass_user_src "${QUTLASS_SRC_DIR}")
13 + cmake_path(ABSOLUTE_PATH _qutlass_user_src
14 + BASE_DIRECTORY "${CMAKE_SOURCE_DIR}"
15 + NORMALIZE)
16 + set(QUTLASS_SRC_DIR "${_qutlass_user_src}")
17 + if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}")
18 + message(FATAL_ERROR
19 + "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'")
20 + endif()
21 + set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}")
22 + set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
23 else()
17 - FetchContent_Declare(
18 - qutlass
19 - GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git
20 - GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65
21 - GIT_PROGRESS TRUE
22 - CONFIGURE_COMMAND ""
23 - BUILD_COMMAND ""
24 - )
25 -endif()
24 + set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
25 + set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
26
27 -FetchContent_Populate(qutlass)
27 + set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
28 + if(NOT _qutlass_fc_root)
29 + set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps")
30 + endif()
31 + set(_qutlass_src "${_qutlass_fc_root}/qutlass-src")
32 + set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build")
33 + set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild")
34 +
35 + if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp")
36 + set(qutlass_SOURCE_DIR "${_qutlass_src}")
37 + set(qutlass_BINARY_DIR "${_qutlass_bin}")
38 + else()
39 + FetchContent_Populate(
40 + qutlass
41 + SUBBUILD_DIR "${_qutlass_sub}"
42 + SOURCE_DIR "${_qutlass_src}"
43 + BINARY_DIR "${_qutlass_bin}"
44 + GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}"
45 + GIT_TAG "${_QUTLASS_UPSTREAM_TAG}"
46 + GIT_PROGRESS TRUE
47 + )
48 + endif()
49 +endif()
50
51 if(NOT qutlass_SOURCE_DIR)
52 message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")
docs/design/attention_backends.md
+6 -14
@@ -133,16 +133,6 @@ Priority is **1 = highest** (tried first).
133 | 7 | `FLASHINFER_MLA_SPARSE`**\*** |
134 | 8 | `FLASHMLA_SPARSE` |
135
136 -**Ampere/Hopper (SM 8.x-9.x):**
137 -
138 -| Priority | Backend |
139 -| -------- | ------- |
140 -| 1 | `FLASH_ATTN_MLA` |
141 -| 2 | `FLASHMLA` |
142 -| 3 | `FLASHINFER_MLA` |
143 -| 4 | `TRITON_MLA` |
144 -| 5 | `FLASHMLA_SPARSE` |
145 -
136 > **\*** 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.
137 >
138 > **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
221 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
222 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
223 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
234 -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
224 +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
225 +| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x |
226 | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
227 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
228 | `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
239 `--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,
240 `FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index
241 pipeline (compressor + SWA + indexer, 256-token blocks, head 512);
251 -default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.
242 +default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and
243 +`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.
244
245 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
246 | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
255 -| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
256 -| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
247 +| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x |
248 +| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
249 | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
tests/model_executor/test_flashinfer_autotune_cache.py deleted
-60
@@ -1,60 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import sys
5 -from hashlib import sha256
6 -from pathlib import Path
7 -from types import SimpleNamespace
8 -
9 -from vllm.model_executor.warmup import kernel_warmup
10 -
11 -
12 -def test_resolve_flashinfer_autotune_file_default_layout(
13 - monkeypatch, tmp_path: Path
14 -) -> None:
15 - fake_jit = SimpleNamespace(
16 - env=SimpleNamespace(
17 - FLASHINFER_WORKSPACE_DIR=Path("/flashinfer-cache/0.6.11.post2/103a")
18 - )
19 - )
20 - fake_flashinfer = SimpleNamespace(jit=fake_jit)
21 - monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer)
22 - monkeypatch.setitem(sys.modules, "flashinfer.jit", fake_jit)
23 - monkeypatch.setattr(
24 - kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
25 - )
26 - monkeypatch.setattr(kernel_warmup.envs, "VLLM_CACHE_ROOT", str(tmp_path))
27 - monkeypatch.setattr(kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None)
28 -
29 - runner = SimpleNamespace(vllm_config=SimpleNamespace())
30 - cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
31 -
32 - path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
33 -
34 - assert path == (
35 - tmp_path
36 - / "flashinfer_autotune_cache"
37 - / "0.6.11.post2"
38 - / "103a"
39 - / cache_hash
40 - / "autotune_configs.json"
41 - )
42 - assert path.parent.is_dir()
43 -
44 -
45 -def test_resolve_flashinfer_autotune_file_uses_override_dir(
46 - monkeypatch, tmp_path: Path
47 -) -> None:
48 - monkeypatch.setattr(
49 - kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path)
50 - )
51 - monkeypatch.setattr(
52 - kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"]
53 - )
54 -
55 - runner = SimpleNamespace(vllm_config=SimpleNamespace())
56 - cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest()
57 -
58 - path = kernel_warmup._resolve_flashinfer_autotune_file(runner)
59 -
60 - assert path == tmp_path / cache_hash / "autotune_configs.json"
tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Behavior checks for FlashInfer SM120 sparse MLA backend selection."""
4 +
5 +from types import SimpleNamespace
6 +
7 +import torch
8 +
9 +from vllm.config import set_current_vllm_config
10 +from vllm.platforms.interface import DeviceCapability
11 +from vllm.utils import flashinfer as fi_utils
12 +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
13 + FlashInferMLASparseSM120Backend,
14 +)
15 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
16 +
17 +
18 +def _fake_vllm_config(model_type: str) -> SimpleNamespace:
19 + return SimpleNamespace(
20 + model_config=SimpleNamespace(
21 + hf_text_config=SimpleNamespace(model_type=model_type, index_topk=2048),
22 + ),
23 + )
24 +
25 +
26 +def test_sm120_backend_uses_dedicated_backend_name() -> None:
27 + assert FlashInferMLASparseSM120Backend.get_name() == "FLASHINFER_MLA_SPARSE_SM120"
28 + assert (
29 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120.get_class()
30 + is FlashInferMLASparseSM120Backend
31 + )
32 +
33 +
34 +def test_v32_glm_sm120_backend_accepts_glm_block_size(
35 + monkeypatch,
36 +) -> None:
37 + monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True)
38 +
39 + with set_current_vllm_config(_fake_vllm_config("glm4_moe")):
40 + invalid_reasons = FlashInferMLASparseSM120Backend.validate_configuration(
41 + head_size=576,
42 + dtype=torch.bfloat16,
43 + kv_cache_dtype="fp8",
44 + block_size=256,
45 + use_mla=True,
46 + has_sink=False,
47 + use_sparse=True,
48 + use_mm_prefix=False,
49 + use_per_head_quant_scales=False,
50 + device_capability=DeviceCapability(12, 0),
51 + attn_type="decoder",
52 + )
53 +
54 + assert invalid_reasons == []
tests/v1/attention/test_sparse_mla_backends.py
+9 -6
@@ -36,7 +36,7 @@ if not current_platform.is_cuda():
36
37 from vllm.utils.math_utils import cdiv
38 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
39 - FlashInferMLASparseBackend,
39 + FlashInferMLASparseTRTLLMBackend,
40 )
41 from vllm.v1.attention.backends.mla.flashmla_sparse import (
42 FlashMLASparseBackend,
@@ -174,8 +174,8 @@ def _quantize_dequantize_fp8_ds_mla(
174
175 @pytest.mark.parametrize(
176 "backend_cls",
177 - [FlashMLASparseBackend, FlashInferMLASparseBackend],
178 - ids=["FlashMLA", "FlashInfer"],
177 + [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],
178 + ids=["FlashMLA", "FlashInferTRTLLM"],
179 )
180 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
181 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])
@@ -217,9 +217,12 @@ def test_sparse_backend_decode_correctness(
217 ok, reason = flashmla.is_flashmla_sparse_supported()
218 if not ok:
219 pytest.skip(reason)
220 - elif backend_cls == FlashInferMLASparseBackend:
221 - if not current_platform.has_device_capability(100):
222 - pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher")
220 + elif backend_cls == FlashInferMLASparseTRTLLMBackend:
221 + device_capability = current_platform.get_device_capability()
222 + if device_capability is None or not backend_cls.supports_compute_capability(
223 + device_capability
224 + ):
225 + pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
226
227 batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
228 use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
tests/v1/spec_decode/test_acceptance_length.py
+3 -3
@@ -132,9 +132,9 @@ def get_available_attention_backends() -> list[str]:
132 )
133
134 return [
135 - backend.name
136 - for backend, _ in valid_backends
137 - if backend not in EXCLUDED_BACKENDS
135 + candidate.backend.name
136 + for candidate in valid_backends
137 + if candidate.backend not in EXCLUDED_BACKENDS
138 ]
139
140
tools/pre_commit/generate_attention_backend_docs.py
+5 -2
@@ -690,7 +690,9 @@ def parse_compute_capability(node: ast.ClassDef) -> str:
690 major_list.sort()
691 if len(major_list) == 1:
692 return f"{major_list[0]}.x"
693 - return f"{major_list[0]}.x-{major_list[-1]}.x"
693 + if major_list == list(range(major_list[0], major_list[-1] + 1)):
694 + return f"{major_list[0]}.x-{major_list[-1]}.x"
695 + return ", ".join(f"{major}.x" for major in major_list)
696
697 if min_cap:
698 if max_cap:
@@ -1668,7 +1670,8 @@ def generate_mla_section(
1670 "`--attention-backend=<BACKEND>` (e.g., `FLASHMLA_SPARSE_DSV4`,",
1671 "`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index",
1672 "pipeline (compressor + SWA + indexer, 256-token blocks, head 512);",
1671 - "default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.",
1673 + "default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and",
1674 + "`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.",
1675 "",
1676 ]
1677 )
vllm/model_executor/layers/attention/mla_attention.py
+29 -14
@@ -208,6 +208,7 @@ from vllm.config import (
208 get_current_vllm_config,
209 get_current_vllm_config_or_none,
210 )
211 +from vllm.config.cache import CacheDType
212 from vllm.distributed.parallel_state import (
213 get_dcp_group,
214 is_global_first_rank,
@@ -319,6 +320,22 @@ def _detect_output_quant_key(
320 return kFp8StaticTensorSym
321
322
323 +def _canonicalize_sparse_mla_kv_cache_dtype(
324 + attn_backend: type[AttentionBackend],
325 + kv_cache_dtype: CacheDType,
326 +) -> CacheDType:
327 + backend_name = attn_backend.get_name()
328 + if backend_name == "FLASHMLA_SPARSE" and is_quantized_kv_cache(kv_cache_dtype):
329 + return "fp8_ds_mla"
330 + if backend_name == "FLASHINFER_MLA_SPARSE_SM120" and kv_cache_dtype in (
331 + "auto",
332 + "fp8",
333 + "fp8_e4m3",
334 + ):
335 + return "fp8_ds_mla"
336 + return kv_cache_dtype
337 +
338 +
339 class MLAAttention(nn.Module, AttentionLayerBase):
340 """Multi-Head Latent Attention layer.
341
@@ -369,7 +386,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
386 self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
387
388 if cache_config is not None:
372 - kv_cache_dtype = cache_config.cache_dtype
389 + kv_cache_dtype: CacheDType = cache_config.cache_dtype
390 calculate_kv_scales = cache_config.calculate_kv_scales
391 else:
392 kv_cache_dtype = "auto"
@@ -393,24 +410,22 @@ class MLAAttention(nn.Module, AttentionLayerBase):
410 num_heads=self.num_heads,
411 )
412
396 - # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format
397 - # Automatically convert fp8 kv-cache format to "fp8_ds_mla"
398 - if (
399 - self.attn_backend.get_name() == "FLASHMLA_SPARSE"
400 - and is_quantized_kv_cache(kv_cache_dtype)
401 - and kv_cache_dtype != "fp8_ds_mla"
402 - ):
403 - assert cache_config is not None
404 - cache_config.cache_dtype = "fp8_ds_mla"
405 - kv_cache_dtype = "fp8_ds_mla"
413 + normalized_kv_cache_dtype = _canonicalize_sparse_mla_kv_cache_dtype(
414 + self.attn_backend, kv_cache_dtype
415 + )
416 + if normalized_kv_cache_dtype != kv_cache_dtype:
417 + if cache_config is not None:
418 + cache_config.cache_dtype = normalized_kv_cache_dtype
419 + kv_cache_dtype = normalized_kv_cache_dtype
420 logger.info_once(
407 - "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "
408 - "fp8 kv-cache format, please set `--attention-backend "
409 - "FLASHINFER_MLA_SPARSE`"
421 + "Using %s KV cache format for %s backend.",
422 + kv_cache_dtype,
423 + self.attn_backend.get_name(),
424 )
425
426 if (
427 self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE"
428 + and kv_cache_dtype != "fp8_ds_mla"
429 and is_quantized_kv_cache(kv_cache_dtype)
430 ):
431 logger.info_once(
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+76 -18
@@ -23,24 +23,82 @@ def expert_num_tokens_round_up_and_sum(
23 return torch.sum(ent).item()
24
25
26 -def compute_aligned_M(
26 +def compute_aligned_M_and_alignment(
27 M: int,
28 num_topk: int,
29 local_num_experts: int,
30 alignment: int,
31 expert_tokens_meta: mk.ExpertTokensMetadata | None,
32 -):
32 +) -> tuple[int, int]:
33 + """Return (M_sum, alignment_used).
34 +
35 + `alignment_used` may be smaller than the caller-supplied `alignment` on
36 + SM100/SM120 when DeepGEMM can JIT a smaller BLOCK_M for the per-call
37 + expected_m. Callers that index by block size (e.g. ``M_sum // block_m``)
38 + or assert workspace alignment must use the returned `alignment_used`,
39 + not their original `alignment` argument.
40 +
41 + Prefer this over the int-returning :func:`compute_aligned_M` when the
42 + GEMM call site needs to wrap itself in ``mk_alignment_scope`` or
43 + otherwise reason about the actual per-expert padding.
44 + """
45 if (expert_tokens_meta is not None) and (
46 expert_tokens_meta.expert_num_tokens_cpu is not None
47 ):
36 - return expert_num_tokens_round_up_and_sum(
37 - expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment
48 + return (
49 + expert_num_tokens_round_up_and_sum(
50 + expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment
51 + ),
52 + alignment,
53 + )
54 +
55 + # expert_num_tokens not on cpu. Cap padding by min(M*num_topk,
56 + # local_num_experts) — at batch=1 decode only `num_topk` experts can be
57 + # active, so the worst-case `local_num_experts*(align-1)` is too loose.
58 + # Also shrink `alignment` to DeepGEMM's per-call theoretical BLOCK_M on
59 + # SM100/SM120 when smaller.
60 + expected_m = M * num_topk
61 + try:
62 + from vllm.utils.deep_gemm import (
63 + get_theoretical_mk_alignment_for_contiguous_layout,
64 + )
65 +
66 + # num_groups=local_num_experts so the helper recovers per-expert em;
67 + # omitting it over-picks BLOCK_M on SM120 (heuristic assumes em is
68 + # already per-expert).
69 + per_call_align = get_theoretical_mk_alignment_for_contiguous_layout(
70 + expected_m=expected_m,
71 + num_groups=local_num_experts,
72 )
73 + if per_call_align and per_call_align <= alignment:
74 + alignment = per_call_align
75 + except Exception:
76 + pass
77
40 - # expert_num_tokens information is not available on the cpu.
41 - # compute the max required size.
42 - M_sum = (M * num_topk) + local_num_experts * (alignment - 1)
78 + max_active_experts = min(M * num_topk, local_num_experts)
79 + M_sum = (M * num_topk) + max_active_experts * (alignment - 1)
80 M_sum = round_up(M_sum, alignment)
81 + return M_sum, alignment
82 +
83 +
84 +def compute_aligned_M(
85 + M: int,
86 + num_topk: int,
87 + local_num_experts: int,
88 + alignment: int,
89 + expert_tokens_meta: mk.ExpertTokensMetadata | None,
90 +) -> int:
91 + """Return ``M_sum`` only (backward-compat wrapper).
92 +
93 + Equivalent to :func:`compute_aligned_M_and_alignment`'s first return
94 + value. Existing downstream callers and the warmup path that only size
95 + a workspace use this. Call sites that need the actual per-expert
96 + alignment (to wrap GEMMs in ``mk_alignment_scope``) should use
97 + :func:`compute_aligned_M_and_alignment` instead.
98 + """
99 + M_sum, _ = compute_aligned_M_and_alignment(
100 + M, num_topk, local_num_experts, alignment, expert_tokens_meta
101 + )
102 return M_sum
103
104
@@ -51,12 +109,6 @@ def apply_expert_map(expert_id, expert_map):
109 return expert_id
110
111
54 -@triton.jit
55 -def round_up_128(x: int) -> int:
56 - y = 128
57 - return ((x + y - 1) // y) * y
58 -
59 -
112 @triton.jit
113 def _fwd_kernel_ep_scatter_1(
114 num_recv_tokens_per_expert,
@@ -65,6 +117,7 @@ def _fwd_kernel_ep_scatter_1(
117 num_experts: tl.constexpr,
118 BLOCK_E: tl.constexpr,
119 BLOCK_EXPERT_NUM: tl.constexpr,
120 + ALIGN_M: tl.constexpr,
121 ):
122 cur_expert = tl.program_id(0)
123
@@ -74,7 +127,8 @@ def _fwd_kernel_ep_scatter_1(
127 mask=offset_cumsum < num_experts,
128 other=0,
129 )
77 - tokens_per_expert = round_up_128(tokens_per_expert)
130 + # Round up to ALIGN_M so cumsum matches the workspace's per-expert slices.
131 + tokens_per_expert = ((tokens_per_expert + ALIGN_M - 1) // ALIGN_M) * ALIGN_M
132 cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert
133
134 # Extract this block's offset from the register vector (warp shuffle,
@@ -227,10 +281,12 @@ def ep_scatter(
281 output_tensor_scale: torch.Tensor,
282 m_indices: torch.Tensor,
283 output_index: torch.Tensor,
284 + align_m: int = 128,
285 block_size: int = 128,
286 pack_ue8m0: bool = False,
287 ):
233 - BLOCK_E = 128 # token num of per expert is aligned to 128
288 + # BLOCK_E is the m_indices fill-loop tile (masked), independent of align_m.
289 + BLOCK_E = 128
290 BLOCK_D = block_size # block size of activation-scale quantization
291 num_warps = 8
292 num_experts = num_recv_tokens_per_expert.shape[0]
@@ -238,7 +294,7 @@ def ep_scatter(
294 # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts)
295 grid = num_experts
296
241 - assert m_indices.shape[0] % BLOCK_E == 0
297 + assert m_indices.shape[0] % align_m == 0
298 assert expert_start_loc.shape[0] == num_experts
299
300 # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is.
@@ -253,6 +309,7 @@ def ep_scatter(
309 num_warps=num_warps,
310 BLOCK_E=BLOCK_E,
311 BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
312 + ALIGN_M=align_m,
313 )
314
315 grid = min(recv_topk.shape[0], 1024 * 8)
@@ -418,7 +475,7 @@ def deepgemm_moe_permute(
475 if block_size is not None:
476 block_k = block_size
477
421 - M_sum = compute_aligned_M(
478 + M_sum, align_used = compute_aligned_M_and_alignment(
479 M=topk_ids.size(0),
480 num_topk=topk_ids.size(1),
481 local_num_experts=local_num_experts,
@@ -482,11 +539,12 @@ def deepgemm_moe_permute(
539 output_tensor_scale=aq_scale_out,
540 m_indices=expert_ids,
541 output_index=inv_perm,
542 + align_m=align_used,
543 block_size=block_k,
544 pack_ue8m0=pack_ue8m0,
545 )
546
489 - return aq_out, aq_scale_out, expert_ids, inv_perm
547 + return aq_out, aq_scale_out, expert_ids, inv_perm, align_used
548
549
550 def deepgemm_unpermute_and_reduce(
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+5 -4
@@ -318,11 +318,12 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):
318
319 def supports_packed_ue8m0_act_scales(self) -> bool:
320 """
321 - DeepGemm supports packed ue8m0 activation scales format in devices == sm100
321 + DeepGemm supports packed ue8m0 activation scales on Blackwell-family
322 + GPUs (SM100 datacenter and SM120 consumer).
323 """
323 - return (
324 - is_deep_gemm_e8m0_used()
325 - and current_platform.is_device_capability_family(100)
324 + return is_deep_gemm_e8m0_used() and (
325 + current_platform.is_device_capability_family(100)
326 + or current_platform.is_device_capability_family(120)
327 )
328
329 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py
+74 -57
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import (
12 FusedMoEQuantConfig,
13 )
14 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
15 - compute_aligned_M,
15 + compute_aligned_M_and_alignment,
16 deepgemm_moe_permute,
17 deepgemm_unpermute_and_reduce,
18 )
@@ -43,6 +43,7 @@ from vllm.utils.deep_gemm import (
43 is_deep_gemm_supported,
44 m_grouped_fp8_fp4_gemm_nt_contiguous,
45 m_grouped_fp8_gemm_nt_contiguous,
46 + mk_alignment_scope,
47 )
48 from vllm.utils.import_utils import has_deep_gemm
49
@@ -210,10 +211,10 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
211 # Use the contiguous-layout M alignment (matches apply()); block_shape[0]
212 # is the quant block (1 for MXFP8) and would under-size the workspace.
213 block_m = get_mk_alignment_for_contiguous_layout()[0]
213 - M_sum = compute_aligned_M(
214 + M_sum, align_used = compute_aligned_M_and_alignment(
215 M, topk, local_num_experts, block_m, expert_tokens_meta
216 )
216 - assert M_sum % block_m == 0
217 + assert M_sum % align_used == 0
218
219 activation_out_dim = self.adjust_N_for_activation(N, activation)
220 workspace1 = (M_sum, max(activation_out_dim, K))
@@ -316,7 +317,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
317
318 assert w2.size(1) == K
319
319 - M_sum = compute_aligned_M(
320 + M_sum, _ = compute_aligned_M_and_alignment(
321 M=topk_ids.size(0),
322 num_topk=topk_ids.size(1),
323 local_num_experts=local_num_experts,
@@ -327,7 +328,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
328 a1q_perm = _resize_cache(
329 workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)
330 )
330 - a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
331 + a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
332 aq=a1q,
333 aq_scale=a1q_scale,
334 topk_ids=topk_ids,
@@ -349,23 +350,35 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
350 else {}
351 )
352
352 - mm1_out = _resize_cache(workspace2, (M_sum, N))
353 - m_grouped_fp8_gemm_nt_contiguous(
354 - (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs
355 - )
353 + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment;
354 + # otherwise the scheduler can pick the wrong expert id from m_indices
355 + # under cudagraph replay.
356 + with mk_alignment_scope(align_used):
357 + mm1_out = _resize_cache(workspace2, (M_sum, N))
358 + m_grouped_fp8_gemm_nt_contiguous(
359 + (a1q, a1q_scale),
360 + (w1, self.w1_scale),
361 + mm1_out,
362 + expert_ids,
363 + **gemm_kwargs,
364 + )
365
357 - activation_out_dim = self.adjust_N_for_activation(N, activation)
358 - quant_out = _resize_cache(
359 - workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
360 - )
361 - a2q, a2q_scale = self._act_mul_quant(
362 - input=mm1_out.view(-1, N), output=quant_out, activation=activation
363 - )
366 + activation_out_dim = self.adjust_N_for_activation(N, activation)
367 + quant_out = _resize_cache(
368 + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
369 + )
370 + a2q, a2q_scale = self._act_mul_quant(
371 + input=mm1_out.view(-1, N), output=quant_out, activation=activation
372 + )
373
365 - mm2_out = _resize_cache(workspace2, (M_sum, K))
366 - m_grouped_fp8_gemm_nt_contiguous(
367 - (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, **gemm_kwargs
368 - )
374 + mm2_out = _resize_cache(workspace2, (M_sum, K))
375 + m_grouped_fp8_gemm_nt_contiguous(
376 + (a2q, a2q_scale),
377 + (w2, self.w2_scale),
378 + mm2_out,
379 + expert_ids,
380 + **gemm_kwargs,
381 + )
382
383 if apply_router_weight_on_input:
384 topk_weights = torch.ones_like(topk_weights)
@@ -384,7 +397,8 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
397 """DeepGemm-based fused MoE expert implementation for FP4 weights.
398
399 Uses m_grouped_fp8_fp4_gemm_nt_contiguous with FP8 activations and
387 - MXFP4 (FP4 E2M1 packed as uint8) weights. Requires SM100+ (Blackwell).
400 + MXFP4 (FP4 E2M1 packed as uint8) weights. Requires Blackwell-family
401 + GPUs (SM100 datacenter or SM120 consumer).
402 """
403
404 # FP8 activation block size (hardcoded since mxfp4_w4a8 quant config
@@ -409,9 +423,9 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
423 def _supports_current_device() -> bool:
424 from vllm.platforms import current_platform
425
412 - return (
413 - is_deep_gemm_supported()
414 - and current_platform.is_device_capability_family(100)
426 + return is_deep_gemm_supported() and (
427 + current_platform.is_device_capability_family(100)
428 + or current_platform.is_device_capability_family(120)
429 )
430
431 @staticmethod
@@ -454,10 +468,10 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
468 activation: MoEActivation,
469 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
470 block_m = get_mk_alignment_for_contiguous_layout()[0]
457 - M_sum = compute_aligned_M(
471 + M_sum, align_used = compute_aligned_M_and_alignment(
472 M, topk, local_num_experts, block_m, expert_tokens_meta
473 )
460 - assert M_sum % block_m == 0
474 + assert M_sum % align_used == 0
475
476 activation_out_dim = self.adjust_N_for_activation(N, activation)
477 workspace1 = (M_sum, max(activation_out_dim, K))
@@ -533,7 +547,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
547 if global_num_experts == -1:
548 global_num_experts = local_num_experts
549
536 - M_sum = compute_aligned_M(
550 + M_sum, _ = compute_aligned_M_and_alignment(
551 M=topk_ids.size(0),
552 num_topk=topk_ids.size(1),
553 local_num_experts=local_num_experts,
@@ -544,7 +558,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
558 a1q_perm = _resize_cache(
559 workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)
560 )
547 - a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
561 + a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
562 aq=a1q,
563 aq_scale=a1q_scale,
564 topk_ids=topk_ids,
@@ -555,37 +569,40 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):
569 )
570 assert a1q.size(0) == M_sum
571
558 - # FC1: FP8 activations x FP4 weights
559 - # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4).
560 - mm1_out = _resize_cache(workspace2, (M_sum, N))
561 - m_grouped_fp8_fp4_gemm_nt_contiguous(
562 - (a1q, a1q_scale),
563 - (w1.view(torch.int8), self.w1_scale),
564 - mm1_out,
565 - expert_ids,
566 - recipe_a=(1, self._ACT_BLOCK_K),
567 - recipe_b=(1, self._WEIGHT_BLOCK_K),
568 - )
572 + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment;
573 + # see DeepGemmExperts.apply for rationale.
574 + with mk_alignment_scope(align_used):
575 + # FC1: FP8 activations x FP4 weights
576 + # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4).
577 + mm1_out = _resize_cache(workspace2, (M_sum, N))
578 + m_grouped_fp8_fp4_gemm_nt_contiguous(
579 + (a1q, a1q_scale),
580 + (w1.view(torch.int8), self.w1_scale),
581 + mm1_out,
582 + expert_ids,
583 + recipe_a=(1, self._ACT_BLOCK_K),
584 + recipe_b=(1, self._WEIGHT_BLOCK_K),
585 + )
586
570 - # SwiGLU activation + FP8 requant
571 - activation_out_dim = self.adjust_N_for_activation(N, activation)
572 - quant_out = _resize_cache(
573 - workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
574 - )
575 - a2q, a2q_scale = self._act_mul_quant(
576 - input=mm1_out.view(-1, N), output=quant_out, activation=activation
577 - )
587 + # SwiGLU activation + FP8 requant
588 + activation_out_dim = self.adjust_N_for_activation(N, activation)
589 + quant_out = _resize_cache(
590 + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim)
591 + )
592 + a2q, a2q_scale = self._act_mul_quant(
593 + input=mm1_out.view(-1, N), output=quant_out, activation=activation
594 + )
595
579 - # FC2: FP8 activations x FP4 weights
580 - mm2_out = _resize_cache(workspace2, (M_sum, K))
581 - m_grouped_fp8_fp4_gemm_nt_contiguous(
582 - (a2q, a2q_scale),
583 - (w2.view(torch.int8), self.w2_scale),
584 - mm2_out,
585 - expert_ids,
586 - recipe_a=(1, self._ACT_BLOCK_K),
587 - recipe_b=(1, self._WEIGHT_BLOCK_K),
588 - )
596 + # FC2: FP8 activations x FP4 weights
597 + mm2_out = _resize_cache(workspace2, (M_sum, K))
598 + m_grouped_fp8_fp4_gemm_nt_contiguous(
599 + (a2q, a2q_scale),
600 + (w2.view(torch.int8), self.w2_scale),
601 + mm2_out,
602 + expert_ids,
603 + recipe_a=(1, self._ACT_BLOCK_K),
604 + recipe_b=(1, self._WEIGHT_BLOCK_K),
605 + )
606
607 if apply_router_weight_on_input:
608 topk_weights = torch.ones_like(topk_weights)
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+48 -10
@@ -57,6 +57,40 @@ if has_triton_kernels():
57 )
58
59
60 +def _pack_deepgemm_mxfp4_scales(
61 + w13_weight: torch.Tensor,
62 + w2_weight: torch.Tensor,
63 + w13_weight_scale: torch.Tensor,
64 + w2_weight_scale: torch.Tensor,
65 +) -> tuple[torch.Tensor, torch.Tensor]:
66 + from vllm.model_executor.layers.quantization.utils.fp8_utils import (
67 + deepgemm_post_process_weight_scale_block,
68 + )
69 +
70 + num_experts = w13_weight.shape[0]
71 + intermediate_size_2 = w13_weight.shape[1] # = intermediate*2
72 + hidden_size = w13_weight.shape[2] * 2 # weight is FP4-packed
73 + intermediate_size = w2_weight.shape[2] * 2 # weight is FP4-packed
74 + block_shape = (1, 32) # MXFP4 block (per-row, K=32)
75 +
76 + return (
77 + deepgemm_post_process_weight_scale_block(
78 + ws=w13_weight_scale.data,
79 + mn=intermediate_size_2,
80 + k=hidden_size,
81 + quant_block_shape=block_shape,
82 + num_groups=num_experts,
83 + ),
84 + deepgemm_post_process_weight_scale_block(
85 + ws=w2_weight_scale.data,
86 + mn=hidden_size,
87 + k=intermediate_size,
88 + quant_block_shape=block_shape,
89 + num_groups=num_experts,
90 + ),
91 + )
92 +
93 +
94 class Mxfp4MoeBackend(Enum):
95 NONE = "None"
96 # DeepGEMM FP8xFP4 backend (SM100+)
@@ -652,15 +686,18 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
686 """Convert loaded weights into backend-specific kernel format."""
687
688 if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
655 - from vllm.model_executor.layers.quantization.utils.fp8_utils import (
656 - _upcast_e8m0_to_fp32,
689 + w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
690 + w13_weight,
691 + w2_weight,
692 + w13_weight_scale,
693 + w2_weight_scale,
694 )
695
696 return (
697 w13_weight.data,
698 w2_weight.data,
662 - _upcast_e8m0_to_fp32(w13_weight_scale.data),
663 - _upcast_e8m0_to_fp32(w2_weight_scale.data),
699 + w13_weight_scale,
700 + w2_weight_scale,
701 w13_bias,
702 w2_bias,
703 )
@@ -1195,17 +1232,18 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1232 """
1233
1234 if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
1198 - from vllm.model_executor.layers.quantization.utils.fp8_utils import (
1199 - _upcast_e8m0_to_fp32,
1235 + w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
1236 + w13_weight,
1237 + w2_weight,
1238 + w13_weight_scale,
1239 + w2_weight_scale,
1240 )
1241
1202 - # Weights stay as uint8 packed FP4 — no layout change needed.
1203 - # Convert E8M0 uint8 scales to float32.
1242 return (
1243 w13_weight.data,
1244 w2_weight.data,
1207 - _upcast_e8m0_to_fp32(w13_weight_scale.data),
1208 - _upcast_e8m0_to_fp32(w2_weight_scale.data),
1245 + w13_weight_scale,
1246 + w2_weight_scale,
1247 w13_bias,
1248 w2_bias,
1249 )
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+36 -22
@@ -1058,6 +1058,34 @@ def _upcast_e8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor:
1058 return fp32_bits.view(torch.float32)
1059
1060
1061 +def deepgemm_post_process_weight_scale_block(
1062 + ws: torch.Tensor,
1063 + mn: int,
1064 + k: int,
1065 + quant_block_shape: tuple[int, ...],
1066 + num_groups: int,
1067 + is_sfa: bool = False,
1068 +) -> torch.Tensor:
1069 + if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
1070 + # Scales already in E8M0 from checkpoint; upcast to fp32 and let
1071 + # DeepGEMM pack the layout expected by the target architecture.
1072 + ws = _upcast_e8m0_to_fp32(ws)
1073 + else:
1074 + assert ws.dtype == torch.float32, (
1075 + f"Expected tensor scales dtype to be torch.float32 or "
1076 + f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
1077 + )
1078 +
1079 + return transform_sf_into_required_layout(
1080 + sf=ws,
1081 + mn=mn,
1082 + k=k,
1083 + recipe=(1, quant_block_shape[0], quant_block_shape[1]),
1084 + num_groups=num_groups,
1085 + is_sfa=is_sfa,
1086 + )
1087 +
1088 +
1089 def deepgemm_post_process_fp8_weight_block(
1090 wq: torch.Tensor,
1091 ws: torch.Tensor,
@@ -1073,13 +1101,13 @@ def deepgemm_post_process_fp8_weight_block(
1101
1102 if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
1103 # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0
1076 - # bits as uint8 for MXFP8) — upcast to fp32 and skip requantization
1104 + # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization
1105 # (weights already have power-of-two scales).
1106 ws = _upcast_e8m0_to_fp32(ws)
1107 else:
1108 assert ws.dtype == torch.float32, (
1109 f"Expected tensor scales dtype to be torch.float32 or "
1082 - f"torch.float8_e8m0fnu, got {ws.dtype} instead"
1110 + f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
1111 )
1112 if use_e8m0:
1113 requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)
@@ -1094,16 +1122,12 @@ def deepgemm_post_process_fp8_weight_block(
1122 r = wq.size(0) // g
1123 wq = wq.view(g, r, d)
1124 ws = ws.view(g, r // quant_block_shape[0], d // quant_block_shape[1])
1097 - # Pre-transform scale with recipe=(1, 128, 128) to broadcast + pack
1098 - # into TMA-aligned UE8M0 (INT32) layout. At runtime fp8_einsum uses
1099 - # recipe=(1, 1, 128) which sees INT dtype and skips re-transform.
1100 - dg_ws = transform_sf_into_required_layout(
1101 - sf=ws,
1125 + dg_ws = deepgemm_post_process_weight_scale_block(
1126 + ws=ws,
1127 mn=r,
1128 k=d,
1104 - recipe=(1, quant_block_shape[0], quant_block_shape[1]),
1129 + quant_block_shape=quant_block_shape,
1130 num_groups=g,
1106 - is_sfa=False,
1131 )
1132 return wq, dg_ws
1133
@@ -1113,22 +1137,12 @@ def deepgemm_post_process_fp8_weight_block(
1137 wq = wq.unsqueeze(0)
1138 ws = ws.unsqueeze(0)
1139
1116 - # From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46
1117 - # (1, block_n, block_k): (1, 128, 128) for FP8 block, (1, 1, 32) for MXFP8.
1118 - recipe = (1, quant_block_shape[0], quant_block_shape[1])
1119 -
1120 - # Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp
1121 - # DeepGemm uses the `transform_sf_into_required_layout` function to
1122 - # represent scales in the correct format.
1123 - dg_ws = transform_sf_into_required_layout(
1124 - sf=ws,
1140 + dg_ws = deepgemm_post_process_weight_scale_block(
1141 + ws=ws,
1142 mn=wq.size(1),
1143 k=wq.size(2),
1127 - recipe=recipe,
1144 + quant_block_shape=quant_block_shape,
1145 num_groups=wq.size(0),
1129 - # is the scale factors for A in (Refers to the argument A in A @ B).
1130 - # Weights are B.
1131 - is_sfa=False,
1146 )
1147
1148 if original_ndim == 2:
vllm/model_executor/warmup/deep_gemm_warmup.py
+58 -25
@@ -12,7 +12,9 @@ from tqdm import tqdm
12 import vllm.envs as envs
13 from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank
14 from vllm.model_executor.layers.fused_moe import MoERunner
15 -from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M
15 +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
16 + compute_aligned_M_and_alignment,
17 +)
18 from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts
19 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
20 TritonOrDeepGemmExperts,
@@ -25,6 +27,7 @@ from vllm.utils.deep_gemm import (
27 fp8_gemm_nt,
28 get_mk_alignment_for_contiguous_layout,
29 m_grouped_fp8_gemm_nt_contiguous,
30 + mk_alignment_scope,
31 )
32 from vllm.utils.math_utils import cdiv
33 from vllm.utils.platform_utils import num_compute_units
@@ -238,7 +241,7 @@ def _get_grouped_gemm_params(
241 w2: torch.Tensor,
242 num_topk: int,
243 max_tokens: int,
241 -) -> tuple[int, int, torch.Tensor]:
244 +) -> tuple[int, int, list[tuple[int, int, torch.Tensor]]]:
245 assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts"
246
247 block_m = get_mk_alignment_for_contiguous_layout()[0]
@@ -248,19 +251,46 @@ def _get_grouped_gemm_params(
251 # Assumes all ranks have the same max_num_batched_tokens
252 max_tokens = get_dp_group().world_size * max_tokens
253
251 - # This is the maximum GroupedGemm M size that we expect to run
252 - # the grouped_gemm with.
253 - MAX_M = compute_aligned_M(
254 - max_tokens, num_topk, num_experts, block_m, expert_tokens_meta=None
255 - )
256 - # Distribute expert-ids evenly.
257 - MAX_BLOCKS = MAX_M // block_m
258 - expert_ids_block = torch.randint(
259 - low=0, high=num_experts, size=(MAX_BLOCKS,), device=device, dtype=torch.int32
254 + request_m_values = _generate_optimal_warmup_m_values(
255 + max_tokens,
256 + max(w1.size(1), w2.size(1)),
257 + device,
258 )
261 - expert_ids = torch.repeat_interleave(expert_ids_block, block_m, dim=0)
259 + request_m_values = sorted({m for m in (*request_m_values, max_tokens) if m > 0})
260 + if not request_m_values:
261 + return 0, block_m, []
262 +
263 + cases_by_shape: dict[tuple[int, int], torch.Tensor] = {}
264 + for request_m in request_m_values:
265 + M_sum, align_used = compute_aligned_M_and_alignment(
266 + M=request_m,
267 + num_topk=num_topk,
268 + local_num_experts=num_experts,
269 + alignment=block_m,
270 + expert_tokens_meta=None,
271 + )
272 + if (M_sum, align_used) in cases_by_shape:
273 + continue
274 +
275 + num_blocks = M_sum // align_used
276 + expert_ids_block = torch.randint(
277 + low=0,
278 + high=num_experts,
279 + size=(num_blocks,),
280 + device=device,
281 + dtype=torch.int32,
282 + )
283 + cases_by_shape[(M_sum, align_used)] = torch.repeat_interleave(
284 + expert_ids_block, align_used, dim=0
285 + )
286
263 - return MAX_M, block_m, expert_ids
287 + max_m = max(M_sum for M_sum, _ in cases_by_shape)
288 + warmup_cases = [
289 + (M_sum, align_used, expert_ids)
290 + for (M_sum, align_used), expert_ids in sorted(cases_by_shape.items())
291 + ]
292 +
293 + return max_m, block_m, warmup_cases
294
295
296 def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
@@ -278,7 +308,11 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
308 ):
309 return
310
281 - MAX_M, block_m, expert_ids = _get_grouped_gemm_params(w1, w2, num_topk, max_tokens)
311 + MAX_M, block_m, warmup_cases = _get_grouped_gemm_params(
312 + w1, w2, num_topk, max_tokens
313 + )
314 + if not warmup_cases:
315 + return
316 device = w1.device
317
318 def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
@@ -289,15 +323,14 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
323 )
324 out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
325
292 - m_values = list(range(block_m, MAX_M + 1, block_m))
293 -
294 - for num_tokens in m_values:
295 - m_grouped_fp8_gemm_nt_contiguous(
296 - (a1q[:num_tokens], a1q_scales[:num_tokens]),
297 - (w, w_scale),
298 - out[:num_tokens],
299 - expert_ids[:num_tokens],
300 - )
326 + for num_tokens, align_used, expert_ids in warmup_cases:
327 + with mk_alignment_scope(align_used):
328 + m_grouped_fp8_gemm_nt_contiguous(
329 + (a1q[:num_tokens], a1q_scales[:num_tokens]),
330 + (w, w_scale),
331 + out[:num_tokens],
332 + expert_ids,
333 + )
334 if pbar is not None:
335 pbar.update(1)
336
@@ -350,8 +383,8 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
383 w13, _, w2, _, num_topk = _extract_data_from_fused_moe_module(m)
384 if w13.size() in seen_grouped_sizes and w2.size() in seen_grouped_sizes:
385 continue
353 - MAX_M, block_m, _ = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
354 - n_values = (MAX_M - block_m) // block_m + 1
386 + _, _, warmup_cases = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens)
387 + n_values = len(warmup_cases)
388 if w13.size() not in seen_grouped_sizes:
389 total += n_values
390 seen_grouped_sizes.add(w13.size())
vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py new
+226
@@ -0,0 +1,226 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Warm up DeepSeek V4 mHC TileLang kernels before serving requests.
4 +
5 +Ported from lucifer1004/vllm-jasl with the two env-var knobs removed
6 +(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`).
7 +Gating is intrinsic: non-DSv4 models and layers without hc_* attributes
8 +return early, so the warmup is a no-op except where it's needed.
9 +"""
10 +
11 +import time
12 +from collections.abc import Iterable
13 +
14 +import torch
15 +
16 +from vllm.logger import init_logger
17 +from vllm.tracing import instrument
18 +from vllm.utils.math_utils import cdiv
19 +
20 +logger = init_logger(__name__)
21 +
22 +_AUTO_WARMUP_MAX_TOKENS = 16_384
23 +_DEFAULT_TOKEN_SIZE_CANDIDATES = (
24 + 1,
25 + 2,
26 + 4,
27 + 8,
28 + 16,
29 + 32,
30 + 64,
31 + 128,
32 + 256,
33 + 512,
34 + 1024,
35 + 2048,
36 + 4096,
37 + 8192,
38 + 16_384,
39 +)
40 +
41 +
42 +def _compute_mhc_pre_num_split(
43 + *,
44 + num_tokens: int,
45 + hidden_size: int,
46 + hc_mult: int,
47 + num_sms: int,
48 +) -> int:
49 + block_k = 64
50 + block_m = 64
51 + k = hc_mult * hidden_size
52 + grid_size = cdiv(num_tokens, block_m)
53 + split_k = num_sms // grid_size
54 + num_block_k = cdiv(k, block_k)
55 + split_k = min(split_k, num_block_k // 4)
56 + return max(split_k, 1)
57 +
58 +
59 +def _normalize_token_sizes(
60 + token_sizes: Iterable[int],
61 + *,
62 + max_tokens: int,
63 +) -> list[int]:
64 + return sorted({size for size in token_sizes if 1 <= size <= max_tokens})
65 +
66 +
67 +def _select_mhc_warmup_token_sizes(
68 + *,
69 + max_tokens: int,
70 + cudagraph_capture_sizes: list[int],
71 +) -> list[int]:
72 + if max_tokens <= 0:
73 + return []
74 +
75 + max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS)
76 + candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES)
77 + candidates.extend(cudagraph_capture_sizes)
78 + candidates.append(max_auto_tokens)
79 + return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens)
80 +
81 +
82 +def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None:
83 + for module in model.modules():
84 + if module.__class__.__name__ != "DeepseekV4DecoderLayer":
85 + continue
86 + if all(
87 + hasattr(module, attr)
88 + for attr in (
89 + "hc_pre",
90 + "hc_post",
91 + "hc_attn_fn",
92 + "hc_attn_scale",
93 + "hc_attn_base",
94 + "hc_ffn_fn",
95 + "hc_ffn_scale",
96 + "hc_ffn_base",
97 + )
98 + ):
99 + return module
100 + return None
101 +
102 +
103 +def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None:
104 + for module in model.modules():
105 + if module.__class__.__name__ != "DeepseekV4Model":
106 + continue
107 + if all(
108 + hasattr(module, attr)
109 + for attr in ("hc_head_fn", "hc_head_scale", "hc_head_base")
110 + ):
111 + return module
112 + return None
113 +
114 +
115 +def _warmup_layer_mhc(
116 + layer: torch.nn.Module,
117 + token_sizes: list[int],
118 +) -> None:
119 + max_tokens = max(token_sizes)
120 + hidden_size = int(layer.hidden_size)
121 + hc_mult = int(layer.hc_mult)
122 + device = layer.hc_attn_fn.device
123 + residual = torch.zeros(
124 + max_tokens,
125 + hc_mult,
126 + hidden_size,
127 + dtype=torch.bfloat16,
128 + device=device,
129 + )
130 +
131 + for size in token_sizes:
132 + residual_slice = residual[:size]
133 + for fn, scale, base in (
134 + (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base),
135 + (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base),
136 + ):
137 + layer_input, post_mix, comb_mix = layer.hc_pre(
138 + residual_slice,
139 + fn,
140 + scale,
141 + base,
142 + )
143 + layer.hc_post(layer_input, residual_slice, post_mix, comb_mix)
144 +
145 +
146 +def _warmup_hc_head(
147 + model: torch.nn.Module,
148 + token_sizes: list[int],
149 +) -> None:
150 + # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored
151 + # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp
152 + # instance attached to the model as ``hc_head_op``. We call through
153 + # that instance so the warmup exercises the same dispatched
154 + # implementation as the inference path.
155 + hc_head_op = getattr(model, "hc_head_op", None)
156 + if hc_head_op is None:
157 + return
158 +
159 + max_tokens = max(token_sizes)
160 + hidden_size = int(model.config.hidden_size)
161 + hc_mult = int(model.hc_mult)
162 + device = model.hc_head_fn.device
163 + hidden_states = torch.zeros(
164 + max_tokens,
165 + hc_mult,
166 + hidden_size,
167 + dtype=torch.bfloat16,
168 + device=device,
169 + )
170 +
171 + for size in token_sizes:
172 + hc_head_op(
173 + hidden_states[:size],
174 + model.hc_head_fn,
175 + model.hc_head_scale,
176 + model.hc_head_base,
177 + model.rms_norm_eps,
178 + model.hc_eps,
179 + )
180 +
181 +
182 +@instrument(span_name="DeepSeek V4 mHC warmup")
183 +def deepseek_v4_mhc_warmup(
184 + model: torch.nn.Module,
185 + *,
186 + max_tokens: int,
187 + cudagraph_capture_sizes: list[int] | None = None,
188 +) -> None:
189 + # Cheap model-type gate before walking ``model.modules()``. The class
190 + # walk below is O(num_layers) and shows up in startup time on very
191 + # large checkpoints; bail out for any model that is not DeepSeek V4.
192 + config = getattr(model, "config", None)
193 + model_type = getattr(config, "model_type", None) if config is not None else None
194 + if model_type is not None and model_type != "deepseek_v4":
195 + return
196 +
197 + layer = _find_first_mhc_layer(model)
198 + if layer is None:
199 + return
200 +
201 + device = layer.hc_attn_fn.device
202 + if device.type != "cuda":
203 + return
204 +
205 + deepseek_model = _find_deepseek_v4_model(model)
206 + token_sizes = _select_mhc_warmup_token_sizes(
207 + max_tokens=max_tokens,
208 + cudagraph_capture_sizes=cudagraph_capture_sizes or [],
209 + )
210 + if not token_sizes:
211 + return
212 +
213 + started = time.perf_counter()
214 + logger.info(
215 + "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s",
216 + token_sizes,
217 + )
218 + with torch.inference_mode():
219 + _warmup_layer_mhc(layer, token_sizes)
220 + if deepseek_model is not None:
221 + _warmup_hc_head(deepseek_model, token_sizes)
222 + torch.accelerator.synchronize()
223 + logger.info(
224 + "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.",
225 + time.perf_counter() - started,
226 + )
vllm/model_executor/warmup/flashinfer_autotune_cache.py new
+56
@@ -0,0 +1,56 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""FlashInfer autotune cache helpers."""
4 +
5 +import hashlib
6 +import os
7 +import tempfile
8 +from contextlib import suppress
9 +from pathlib import Path
10 +from typing import TYPE_CHECKING
11 +
12 +import vllm.envs as envs
13 +from vllm.compilation.caching import aot_compile_hash_factors
14 +
15 +if TYPE_CHECKING:
16 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
17 +
18 +
19 +def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str:
20 + factors = aot_compile_hash_factors(runner.vllm_config)
21 + return hashlib.sha256(str(factors).encode()).hexdigest()
22 +
23 +
24 +def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
25 + override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR
26 + if override_dir:
27 + root = Path(override_dir).expanduser()
28 + else:
29 + from flashinfer.jit import env as flashinfer_jit_env
30 +
31 + flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR
32 + root = (
33 + Path(envs.VLLM_CACHE_ROOT)
34 + / "flashinfer_autotune_cache"
35 + / flashinfer_workspace.parent.name
36 + / flashinfer_workspace.name
37 + )
38 +
39 + output_dir = root / flashinfer_autotune_cache_hash(runner)
40 + output_dir.mkdir(parents=True, exist_ok=True)
41 + return output_dir / "autotune_configs.json"
42 +
43 +
44 +def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None:
45 + cache_path.parent.mkdir(parents=True, exist_ok=True)
46 + fd, tmp_path = tempfile.mkstemp(
47 + dir=cache_path.parent, suffix=".tmp", prefix=f".{cache_path.name}."
48 + )
49 + try:
50 + with os.fdopen(fd, "wb") as f:
51 + f.write(contents)
52 + os.replace(tmp_path, cache_path)
53 + except BaseException:
54 + with suppress(OSError):
55 + os.unlink(tmp_path)
56 + raise
vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py new
+255
@@ -0,0 +1,255 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Warmup and autotune helpers for FlashInfer sparse MLA backends."""
4 +
5 +from typing import TYPE_CHECKING, cast
6 +
7 +import torch
8 +
9 +from vllm.logger import init_logger
10 +from vllm.model_executor.warmup.flashinfer_autotune_cache import (
11 + resolve_flashinfer_autotune_file,
12 + write_flashinfer_autotune_cache,
13 +)
14 +from vllm.platforms import current_platform
15 +from vllm.utils.flashinfer import autotune as flashinfer_autotune
16 +from vllm.utils.flashinfer import has_flashinfer
17 +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup
18 +
19 +if TYPE_CHECKING:
20 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2GPUModelRunner
21 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
22 + from vllm.v1.worker.gpu_worker import Worker
23 +
24 +logger = init_logger(__name__)
25 +
26 +_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
27 + {
28 + "FLASHMLA_SPARSE_DSV4",
29 + "FLASHINFER_MLA_SPARSE_DSV4",
30 + "ROCM_FLASHMLA_SPARSE_DSV4",
31 + "DEEPSEEK_SPARSE_SWA",
32 + }
33 +)
34 +_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"})
35 +_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"})
36 +
37 +_FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = {
38 + "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2",
39 + "FLASHINFER_MLA_SPARSE_DSV4": "DSv4",
40 +}
41 +
42 +_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16
43 +
44 +
45 +def _attention_backend_name(backend: object) -> str | None:
46 + get_name = getattr(backend, "get_name", None)
47 + if get_name is None:
48 + return None
49 + try:
50 + return get_name()
51 + except NotImplementedError:
52 + return None
53 +
54 +
55 +def _has_deepseek_v4_sparse_mla_backend(runner: "GPUModelRunner") -> bool:
56 + for groups in getattr(runner, "attn_groups", []) or ():
57 + for group in groups:
58 + name = _attention_backend_name(getattr(group, "backend", None))
59 + if name in _DEEPSEEK_V4_SPARSE_MLA_BACKENDS:
60 + return True
61 + return False
62 +
63 +
64 +def _flashinfer_sparse_mla_decode_label(
65 + runner: "GPUModelRunner",
66 + allowed_backends: frozenset[str],
67 +) -> str | None:
68 + for groups in getattr(runner, "attn_groups", []) or ():
69 + for group in groups:
70 + name = _attention_backend_name(getattr(group, "backend", None))
71 + if name in allowed_backends:
72 + return _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS.get(name)
73 + return None
74 +
75 +
76 +def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int:
77 + return max(0, min(num_tokens, max_tokens))
78 +
79 +
80 +def _uses_v2_model_runner(runner: "GPUModelRunner") -> bool:
81 + vllm_config = getattr(runner, "vllm_config", None)
82 + return bool(getattr(vllm_config, "use_v2_model_runner", False))
83 +
84 +
85 +def _run_flashinfer_sparse_mla_decode_autotune(
86 + worker: "Worker",
87 + num_tokens: int,
88 + allowed_backends: frozenset[str],
89 +) -> bool:
90 + """Autotune FlashInfer's SM120 sparse-MLA decode path."""
91 + runner = worker.model_runner
92 + log_label = _flashinfer_sparse_mla_decode_label(runner, allowed_backends)
93 + if log_label is None:
94 + return False
95 + if worker.vllm_config.kernel_config.enable_flashinfer_autotune is not True:
96 + return False
97 + if not has_flashinfer() or not current_platform.is_device_capability_family(120):
98 + return False
99 +
100 + try:
101 + from flashinfer.autotuner import AutoTuner
102 + except ImportError:
103 + logger.warning(
104 + "Skipping FlashInfer SM120 sparse MLA decode autotune because "
105 + "FlashInfer autotuner is unavailable."
106 + )
107 + return False
108 +
109 + from vllm.distributed.parallel_state import get_world_group
110 +
111 + world = get_world_group()
112 + is_leader = world.rank_in_group == 0
113 + cache_path = resolve_flashinfer_autotune_file(runner)
114 +
115 + dummy_run_kwargs = dict(
116 + num_tokens=num_tokens,
117 + skip_eplb=True,
118 + is_profile=True,
119 + force_attention=True,
120 + create_mixed_batch=True,
121 + )
122 +
123 + if is_leader:
124 + logger.info(
125 + "Autotuning FlashInfer SM120 sparse MLA %s decode with cache: %s",
126 + log_label,
127 + cache_path,
128 + )
129 +
130 + with torch.inference_mode():
131 + warmup_executed = True
132 + if is_leader:
133 + if _uses_v2_model_runner(runner):
134 + v2_runner = cast("V2GPUModelRunner", runner)
135 + warmup_executed = run_mixed_prefill_decode_warmup(
136 + v2_runner,
137 + worker.execute_model,
138 + worker.sample_tokens,
139 + num_tokens,
140 + mixed_step_context=flashinfer_autotune(True, cache=str(cache_path)),
141 + req_id_prefix="_sparse_mla_v2_warmup",
142 + )
143 + else:
144 + with flashinfer_autotune(True, cache=str(cache_path)):
145 + runner._dummy_run(**dummy_run_kwargs)
146 + else:
147 + if _uses_v2_model_runner(runner):
148 + v2_runner = cast("V2GPUModelRunner", runner)
149 + warmup_executed = run_mixed_prefill_decode_warmup(
150 + v2_runner,
151 + worker.execute_model,
152 + worker.sample_tokens,
153 + num_tokens,
154 + req_id_prefix="_sparse_mla_v2_warmup",
155 + )
156 + else:
157 + runner._dummy_run(**dummy_run_kwargs)
158 +
159 + if not warmup_executed:
160 + return False
161 +
162 + tune_results: bytes | None = None
163 + if is_leader and cache_path.exists():
164 + with open(cache_path, "rb") as f:
165 + tune_results = f.read()
166 +
167 + tune_results = world.broadcast_object(tune_results, src=0)
168 + if tune_results is None:
169 + logger.warning(
170 + "No FlashInfer SM120 sparse MLA %s decode autotune cache entries found. "
171 + "Falling back to FlashInfer's default tactic heuristic.",
172 + log_label,
173 + )
174 + world.barrier()
175 + return True
176 +
177 + write_flashinfer_autotune_cache(cache_path, tune_results)
178 + world.barrier()
179 +
180 + AutoTuner.get().load_configs(str(cache_path))
181 + logger.info(
182 + "FlashInfer SM120 sparse MLA %s decode autotune cache loaded on rank %d "
183 + "from %s.",
184 + log_label,
185 + world.rank_in_group,
186 + cache_path,
187 + )
188 + return True
189 +
190 +
191 +def _flashinfer_sparse_mla_decode_autotune(
192 + worker: "Worker",
193 + num_tokens: int,
194 +) -> bool:
195 + return _run_flashinfer_sparse_mla_decode_autotune(
196 + worker, num_tokens, _FLASHINFER_MLA_SPARSE_BACKENDS
197 + )
198 +
199 +
200 +def _deepseek_v4_sparse_mla_decode_autotune(
201 + worker: "Worker",
202 + num_tokens: int,
203 +) -> bool:
204 + return _run_flashinfer_sparse_mla_decode_autotune(
205 + worker, num_tokens, _DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS
206 + )
207 +
208 +
209 +def flashinfer_sparse_mla_decode_autotune_warmup(worker: "Worker") -> None:
210 + """Autotune generic FlashInfer sparse MLA decode when selected."""
211 + runner = worker.model_runner
212 + if runner.is_pooling_model:
213 + return
214 +
215 + max_tokens = worker.scheduler_config.max_num_batched_tokens
216 + mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens)
217 + if mixed_tokens <= 0:
218 + return
219 + _flashinfer_sparse_mla_decode_autotune(worker, mixed_tokens)
220 +
221 +
222 +def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None:
223 + """Warm DSv4 sparse-MLA mixed prefill+decode attention."""
224 + runner = worker.model_runner
225 + if runner.is_pooling_model or not _has_deepseek_v4_sparse_mla_backend(runner):
226 + return
227 +
228 + max_tokens = worker.scheduler_config.max_num_batched_tokens
229 + mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens)
230 + if mixed_tokens <= 0:
231 + return
232 +
233 + logger.info(
234 + "Warming up DeepSeek V4 sparse MLA attention for mixed tokens=%s.",
235 + mixed_tokens,
236 + )
237 + mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens)
238 + if not mixed_warmup_done:
239 + if _uses_v2_model_runner(runner):
240 + v2_runner = cast("V2GPUModelRunner", runner)
241 + run_mixed_prefill_decode_warmup(
242 + v2_runner,
243 + worker.execute_model,
244 + worker.sample_tokens,
245 + mixed_tokens,
246 + req_id_prefix="_sparse_mla_v2_warmup",
247 + )
248 + else:
249 + runner._dummy_run(
250 + num_tokens=mixed_tokens,
251 + skip_eplb=True,
252 + is_profile=True,
253 + force_attention=True,
254 + create_mixed_batch=True,
255 + )
vllm/model_executor/warmup/kernel_warmup.py
+28 -32
@@ -6,16 +6,24 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to
6 happen during model execution.
7 """
8
9 -import hashlib
10 -from pathlib import Path
9 from typing import TYPE_CHECKING
10
11 import torch
12
13 import vllm.envs as envs
16 -from vllm.compilation.caching import aot_compile_hash_factors
14 from vllm.logger import init_logger
15 from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
16 +from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
17 + deepseek_v4_mhc_warmup,
18 +)
19 +from vllm.model_executor.warmup.flashinfer_autotune_cache import (
20 + resolve_flashinfer_autotune_file,
21 + write_flashinfer_autotune_cache,
22 +)
23 +from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
24 + deepseek_v4_sparse_mla_attention_warmup,
25 + flashinfer_sparse_mla_decode_autotune_warmup,
26 +)
27 from vllm.platforms import current_platform
28 from vllm.utils.deep_gemm import is_deep_gemm_supported
29 from vllm.utils.flashinfer import has_flashinfer
@@ -27,36 +35,26 @@ if TYPE_CHECKING:
35 logger = init_logger(__name__)
36
37
30 -def _flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str:
31 - factors = aot_compile_hash_factors(runner.vllm_config)
32 - return hashlib.sha256(str(factors).encode()).hexdigest()
33 -
34 -
35 -def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path:
36 - override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR
37 - if override_dir:
38 - root = Path(override_dir).expanduser()
39 - else:
40 - from flashinfer.jit import env as flashinfer_jit_env
41 -
42 - flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR
43 - root = (
44 - Path(envs.VLLM_CACHE_ROOT)
45 - / "flashinfer_autotune_cache"
46 - / flashinfer_workspace.parent.name
47 - / flashinfer_workspace.name
48 - )
49 -
50 - output_dir = root / _flashinfer_autotune_cache_hash(runner)
51 - output_dir.mkdir(parents=True, exist_ok=True)
52 - return output_dir / "autotune_configs.json"
53 -
54 -
38 def kernel_warmup(worker: "Worker"):
39 from vllm.model_executor.warmup.minimax_m3_msa_warmup import (
40 minimax_m3_msa_warmup,
41 )
42
43 + # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
44 + # layer per token; warm them across token sizes first so the first real
45 + # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside).
46 + deepseek_v4_mhc_warmup(
47 + worker.get_model(),
48 + max_tokens=worker.scheduler_config.max_num_batched_tokens,
49 + cudagraph_capture_sizes=(
50 + worker.vllm_config.compilation_config.cudagraph_capture_sizes or []
51 + ),
52 + )
53 +
54 + # Run next so input-prep kernels JIT against pristine runner state.
55 + flashinfer_sparse_mla_decode_autotune_warmup(worker)
56 + deepseek_v4_sparse_mla_attention_warmup(worker)
57 +
58 # Deep GEMM warmup
59 do_deep_gemm_warmup = (
60 envs.VLLM_USE_DEEP_GEMM
@@ -147,7 +145,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
145 world = get_world_group()
146 is_leader = world.rank_in_group == 0
147
150 - cache_path = _resolve_flashinfer_autotune_file(runner)
148 + cache_path = resolve_flashinfer_autotune_file(runner)
149 if is_leader:
150 logger.info("Using FlashInfer autotune cache file: %s", cache_path)
151
@@ -183,9 +181,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
181 "Falling back to default tactics."
182 )
183 else:
186 - if not is_leader and world.local_rank == 0:
187 - with open(cache_path, "wb") as f:
188 - f.write(tune_results)
184 + write_flashinfer_autotune_cache(cache_path, tune_results)
185 world.barrier()
186 from flashinfer.autotuner import AutoTuner
187
vllm/models/deepseek_v4/attention.py
+34 -32
@@ -62,23 +62,22 @@ logger = init_logger(__name__)
62
63
64 def _resolve_dsv4_kv_cache_dtype(
65 - use_flashmla_fp8_layout: bool,
65 + use_fp8_ds_mla_layout: bool,
66 kv_cache_dtype: str,
67 cache_config: CacheConfig | None,
68 ) -> tuple[str, torch.dtype]:
69 """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``.
70
71 Both layouts are paged; they differ in the per-token block format. The
72 - FlashMLA fp8 layout (FlashMLA / ROCm Aiter) is the ``fp8_ds_mla`` format:
73 - UE8M0 block-scaled fp8 packed as ``uint8`` (the canonical ``fp8_ds_mla``
74 - string is written back onto ``cache_config`` so the page-size specs pick
75 - the 576B per-token slot). Otherwise (FlashInfer) each token's KV row is
76 - stored in its plain element dtype — bf16 or per-tensor FP8 E4M3.
72 + ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the
73 + canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the
74 + page-size specs pick the 576B per-token slot). Plain-row backends store each
75 + token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3.
76 """
78 - if use_flashmla_fp8_layout:
77 + if use_fp8_ds_mla_layout:
78 # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
79 assert kv_cache_dtype.startswith("fp8"), (
81 - f"DeepseekV4 FlashMLA fp8 layout only supports fp8 kv-cache, "
80 + f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, "
81 f"got {kv_cache_dtype}"
82 )
83 if kv_cache_dtype != "fp8_ds_mla":
@@ -100,18 +99,20 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
99
100 The platform-specific sparse-MLA forward (``forward_mqa`` /
101 ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a
103 - subclass — ``DeepseekV4FlashMLAAttention`` / ``DeepseekV4FlashInferMLAAttention``
104 - (CUDA) or ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the
105 - platform-specific deepseek_v4 model module. The base is never instantiated
106 - directly.
102 + subclass — ``DeepseekV4FlashMLAAttention`` /
103 + ``DeepseekV4FlashInferSM120Attention`` /
104 + ``DeepseekV4FlashInferMLAAttention`` (CUDA) or
105 + ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific
106 + deepseek_v4 model module. The base is never instantiated directly.
107 """
108
109 # Provided by the platform subclass.
110 backend_cls: ClassVar[type[AttentionBackend]]
111 # KV-cache per-token block format (both layouts are paged). True (default)
112 - # = FlashMLA / ROCm fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8);
113 - # False = FlashInfer plain bf16 / per-tensor fp8 KV row.
114 - use_flashmla_fp8_layout: ClassVar[bool] = True
112 + # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain
113 + # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when
114 + # a single attention class dispatches across arch-specific layouts.
115 + use_fp8_ds_mla_layout: ClassVar[bool] = True
116 # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
117 # workspace allocated in _forward_prefill and is also read by the dummy-run
118 # path to pre-reserve that workspace.
@@ -145,6 +146,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
146 """Inverse-RoPE + wo_a + wo_b output projection (platform-specific)."""
147 raise NotImplementedError
148
149 + def _uses_fp8_ds_mla_layout(self) -> bool:
150 + """Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
151 + return self.use_fp8_ds_mla_layout
152 +
153 def __init__(
154 self,
155 vllm_config: VllmConfig,
@@ -276,13 +281,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
281 )
282 self.max_model_len = vllm_config.model_config.max_model_len
283
279 - # Resolve the kv-cache dtype from this backend's block format (a
280 - # ClassVar set by the subclass): fp8_ds_mla (UE8M0 block-scaled fp8 as
281 - # uint8) for FlashMLA / ROCm, vs a plain bf16 / per-tensor fp8 row for
282 - # FlashInfer. The same resolution drives the SWA cache tensor dtype
283 - # below.
284 + # Resolve the kv-cache dtype from this backend's block format. The same
285 + # resolution drives the SWA cache tensor dtype below.
286 self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype(
285 - self.use_flashmla_fp8_layout, cache_config.cache_dtype, cache_config
287 + self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config
288 )
289
290 self.swa_cache_layer = DeepseekV4SWACache(
@@ -539,7 +541,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
541
542 # kv is unchanged; attention reads kv solely via swa_kv_cache.
543 if cache_dtype == torch.uint8:
542 - # Legacy FlashMLA UE8M0 paged path. Horizontally fused:
544 + # fp8_ds_mla UE8M0 paged path. Horizontally fused:
545 # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling
546 # the padding head slots; the kernel allocates and returns
547 # the padded q tensor.
@@ -557,10 +559,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
559 swa_metadata.block_size,
560 )
561
560 - # FlashInfer full-cache path: the [num_blocks, block_size, 512] cache
561 - # stores the KV row in its plain dtype (no Q padding). bf16 rewrites q
562 - # in place; per-tensor fp8 writes a separately-allocated fp8 q and
563 - # quantizes the KV row.
562 + # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV
563 + # row in its element dtype (no Q padding). bf16 rewrites q in place;
564 + # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
565 + # KV row.
566 block_size = swa_metadata.block_size
567 swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
568 if cache_dtype == torch.bfloat16:
@@ -601,18 +603,18 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
603 self.compress_ratio <= 1
604 ): # SWA part. Allocated separately as DeepseekV4SWACache.
605 return None
604 - # FlashMLA uses the fp8_ds_mla block format (UE8M0 block-scaled fp8 as
605 - # uint8, 576B aligned); FlashInfer stores a plain bf16 / per-tensor fp8
606 - # row with no extra alignment.
607 - is_flashmla = self.kv_cache_dtype == "fp8_ds_mla"
606 + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
607 + # alignment; plain bf16 / per-tensor fp8 rows use natural element-size
608 + # pages.
609 + uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
610 return MLAAttentionSpec(
611 block_size=vllm_config.cache_config.block_size,
612 num_kv_heads=1,
613 head_size=self.head_dim,
612 - dtype=torch.uint8 if is_flashmla else self.kv_cache_torch_dtype,
614 + dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
615 compress_ratio=self.compress_ratio,
616 cache_dtype_str=self.kv_cache_dtype,
615 - alignment=576 if is_flashmla else None, # FlashMLA needs 576B
617 + alignment=576 if uses_fp8_ds_mla_layout else None,
618 model_version="deepseek_v4",
619 )
620
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+10 -1
@@ -40,9 +40,18 @@ def _fused_inv_rope_fp8_quant_per_head(
40 USE_GDC: tl.constexpr,
41 launch_pdl: tl.constexpr, # triton metadata
42 ):
43 - # int64: stride multiply overflows int32 past num_tokens=32768 (IMA).
43 + # Cast every stride to int64 — without this, Python-int strides are
44 + # inferred as int32 and `pid_token(int64) × stride(int32)` can lower to
45 + # int32 arithmetic, wrapping past 2³¹ for large prefill batches → IMA.
46 pid_token = tl.program_id(0).to(tl.int64)
47 pid_gh = tl.program_id(1).to(tl.int64)
48 + o_stride_token = o_stride_token.to(tl.int64)
49 + o_stride_head = o_stride_head.to(tl.int64)
50 + cache_stride_pos = cache_stride_pos.to(tl.int64)
51 + fp8_stride_group = fp8_stride_group.to(tl.int64)
52 + fp8_stride_token = fp8_stride_token.to(tl.int64)
53 + scale_stride_group = scale_stride_group.to(tl.int64)
54 + scale_stride_k = scale_stride_k.to(tl.int64)
55
56 g = pid_gh // heads_per_group
57 head_in_group = pid_gh % heads_per_group
vllm/models/deepseek_v4/compressor.py
+9 -9
@@ -155,17 +155,17 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
155 raise ValueError(f"Invalid compress ratio: {compress_ratio}")
156
157 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
158 - # FlashMLA's UE8M0 paged layout needs 576B alignment; the FlashInfer
159 - # full-cache path shares state pages with contiguous KV pages, so
160 - # padding would break page matching.
161 - is_flashmla = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
158 + # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain
159 + # full-cache rows share state pages with contiguous KV pages, so padding
160 + # would break page matching.
161 + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
162 return SlidingWindowMLASpec( # only has one vector instead of K + V
163 block_size=self.block_size,
164 num_kv_heads=1,
165 head_size=self.state_dim,
166 dtype=self.dtype,
167 sliding_window=self.sliding_window,
168 - alignment=576 if is_flashmla else None,
168 + alignment=576 if uses_fp8_ds_mla_layout else None,
169 )
170
171 def forward(self): ...
@@ -340,8 +340,8 @@ class DeepseekCompressor(nn.Module):
340 k_cache_layer = self._static_forward_context[self.k_cache_prefix]
341 kv_cache = k_cache_layer.kv_cache
342
343 - # FlashInfer V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
344 - # legacy FlashMLA path uses the UE8M0 paged uint8 layout.
343 + # Plain-row V4 reads a contiguous bf16 / per-tensor fp8 cache row; the
344 + # fp8_ds_mla path uses the UE8M0 paged uint8 layout.
345 store_full_kv = self.head_dim == 512 and kv_cache.dtype != torch.uint8
346 store_full_fp8 = kv_cache.dtype == torch.float8_e4m3fn
347 fp8_scale = (
@@ -358,8 +358,8 @@ class DeepseekCompressor(nn.Module):
358 compress_norm_rope_store_cutedsl,
359 )
360
361 - # head=512 on CUDA always uses cutedsl, for both the legacy UE8M0
362 - # layout and the FlashInfer full-cache layout. The full-cache flags
361 + # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla
362 + # layout and the plain full-cache layout. The full-cache flags
363 # are consumed only here.
364 compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
365 extra_kwargs: dict[str, Any] = dict(
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+473 -27
@@ -1,13 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""DeepSeek V4 FlashInfer TRTLLM-gen sparse MLA backend.
4 -
5 -Uses FlashInfer's public ``trtllm_batch_decode_sparse_mla_dsv4`` launcher with a
6 -plain bf16 / per-tensor FP8 KV row (vs FlashMLA's packed ``fp8_ds_mla`` block
7 -format). Shares the V4 sparse-index pipeline (SWA cache + compressor + indexer,
8 -256-token blocks, head_size 512) with the FlashMLA V4 backend; only the
9 -attention forward differs.
10 -"""
3 +"""DeepSeek V4 FlashInfer sparse MLA backend."""
4
5 from typing import TYPE_CHECKING, ClassVar, cast
6
@@ -18,6 +11,7 @@ from vllm.forward_context import get_forward_context
11 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
12 from vllm.models.deepseek_v4.common.ops import (
13 build_flashinfer_mixed_sparse_indices,
14 + compute_global_topk_indices_and_lens,
15 )
16 from vllm.models.deepseek_v4.nvidia.ops.o_proj import (
17 compute_fp8_einsum_recipe,
@@ -27,13 +21,14 @@ from vllm.models.deepseek_v4.sparse_mla import (
21 DeepseekV4FlashMLABackend,
22 DeepseekV4FlashMLAMetadata,
23 )
24 +from vllm.platforms import current_platform
25 +from vllm.platforms.interface import DeviceCapability
26 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
27 +from vllm.v1.attention.backend import MultipleOf
28
29 if TYPE_CHECKING:
30 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
31
35 -# 128 MB TRTLLM-gen workspace, allocated once per device and zero-initialized
36 -# (required for first use). Reused across all FlashInfer V4 layers.
32 _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
33 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
34
@@ -51,34 +46,113 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
46
47
48 class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
54 - """Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl.
49 + """FlashInfer backend using the DSv4 sparse metadata/cache layout.
50
56 - Inheriting from the FlashMLA V4 backend reuses its ``DeepseekV4FlashMLAMetadata``
57 - builder.
51 + Inheriting from the FlashMLA V4 backend reuses its
52 + ``DeepseekV4FlashMLAMetadata`` builder.
53 """
54
55 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
61 - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16", "fp8"]
56 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
57 + "auto",
58 + "bfloat16",
59 + "fp8",
60 + "fp8_e4m3",
61 + "fp8_ds_mla",
62 + ]
63 +
64 + @staticmethod
65 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
66 + return [256]
67
68 @staticmethod
69 def get_name() -> str:
70 return "FLASHINFER_MLA_SPARSE_DSV4"
71
72 + @classmethod
73 + def get_supported_head_sizes(cls) -> list[int]:
74 + return [512]
75 +
76 + @classmethod
77 + def supports_sink(cls) -> bool:
78 + return True
79 +
80 + @classmethod
81 + def is_sparse(cls) -> bool:
82 + return True
83 +
84 + @classmethod
85 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
86 + return capability.major in [10, 12]
87 +
88 + @classmethod
89 + def supports_combination(
90 + cls,
91 + head_size: int,
92 + dtype: torch.dtype,
93 + kv_cache_dtype: CacheDType | None,
94 + block_size: int | None,
95 + use_mla: bool,
96 + has_sink: bool,
97 + use_sparse: bool,
98 + use_mm_prefix: bool,
99 + device_capability: DeviceCapability,
100 + ) -> str | None:
101 + if device_capability.major == 10:
102 + if kv_cache_dtype == "fp8_ds_mla":
103 + return (
104 + "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain "
105 + "per-tensor FP8 KV layout, not fp8_ds_mla"
106 + )
107 + if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"):
108 + return "kv_cache_dtype not supported"
109 + return None
110 + if device_capability.major == 12:
111 + if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"):
112 + return "kv_cache_dtype not supported"
113 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
114 +
115 + if not has_flashinfer_sparse_mla_sm120():
116 + return (
117 + "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's "
118 + "sparse MLA decode API"
119 + )
120 + return None
121 + return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
122 +
123 + @staticmethod
124 + def get_kv_cache_shape(
125 + num_blocks: int,
126 + block_size: int,
127 + num_kv_heads: int,
128 + head_size: int,
129 + cache_dtype_str: str = "auto",
130 + ) -> tuple[int, ...]:
131 + device_capability = current_platform.get_device_capability()
132 + if device_capability is not None and device_capability.major == 12:
133 + return DeepseekV4FlashMLABackend.get_kv_cache_shape(
134 + num_blocks,
135 + block_size,
136 + num_kv_heads,
137 + head_size,
138 + cache_dtype_str,
139 + )
140 + assert num_kv_heads == 1
141 + return (num_blocks, block_size, head_size)
142 +
143
144 class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
69 - """FlashInfer TRTLLM-gen sparse MLA attention layer for DeepSeek V4."""
145 + """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
146
147 backend_cls = DeepseekV4FlashInferMLASparseBackend
72 - # FlashInfer stores a plain bf16 / per-tensor fp8 KV row, not the FlashMLA
73 - # packed fp8_ds_mla block format (UE8M0 block-scaled fp8 as uint8).
74 - use_flashmla_fp8_layout: ClassVar[bool] = False
148 + use_fp8_ds_mla_layout: ClassVar[bool] = False
149
150 @classmethod
151 def get_padded_num_q_heads(cls, num_heads: int) -> int:
152 # FP8 decode kernel only supports h_q = 64 or 128.
153 if num_heads > 128:
154 raise ValueError(
81 - f"DeepseekV4 Flashinfer MLA Sparse does not support {num_heads} heads "
155 + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
156 "(FP8 decode kernel requires h_q in {64, 128})."
157 )
158 return 64 if num_heads <= 64 else 128
@@ -106,8 +180,6 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
180 # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
181 if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
182 return
109 - # TODO: load real per-tensor Q/KV scales from the checkpoint; unit
110 - # scales until the scale tensor names are wired.
183 fp8_q_scale = 1.0
184 fp8_kv_scale = 1.0
185 self.register_buffer(
@@ -125,9 +197,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
197 torch.tensor([fp8_kv_scale], dtype=torch.float32),
198 persistent=False,
199 )
128 - # TRTLLM-gen takes scalar scale args on a distinct (correct) C++ path
129 - # vs 1-elem tensors, so these are Python floats. bmm1 folds the softmax
130 - # scale and the Q/KV per-tensor scales; bmm2 is the KV scale.
200 + # TRTLLM-gen takes scalar scale args on a distinct C++ path vs
201 + # one-element tensors, so these are Python floats.
202 self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
203 self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
204
@@ -387,9 +458,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
458 query_start_loc_cpu = swa_metadata.query_start_loc_cpu
459 assert query_start_loc is not None and query_start_loc_cpu is not None
460
390 - # Keep Perkz's two-call decode/prefill split: the TRTLLM-gen launcher is
391 - # tuned for uniform-q batches, and collapsing the mixed batch into a
392 - # single call is the suspected source of the prior IMA.
461 + # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for
462 + # uniform-q batches, and this avoids flattening mixed batches into one call.
463 if num_decode_tokens > 0:
464 decode_cu = query_start_loc[: num_decodes + 1]
465 decode_cu_cpu = query_start_loc_cpu[: num_decodes + 1]
@@ -434,3 +504,379 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
504 cum_seq_lens_q=prefill_cu,
505 max_q_len=int(prefill_lens_cpu.max().item()),
506 )
507 +
508 +
509 +class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
510 + """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels."""
511 +
512 + backend_cls = DeepseekV4FlashInferMLASparseBackend
513 + use_fp8_ds_mla_layout: ClassVar[bool] = True
514 +
515 + @staticmethod
516 + def _get_workspace(device: torch.device) -> torch.Tensor:
517 + return _get_flashinfer_dsv4_workspace(device)
518 +
519 + @staticmethod
520 + def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor:
521 + if kv_cache.dtype == torch.float8_e4m3fn:
522 + kv_cache = kv_cache.view(torch.uint8)
523 + if kv_cache.dim() == 4:
524 + return kv_cache
525 + return kv_cache.unsqueeze(-2)
526 +
527 + @classmethod
528 + def get_padded_num_q_heads(cls, num_heads: int) -> int:
529 + if num_heads <= 16:
530 + return 16
531 + if num_heads <= 32:
532 + return 32
533 + if num_heads <= 64:
534 + return 64
535 + if num_heads <= 128:
536 + return 128
537 + raise ValueError(
538 + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
539 + "(SM120 kernel requires h_q in {16, 32, 64, 128})."
540 + )
541 +
542 + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
543 + return deep_gemm_fp8_o_proj(
544 + o,
545 + positions,
546 + self.rotary_emb.cos_sin_cache,
547 + self.wo_a,
548 + self.wo_b,
549 + n_groups=self.n_local_groups,
550 + heads_per_group=self.n_local_heads // self.n_local_groups,
551 + nope_dim=self.nope_head_dim,
552 + rope_dim=self.rope_head_dim,
553 + o_lora_rank=self.o_lora_rank,
554 + einsum_recipe=self._einsum_recipe,
555 + tma_aligned_scales=self._tma_aligned_scales,
556 + )
557 +
558 + def __init__(self, *args, **kwargs) -> None:
559 + super().__init__(*args, **kwargs)
560 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
561 +
562 + if not has_flashinfer_sparse_mla_sm120():
563 + raise RuntimeError(
564 + "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires FlashInfer's "
565 + "sparse MLA decode API."
566 + )
567 + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
568 + # Per-tensor FP8 cache path scales.
569 + if self.kv_cache_torch_dtype != torch.float8_e4m3fn:
570 + return
571 + fp8_q_scale = 1.0
572 + fp8_kv_scale = 1.0
573 + self.register_buffer(
574 + "_flashinfer_fp8_q_scale",
575 + torch.tensor([fp8_q_scale], dtype=torch.float32),
576 + persistent=False,
577 + )
578 + self.register_buffer(
579 + "_flashinfer_fp8_q_scale_inv",
580 + torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32),
581 + persistent=False,
582 + )
583 + self.register_buffer(
584 + "_flashinfer_fp8_kv_scale",
585 + torch.tensor([fp8_kv_scale], dtype=torch.float32),
586 + persistent=False,
587 + )
588 + # FlashInfer expects scalar scale arguments for this path.
589 + self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale
590 + self._flashinfer_fp8_bmm2_scale = fp8_kv_scale
591 +
592 + def _reserve_empty_forward_workspace(self) -> None:
593 + self._get_workspace(
594 + torch.device("cuda", torch.accelerator.current_device_index())
595 + )
596 +
597 + def _forward_sparse_impl(
598 + self,
599 + q: torch.Tensor,
600 + output: torch.Tensor,
601 + flashmla_metadata: DeepseekV4FlashMLAMetadata | None,
602 + swa_metadata: "DeepseekSparseSWAMetadata",
603 + self_kv_cache: torch.Tensor | None,
604 + swa_kv_cache: torch.Tensor,
605 + swa_only: bool,
606 + ) -> None:
607 + num_decode_tokens = swa_metadata.num_decode_tokens
608 + if swa_metadata.num_prefills > 0:
609 + self._forward_prefill(
610 + q=q[num_decode_tokens:],
611 + compressed_k_cache=self_kv_cache,
612 + swa_k_cache=swa_kv_cache,
613 + output=output[num_decode_tokens:],
614 + attn_metadata=flashmla_metadata,
615 + swa_metadata=swa_metadata,
616 + )
617 + if swa_metadata.num_decodes > 0:
618 + self._forward_decode(
619 + q=q[:num_decode_tokens],
620 + kv_cache=self_kv_cache,
621 + swa_metadata=swa_metadata,
622 + attn_metadata=flashmla_metadata,
623 + swa_only=swa_only,
624 + output=output[:num_decode_tokens],
625 + )
626 +
627 + def forward_mqa(
628 + self,
629 + q: torch.Tensor,
630 + kv: torch.Tensor,
631 + positions: torch.Tensor,
632 + output: torch.Tensor,
633 + ) -> None:
634 + # Output may be padded to backend-supported head counts.
635 + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], (
636 + f"output buffer shape {output.shape} incompatible with q shape {q.shape}"
637 + )
638 + assert output.shape[1] >= q.shape[1], (
639 + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}"
640 + )
641 + # Per-tensor FP8 q produces a bf16 attention output.
642 + expected_output_dtype = (
643 + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype
644 + )
645 + assert output.dtype == expected_output_dtype, (
646 + f"output dtype {output.dtype} must match expected {expected_output_dtype} "
647 + f"for q dtype {q.dtype}"
648 + )
649 +
650 + forward_context = get_forward_context()
651 + attn_metadata = forward_context.attn_metadata
652 + if attn_metadata is None:
653 + self._reserve_empty_forward_workspace()
654 + output.zero_()
655 + return
656 +
657 + assert isinstance(attn_metadata, dict)
658 + flashmla_metadata = cast(
659 + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix)
660 + )
661 + swa_metadata = cast(
662 + "DeepseekSparseSWAMetadata | None",
663 + attn_metadata.get(self.swa_cache_layer.prefix),
664 + )
665 + assert swa_metadata is not None
666 +
667 + swa_only = self.compress_ratio <= 1
668 + # SWA-only layers don't allocate their own compressed KV cache.
669 + self_kv_cache = self.kv_cache if not swa_only else None
670 + swa_kv_cache = self.swa_cache_layer.kv_cache
671 +
672 + self._forward_sparse_impl(
673 + q=q,
674 + output=output,
675 + flashmla_metadata=flashmla_metadata,
676 + swa_metadata=swa_metadata,
677 + self_kv_cache=self_kv_cache,
678 + swa_kv_cache=swa_kv_cache,
679 + swa_only=swa_only,
680 + )
681 +
682 + def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
683 + if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
684 + assert q.dtype == torch.float8_e4m3fn
685 + q = q.to(torch.bfloat16)
686 + else:
687 + assert q.dtype == torch.bfloat16
688 + padded_heads = output.shape[1]
689 + if q.shape[1] < padded_heads:
690 + padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2]))
691 + padded_query[:, : q.shape[1], :] = q
692 + q = padded_query
693 + return q.contiguous()
694 +
695 + def _forward_decode(
696 + self,
697 + q: torch.Tensor,
698 + kv_cache: torch.Tensor | None,
699 + swa_metadata: "DeepseekSparseSWAMetadata",
700 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
701 + swa_only: bool,
702 + output: torch.Tensor,
703 + ) -> None:
704 + num_decodes = swa_metadata.num_decodes
705 + num_decode_tokens = swa_metadata.num_decode_tokens
706 +
707 + extra_sparse_indices = None
708 + extra_sparse_lengths = None
709 + if not swa_only:
710 + if attn_metadata is None:
711 + raise RuntimeError(
712 + "Sparse MLA metadata is required for compressed layers."
713 + )
714 + if swa_metadata.is_valid_token is None:
715 + raise RuntimeError(
716 + "SWA validity metadata is required for compressed layers."
717 + )
718 + is_valid = swa_metadata.is_valid_token[:num_decode_tokens]
719 + if self.compress_ratio == 4:
720 + if self.topk_indices_buffer is None:
721 + raise RuntimeError(
722 + "C4A decode requires top-k indices from the indexer."
723 + )
724 + block_size = attn_metadata.block_size // self.compress_ratio
725 + global_indices, extra_sparse_lengths = (
726 + compute_global_topk_indices_and_lens(
727 + self.topk_indices_buffer[:num_decode_tokens],
728 + swa_metadata.token_to_req_indices,
729 + attn_metadata.block_table[:num_decodes],
730 + block_size,
731 + is_valid,
732 + )
733 + )
734 + extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1)
735 + else:
736 + extra_sparse_indices = attn_metadata.c128a_global_decode_topk_indices
737 + extra_sparse_lengths = attn_metadata.c128a_decode_topk_lens
738 +
739 + swa_indices = swa_metadata.decode_swa_indices
740 + swa_lens = swa_metadata.decode_swa_lens
741 + assert swa_indices is not None
742 + assert swa_lens is not None
743 + q = self._prepare_query(q, output)
744 + swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache)
745 + extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None
746 + if extra_cache is not None and extra_sparse_indices is None:
747 + raise RuntimeError(
748 + "Compressed sparse MLA decode requires compressed sparse indices."
749 + )
750 + flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
751 + query=q,
752 + swa_kv_cache=swa_cache,
753 + workspace_buffer=self._get_workspace(q.device),
754 + sparse_indices=swa_indices,
755 + compressed_kv_cache=extra_cache,
756 + out=output,
757 + bmm1_scale=self.scale,
758 + sinks=self.attn_sink,
759 + kv_layout="NHD",
760 + swa_topk_lens=swa_lens,
761 + extra_sparse_indices=extra_sparse_indices,
762 + extra_sparse_topk_lens=extra_sparse_lengths,
763 + )
764 +
765 + def _forward_prefill(
766 + self,
767 + q: torch.Tensor,
768 + compressed_k_cache: torch.Tensor | None,
769 + swa_k_cache: torch.Tensor,
770 + output: torch.Tensor,
771 + attn_metadata: DeepseekV4FlashMLAMetadata | None,
772 + swa_metadata: "DeepseekSparseSWAMetadata",
773 + ) -> None:
774 + swa_only = self.compress_ratio <= 1
775 +
776 + num_prefills = swa_metadata.num_prefills
777 + num_decodes = swa_metadata.num_decodes
778 + num_decode_tokens = swa_metadata.num_decode_tokens
779 + num_prefill_tokens = swa_metadata.num_prefill_tokens
780 +
781 + query_start_loc_cpu = swa_metadata.query_start_loc_cpu
782 + assert query_start_loc_cpu is not None
783 + prefill_token_base = query_start_loc_cpu[num_decodes]
784 +
785 + local_topk_indices: torch.Tensor | None
786 + if swa_only:
787 + local_topk_indices = None
788 + elif self.compress_ratio == 4:
789 + if self.topk_indices_buffer is None:
790 + raise RuntimeError(
791 + "C4A prefill requires top-k indices from the indexer."
792 + )
793 + local_topk_indices = self.topk_indices_buffer[
794 + num_decode_tokens : num_decode_tokens + num_prefill_tokens
795 + ]
796 + else:
797 + if attn_metadata is None:
798 + raise RuntimeError("C128A prefill metadata is missing.")
799 + local_topk_indices = attn_metadata.c128a_prefill_topk_indices
800 +
801 + extra_sparse_indices: torch.Tensor | None = None
802 + extra_sparse_lengths: torch.Tensor | None = None
803 + if local_topk_indices is not None:
804 + if attn_metadata is None:
805 + raise RuntimeError("C4A prefill metadata is missing.")
806 + if swa_metadata.token_to_req_indices is None:
807 + raise RuntimeError("C4A prefill request mapping is missing.")
808 + if swa_metadata.is_valid_token is None:
809 + raise RuntimeError("C4A prefill validity metadata is missing.")
810 + prefill_token_slice = slice(
811 + num_decode_tokens, num_decode_tokens + num_prefill_tokens
812 + )
813 + block_size = attn_metadata.block_size // self.compress_ratio
814 + extra_sparse_indices, extra_sparse_lengths = (
815 + compute_global_topk_indices_and_lens(
816 + local_topk_indices,
817 + swa_metadata.token_to_req_indices[prefill_token_slice],
818 + attn_metadata.block_table,
819 + block_size,
820 + swa_metadata.is_valid_token[prefill_token_slice],
821 + )
822 + )
823 +
824 + assert swa_metadata.prefill_swa_indices is not None
825 + assert swa_metadata.prefill_swa_lens is not None
826 +
827 + q = self._prepare_query(q, output)
828 + swa_kv_paged = self._as_sparse_cache(swa_k_cache)
829 + if swa_only:
830 + extra_kv_paged = None
831 + else:
832 + if compressed_k_cache is None:
833 + raise RuntimeError(
834 + "Compressed sparse MLA layers require their compressed KV cache."
835 + )
836 + extra_kv_paged = self._as_sparse_cache(compressed_k_cache)
837 +
838 + num_chunks = (
839 + num_prefills + self.PREFILL_CHUNK_SIZE - 1
840 + ) // self.PREFILL_CHUNK_SIZE
841 + for chunk_idx in range(num_chunks):
842 + chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE
843 + chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills)
844 + query_start = (
845 + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base
846 + )
847 + query_end = (
848 + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base
849 + )
850 +
851 + extra_sparse_indices_chunk = (
852 + extra_sparse_indices[query_start:query_end]
853 + if extra_sparse_indices is not None
854 + else None
855 + )
856 + extra_sparse_lengths_chunk = (
857 + extra_sparse_lengths[query_start:query_end]
858 + if extra_sparse_lengths is not None
859 + else None
860 + )
861 +
862 + q_chunk = q[query_start:query_end]
863 + swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end]
864 + swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end]
865 + if extra_kv_paged is not None and extra_sparse_indices_chunk is None:
866 + raise RuntimeError(
867 + "Compressed sparse MLA prefill requires compressed sparse indices."
868 + )
869 + flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
870 + query=q_chunk,
871 + swa_kv_cache=swa_kv_paged,
872 + workspace_buffer=self._get_workspace(q.device),
873 + sparse_indices=swa_indices_chunk,
874 + compressed_kv_cache=extra_kv_paged,
875 + out=output[query_start:query_end],
876 + bmm1_scale=self.scale,
877 + sinks=self.attn_sink,
878 + kv_layout="NHD",
879 + swa_topk_lens=swa_lens_chunk,
880 + extra_sparse_indices=extra_sparse_indices_chunk,
881 + extra_sparse_topk_lens=extra_sparse_lengths_chunk,
882 + )
vllm/models/deepseek_v4/nvidia/model.py
+27 -5
@@ -60,9 +60,11 @@ from vllm.model_executor.utils import set_weight_attrs
60 from vllm.models.deepseek_v4.attention import DeepseekV4Attention
61 from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
62 DeepseekV4FlashInferMLAAttention,
63 + DeepseekV4FlashInferSM120Attention,
64 )
65 from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
66 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
67 +from vllm.platforms import current_platform
68 from vllm.sequence import IntermediateTensors
69 from vllm.utils.math_utils import cdiv
70 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -736,14 +738,34 @@ class DeepseekV4MoE(nn.Module):
738 def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:
739 """Pick the CUDA sparse-MLA attention class for the configured backend.
740
739 - An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the
740 - FlashInfer TRTLLM-gen path; otherwise the FlashMLA path is used.
741 + The generic CUDA backend selector does not instantiate DSv4 layers directly,
742 + so map generic sparse-MLA choices to the DSv4-specialized attention class.
743 + Without an explicit backend, SM12 defaults to FlashInfer while the other
744 + CUDA arches keep the FlashMLA path.
745 """
742 - if (
743 - vllm_config.attention_config.backend
744 - == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4
746 + backend = vllm_config.attention_config.backend
747 + device_capability = current_platform.get_device_capability()
748 + if backend in (
749 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
750 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
751 ):
752 + raise ValueError(
753 + f"{backend.name} is not a DeepSeek V4 attention backend. "
754 + "Use FLASHINFER_MLA_SPARSE_DSV4 for DeepSeek V4 FlashInfer "
755 + "sparse MLA."
756 + )
757 + if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4:
758 + if device_capability is not None and device_capability.major == 12:
759 + return DeepseekV4FlashInferSM120Attention
760 return DeepseekV4FlashInferMLAAttention
761 + if backend in (
762 + AttentionBackendEnum.FLASHMLA_SPARSE,
763 + AttentionBackendEnum.FLASHMLA_SPARSE_DSV4,
764 + ):
765 + return DeepseekV4FlashMLAAttention
766 +
767 + if device_capability is not None and device_capability.major == 12:
768 + return DeepseekV4FlashInferSM120Attention
769 return DeepseekV4FlashMLAAttention
770
771
vllm/models/deepseek_v4/sparse_mla.py
+4
@@ -86,6 +86,10 @@ class DeepseekV4FlashMLABackend(AttentionBackend):
86 def is_sparse(cls) -> bool:
87 return True
88
89 + @classmethod
90 + def supports_sink(cls) -> bool:
91 + return True
92 +
93 @classmethod
94 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
95 return capability.major in [9, 10]
vllm/platforms/cuda.py
+46 -15
@@ -11,7 +11,7 @@ import platform
11 from collections.abc import Callable
12 from datetime import timedelta
13 from functools import cache, lru_cache, wraps
14 -from typing import TYPE_CHECKING, TypeVar
14 +from typing import TYPE_CHECKING, NamedTuple, TypeVar
15
16 import torch
17 from torch.distributed import PrefixStore, ProcessGroup
@@ -31,6 +31,7 @@ if TYPE_CHECKING:
31 from vllm.config import VllmConfig
32 from vllm.config.cache import CacheDType
33 from vllm.config.kernel import IrOpPriorityConfig
34 + from vllm.v1.attention.backend import AttentionBackend
35 from vllm.v1.attention.selector import AttentionSelectorConfig
36 else:
37 VllmConfig = None
@@ -126,6 +127,11 @@ def _get_backend_priorities(
127 AttentionBackendEnum.TRITON_MLA,
128 *sparse_backends,
129 ]
130 + elif device_capability.major == 12:
131 + return [
132 + AttentionBackendEnum.TRITON_MLA,
133 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
134 + ]
135 else:
136 return [
137 AttentionBackendEnum.FLASH_ATTN_MLA,
@@ -153,6 +159,21 @@ def _get_backend_priorities(
159 ]
160
161
162 +def _backend_cls_path(backend_cls: type[AttentionBackend]) -> str:
163 + module, qualname = backend_cls.full_cls_name()
164 + return f"{module}.{qualname}"
165 +
166 +
167 +def _get_attn_backend_class(backend: AttentionBackendEnum) -> type[AttentionBackend]:
168 + return backend.get_class()
169 +
170 +
171 +class _BackendCandidate(NamedTuple):
172 + backend_class: type[AttentionBackend]
173 + backend: AttentionBackendEnum
174 + priority: int
175 +
176 +
177 def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
178 @wraps(fn)
179 def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
@@ -340,7 +361,7 @@ class CudaPlatformBase(Platform):
361 attn_selector_config: AttentionSelectorConfig,
362 num_heads: int | None = None,
363 ) -> tuple[
343 - list[tuple[AttentionBackendEnum, int]],
364 + list[_BackendCandidate],
365 dict[AttentionBackendEnum, tuple[int, list[str]]],
366 ]:
367 valid_backends_priorities = []
@@ -354,7 +375,7 @@ class CudaPlatformBase(Platform):
375 )
376 for priority, backend in enumerate(backend_priorities):
377 try:
357 - backend_class = backend.get_class()
378 + backend_class = _get_attn_backend_class(backend)
379 invalid_reasons_i = backend_class.validate_configuration(
380 device_capability=device_capability,
381 **attn_selector_config._asdict(),
@@ -364,7 +385,9 @@ class CudaPlatformBase(Platform):
385 if invalid_reasons_i:
386 invalid_reasons[backend] = (priority, invalid_reasons_i)
387 else:
367 - valid_backends_priorities.append((backend, priority))
388 + valid_backends_priorities.append(
389 + _BackendCandidate(backend_class, backend, priority)
390 + )
391
392 return valid_backends_priorities, invalid_reasons
393
@@ -381,7 +404,7 @@ class CudaPlatformBase(Platform):
404 # First try checking just the selected backend, if there is one.
405 if selected_backend is not None:
406 try:
384 - backend_class = selected_backend.get_class()
407 + backend_class = _get_attn_backend_class(selected_backend)
408 invalid_reasons = backend_class.validate_configuration(
409 device_capability=device_capability,
410 **attn_selector_config._asdict(),
@@ -395,7 +418,7 @@ class CudaPlatformBase(Platform):
418 )
419 else:
420 logger.info("Using %s backend.", selected_backend)
398 - return selected_backend.get_path()
421 + return _backend_cls_path(backend_class)
422
423 # No selected backend or the selected backend is invalid,
424 # so we try finding a valid backend.
@@ -425,13 +448,13 @@ class CudaPlatformBase(Platform):
448
449 # We have found some valid backends. Select the one with the
450 # highest priority.
428 - sorted_indices = sorted(
429 - range(len(valid_backends_priorities)),
430 - key=lambda i: valid_backends_priorities[i][1],
451 + selected_candidate = min(
452 + valid_backends_priorities,
453 + key=lambda candidate: candidate.priority,
454 )
432 - selected_index = sorted_indices[0]
433 - selected_backend = valid_backends_priorities[selected_index][0]
434 - selected_priority = valid_backends_priorities[selected_index][1]
455 + selected_backend_class = selected_candidate.backend_class
456 + selected_backend = selected_candidate.backend
457 + selected_priority = selected_candidate.priority
458
459 # If the user specified --block-size (but not --attention-backend),
460 # check whether that constraint precluded any higher-priority backends.
@@ -457,10 +480,14 @@ class CudaPlatformBase(Platform):
480 logger.info_once(
481 "Using %s attention backend out of potential backends: %s.",
482 selected_backend.name,
460 - "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]",
483 + "["
484 + + ", ".join(
485 + f"'{candidate.backend.name}'" for candidate in valid_backends_priorities
486 + )
487 + + "]",
488 )
489
463 - return selected_backend.get_path()
490 + return _backend_cls_path(selected_backend_class)
491
492 @classmethod
493 def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]:
@@ -635,7 +662,11 @@ class CudaPlatformBase(Platform):
662 @classmethod
663 def support_deep_gemm(cls) -> bool:
664 """Currently, only Hopper and Blackwell GPUs are supported."""
638 - return cls.is_device_capability(90) or cls.is_device_capability_family(100)
665 + return (
666 + cls.is_device_capability(90)
667 + or cls.is_device_capability_family(100)
668 + or cls.is_device_capability_family(120)
669 + )
670
671 @classmethod
672 def is_integrated_gpu(cls, device_id: int = 0) -> bool:
vllm/utils/deep_gemm.py
+144 -3
@@ -5,6 +5,7 @@
5 Users of vLLM should always import **only** these wrappers.
6 """
7
8 +import contextlib
9 import functools
10 import importlib
11 import os
@@ -37,7 +38,10 @@ def should_auto_disable_deep_gemm(model_type: str | None) -> bool:
38 """
39 if model_type is None:
40 return False
40 - if not current_platform.is_device_capability_family(100):
41 + if not (
42 + current_platform.is_device_capability_family(100)
43 + or current_platform.is_device_capability_family(120)
44 + ):
45 return False
46 return model_type in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES
47
@@ -71,7 +75,10 @@ class DeepGemmQuantScaleFMT(Enum):
75
76 cls._oracle_cache = ( # type: ignore
77 cls.UE8M0
74 - if current_platform.is_device_capability_family(100)
78 + if (
79 + current_platform.is_device_capability_family(100)
80 + or current_platform.is_device_capability_family(120)
81 + )
82 else cls.FLOAT32_CEIL_UE8M0
83 )
84
@@ -138,7 +145,15 @@ _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
145 _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
146 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
147 _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
148 +_get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = (
149 + None
150 +)
151 _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None
152 +_pack_ue8m0_to_int_impl: Callable[..., Any] | None = None
153 +_get_mn_major_tma_aligned_packed_ue8m0_tensor_impl: Callable[..., Any] | None = None
154 +_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl: (
155 + Callable[..., Any] | None
156 +) = None
157
158
159 @functools.cache
@@ -203,7 +218,11 @@ def _lazy_init() -> None:
218 global _tf32_hc_prenorm_gemm_impl
219 global _get_mn_major_tma_aligned_tensor_impl
220 global _get_mk_alignment_for_contiguous_layout_impl
221 + global _get_theoretical_mk_alignment_for_contiguous_layout_impl
222 global _transform_sf_into_required_layout_impl
223 + global _pack_ue8m0_to_int_impl
224 + global _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl
225 + global _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl
226 # fast path
227 if (
228 _cublaslt_gemm_nt_impl is not None
@@ -218,6 +237,9 @@ def _lazy_init() -> None:
237 or _tf32_hc_prenorm_gemm_impl is not None
238 or _get_mk_alignment_for_contiguous_layout_impl is not None
239 or _transform_sf_into_required_layout_impl is not None
240 + or _pack_ue8m0_to_int_impl is not None
241 + or _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None
242 + or _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None
243 ):
244 return
245
@@ -258,9 +280,19 @@ def _lazy_init() -> None:
280 _get_mk_alignment_for_contiguous_layout_impl = getattr(
281 _dg, "get_mk_alignment_for_contiguous_layout", None
282 )
283 + _get_theoretical_mk_alignment_for_contiguous_layout_impl = getattr(
284 + _dg, "get_theoretical_mk_alignment_for_contiguous_layout", None
285 + )
286 _transform_sf_into_required_layout_impl = getattr(
287 _dg, "transform_sf_into_required_layout", None
288 )
289 + _pack_ue8m0_to_int_impl = getattr(_dg, "pack_ue8m0_to_int", None)
290 + _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr(
291 + _dg, "get_mn_major_tma_aligned_packed_ue8m0_tensor", None
292 + )
293 + _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr(
294 + _dg, "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", None
295 + )
296 DeepGemmQuantScaleFMT.init_oracle_cache()
297
298
@@ -280,7 +312,6 @@ def set_num_sms(num_sms: int) -> None:
312 dg.set_num_sms(num_sms)
313
314
283 -@functools.cache
315 def get_mk_alignment_for_contiguous_layout() -> list[int]:
316 _lazy_init()
317 if _get_mk_alignment_for_contiguous_layout_impl is None:
@@ -289,6 +320,70 @@ def get_mk_alignment_for_contiguous_layout() -> list[int]:
320 return [mk_align_size, mk_align_size]
321
322
323 +def get_theoretical_mk_alignment_for_contiguous_layout(
324 + expected_m: int | None = None,
325 + num_groups: int | None = None,
326 +) -> int:
327 + """Per-call optimal M alignment for grouped contiguous GEMMs.
328 +
329 + `expected_m` is the TOTAL routed tokens (sum across experts, typically
330 + M × num_topk). `num_groups` is the number of experts on this rank.
331 + The helper divides to recover per-expert em and picks an alignment based
332 + on data-driven thresholds (see deep_gemm runtime.hpp comments).
333 +
334 + Older callers that omit `num_groups` are interpreted as passing already
335 + per-expert em (legacy behaviour preserved for backward compat).
336 + """
337 + _lazy_init()
338 + if _get_theoretical_mk_alignment_for_contiguous_layout_impl is None:
339 + return _missing()
340 + if num_groups is None:
341 + return _get_theoretical_mk_alignment_for_contiguous_layout_impl(expected_m)
342 + if num_groups <= 0:
343 + raise ValueError(f"num_groups must be positive, got {num_groups}")
344 + try:
345 + return _get_theoretical_mk_alignment_for_contiguous_layout_impl(
346 + expected_m, num_groups
347 + )
348 + except TypeError:
349 + per_group_m = None if expected_m is None else cdiv(expected_m, num_groups)
350 + return _get_theoretical_mk_alignment_for_contiguous_layout_impl(per_group_m)
351 +
352 +
353 +def set_mk_alignment_for_contiguous_layout(value: int) -> None:
354 + """Set DeepGEMM's BLOCK_M cap for grouped contiguous GEMMs.
355 +
356 + The DG heuristic constrains BLOCK_M ≤ this value when picking a kernel
357 + layout. Use this in concert with `compute_aligned_M_and_alignment`'s
358 + per-call alignment so the workspace's per-expert padding matches the
359 + kernel's BLOCK_M; a mismatch leads to the scheduler reading the wrong
360 + expert_id from `m_indices` at `m_block_idx * BLOCK_M` stride and
361 + OOB-indexing the B-weights tensor (manifests as IMA under CUDA-graph
362 + replay).
363 + """
364 + _lazy_init()
365 + dg = _import_deep_gemm()
366 + if dg is None:
367 + raise RuntimeError("DeepGEMM is not available")
368 + dg.set_mk_alignment_for_contiguous_layout(value)
369 +
370 +
371 +@contextlib.contextmanager
372 +def mk_alignment_scope(value: int):
373 + """Temporarily set DeepGEMM's BLOCK_M cap, restoring on exit.
374 +
375 + Use around a sequence of grouped-contiguous GEMM calls whose workspace
376 + is padded to `value` (typically the per_call_align returned by
377 + `compute_aligned_M_and_alignment`).
378 + """
379 + prev = get_mk_alignment_for_contiguous_layout()[0]
380 + set_mk_alignment_for_contiguous_layout(value)
381 + try:
382 + yield
383 + finally:
384 + set_mk_alignment_for_contiguous_layout(prev)
385 +
386 +
387 def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
388 """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor"""
389 _lazy_init()
@@ -297,6 +392,48 @@ def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
392 return _get_mn_major_tma_aligned_tensor_impl(x)
393
394
395 +def pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor:
396 + """Pack 4 UE8M0 (uint8) scales into one int32.
397 +
398 + DeepGEMM's SM100/SM120 FP8/FP4 kernels accept either ``float32`` scales
399 + (legacy format, 4 B/scale) or ``int32`` packed UE8M0 scales (1 B/scale
400 + after 4:1 packing — 4× smaller than the legacy fp32 representation).
401 + """
402 + _lazy_init()
403 + if _pack_ue8m0_to_int_impl is None:
404 + return _missing()
405 + return _pack_ue8m0_to_int_impl(x)
406 +
407 +
408 +def get_mn_major_tma_aligned_packed_ue8m0_tensor(x: torch.Tensor) -> torch.Tensor:
409 + """Pack UE8M0 (uint8) → int32 with the MN-major TMA-aligned layout the
410 + DeepGEMM kernels consume directly. 16× smaller than the fp32 legacy SF
411 + format. Use for non-grouped 2D scale tensors.
412 + """
413 + _lazy_init()
414 + if _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None:
415 + return _missing()
416 + return _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl(x)
417 +
418 +
419 +def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(
420 + sf: torch.Tensor,
421 + ks_tensor: torch.Tensor,
422 + ks: list[int],
423 + gran_k: int,
424 +) -> torch.Tensor:
425 + """Grouped (3D, expert-batched) variant of
426 + ``get_mn_major_tma_aligned_packed_ue8m0_tensor``. Use for MoE weight
427 + scale tensors of shape ``(num_experts, mn, k_scale)``.
428 + """
429 + _lazy_init()
430 + if _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None:
431 + return _missing()
432 + return _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl(
433 + sf, ks_tensor, ks, gran_k
434 + )
435 +
436 +
437 def cublaslt_gemm_nt(*args, **kwargs):
438 _lazy_init()
439 if _cublaslt_gemm_nt_impl is None:
@@ -601,4 +738,8 @@ __all__ = [
738 "should_use_deepgemm_for_fp8_linear",
739 "get_col_major_tma_aligned_tensor",
740 "get_mk_alignment_for_contiguous_layout",
741 + "get_theoretical_mk_alignment_for_contiguous_layout",
742 + "pack_ue8m0_to_int",
743 + "get_mn_major_tma_aligned_packed_ue8m0_tensor",
744 + "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor",
745 ]
vllm/utils/flashinfer.py
+33 -9
@@ -72,11 +72,10 @@ def _missing(*_: Any, **__: Any) -> NoReturn:
72 )
73
74
75 -def _missing_dsv4_sparse_mla(*_: Any, **__: Any) -> NoReturn:
75 +def _missing_sparse_mla(*_: Any, **__: Any) -> NoReturn:
76 raise RuntimeError(
77 - "flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4 is not available. "
78 - "Install a FlashInfer build that includes DeepSeek V4 sparse MLA "
79 - "TRTLLM-GEN support."
77 + "FlashInfer sparse MLA decode APIs are not available. "
78 + "Install a FlashInfer build that includes sparse MLA decode support."
79 )
80
81
@@ -149,14 +148,18 @@ flashinfer_b12x_fused_moe = _lazy_import_wrapper(
148 trtllm_fp4_block_scale_moe = _lazy_import_wrapper(
149 "flashinfer", "trtllm_fp4_block_scale_moe"
150 )
152 -# DeepSeek V4 sparse MLA TRTLLM-GEN decode launcher (public wrapper). Handles
153 -# the SWA + compressed KV pools, the concatenated sparse-index matrix, and
154 -# per-tensor FP8 / BF16 inputs with BF16 output.
151 +flashinfer_trtllm_batch_decode_with_kv_cache_mla = _lazy_import_wrapper(
152 + "flashinfer.decode",
153 + "trtllm_batch_decode_with_kv_cache_mla",
154 + fallback_fn=_missing_sparse_mla,
155 +)
156 flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper(
156 - "flashinfer.mla",
157 + "flashinfer.decode",
158 "trtllm_batch_decode_sparse_mla_dsv4",
158 - fallback_fn=_missing_dsv4_sparse_mla,
159 + fallback_fn=_missing_sparse_mla,
160 )
161 +
162 +
163 # Special case for autotune since it returns a context manager
164 autotune = _lazy_import_wrapper(
165 "flashinfer.autotuner",
@@ -209,6 +212,26 @@ def has_flashinfer_moe() -> bool:
212 )
213
214
215 +@functools.cache
216 +def has_flashinfer_sparse_mla_sm120() -> bool:
217 + """Return ``True`` if FlashInfer sparse MLA decode support is available."""
218 + if not has_flashinfer():
219 + return False
220 + try:
221 + from flashinfer.autotuner import autotune
222 + from flashinfer.decode import (
223 + trtllm_batch_decode_sparse_mla_dsv4,
224 + trtllm_batch_decode_with_kv_cache_mla,
225 + )
226 + except ImportError:
227 + return False
228 + return (
229 + callable(trtllm_batch_decode_sparse_mla_dsv4)
230 + and callable(trtllm_batch_decode_with_kv_cache_mla)
231 + and callable(autotune)
232 + )
233 +
234 +
235 @functools.cache
236 def has_flashinfer_cutedsl() -> bool:
237 """Return ``True`` if FlashInfer cutedsl module is available."""
@@ -988,6 +1011,7 @@ __all__ = [
1011 "flashinfer_b12x_fused_moe",
1012 "flashinfer_convert_sf_to_mma_layout",
1013 "trtllm_fp4_block_scale_moe",
1014 + "flashinfer_trtllm_batch_decode_with_kv_cache_mla",
1015 "flashinfer_trtllm_batch_decode_sparse_mla_dsv4",
1016 "autotune",
1017 "has_flashinfer_moe",
vllm/v1/attention/backend.py
+1 -1
@@ -436,7 +436,7 @@ class CommonAttentionMetadata:
436 positions: torch.Tensor | None = None
437 """(num_actual_tokens,) token positions. Optional; set when the caller
438 has positions available so that builders can pre-compute position-dependent
439 - metadata (e.g. C128A topk indices for DeepSeek V4)."""
439 + sparse metadata for DeepSeek V4 C128A layers."""
440
441 is_prefilling: torch.Tensor | None = None
442 """(batch_size,) bool tensor: True if request is still in prefill phase
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+128 -37
@@ -1,23 +1,12 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""FlashInfer MLA Sparse Attention Backend.
4 -
5 -This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k
6 -for models like DeepSeek-V3.2 that use index-based sparse attention.
7 -
8 -For sparse MLA:
9 -- block_tables shape changes from [batch_size, max_num_blocks] (dense)
10 - to [batch_size, q_len_per_request, sparse_mla_top_k] (sparse)
11 -- The sparse indices represent physical cache slot positions to attend to
12 -- sparse_mla_top_k parameter must be set to the topk value
13 -"""
3 +"""FlashInfer sparse MLA attention backend."""
4
5 from dataclasses import dataclass
6 from typing import TYPE_CHECKING, ClassVar
7
8 import numpy as np
9 import torch
20 -from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
10
11 from vllm.config import VllmConfig
12 from vllm.config.cache import CacheDType
@@ -52,34 +41,13 @@ logger = init_logger(__name__)
41 FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
42
43
55 -class FlashInferMLASparseBackend(AttentionBackend):
56 - """FlashInfer MLA backend with sparse attention support.
57 -
58 - This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k
59 - for models like DeepSeek-V3.2 that use index-based sparse attention.
60 - """
61 -
62 - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
63 - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
64 - "auto",
65 - "float16",
66 - "bfloat16",
67 - "fp8",
68 - "fp8_e4m3",
69 - ]
70 -
71 - @staticmethod
72 - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
73 - return [32, 64]
44 +class _FlashInferMLASparseBackendBase(AttentionBackend):
45 + """Common metadata for concrete FlashInfer sparse MLA backends."""
46
47 @staticmethod
48 def get_name() -> str:
49 return "FLASHINFER_MLA_SPARSE"
50
79 - @staticmethod
80 - def get_impl_cls() -> type["FlashInferMLASparseImpl"]:
81 - return FlashInferMLASparseImpl
82 -
51 @staticmethod
52 def get_builder_cls() -> type["FlashInferMLASparseMetadataBuilder"]:
53 return FlashInferMLASparseMetadataBuilder
@@ -96,9 +64,29 @@ class FlashInferMLASparseBackend(AttentionBackend):
64 def is_sparse(cls) -> bool:
65 return True
66
67 +
68 +class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
69 + """FlashInfer sparse MLA backend using the TRTLLM-gen launcher."""
70 +
71 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
72 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
73 + "auto",
74 + "float16",
75 + "bfloat16",
76 + "fp8",
77 + "fp8_e4m3",
78 + ]
79 +
80 + @staticmethod
81 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
82 + return [32, 64]
83 +
84 + @staticmethod
85 + def get_impl_cls() -> type[SparseMLAAttentionImpl]:
86 + return FlashInferMLASparseImpl
87 +
88 @classmethod
89 def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
101 - # FlashInfer sparse MLA targets Blackwell (SM 10.x)
90 return capability.major == 10
91
92 @classmethod
@@ -114,10 +102,15 @@ class FlashInferMLASparseBackend(AttentionBackend):
102 use_mm_prefix: bool,
103 device_capability: DeviceCapability,
104 ) -> str | None:
117 - # FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192]
105 from vllm.config import get_current_vllm_config
106
107 vllm_config = get_current_vllm_config()
108 + if kv_cache_dtype == "fp8_ds_mla":
109 + return (
110 + "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype"
111 + )
112 +
113 + # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192].
114 if vllm_config.model_config is not None:
115 hf_text_config = vllm_config.model_config.hf_text_config
116 qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
@@ -146,6 +139,102 @@ class FlashInferMLASparseBackend(AttentionBackend):
139 return "HND"
140
141
142 +class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
143 + """FlashInfer sparse MLA backend for SM120."""
144 +
145 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
146 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
147 + "auto",
148 + "fp8",
149 + "fp8_e4m3",
150 + "fp8_ds_mla",
151 + ]
152 +
153 + @staticmethod
154 + def get_name() -> str:
155 + return "FLASHINFER_MLA_SPARSE_SM120"
156 +
157 + @staticmethod
158 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
159 + return [64, 256]
160 +
161 + @staticmethod
162 + def get_impl_cls() -> type[SparseMLAAttentionImpl]:
163 + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
164 + FlashInferMLASparseSM120Impl,
165 + )
166 +
167 + return FlashInferMLASparseSM120Impl
168 +
169 + @classmethod
170 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
171 + return capability.major == 12
172 +
173 + @classmethod
174 + def supports_combination(
175 + cls,
176 + head_size: int,
177 + dtype: torch.dtype,
178 + kv_cache_dtype: CacheDType | None,
179 + block_size: int | None,
180 + use_mla: bool,
181 + has_sink: bool,
182 + use_sparse: bool,
183 + use_mm_prefix: bool,
184 + device_capability: DeviceCapability,
185 + ) -> str | None:
186 + from vllm.config import get_current_vllm_config
187 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
188 +
189 + if not has_flashinfer_sparse_mla_sm120():
190 + return (
191 + "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
192 + "sparse MLA decode API"
193 + )
194 + if dtype != torch.bfloat16:
195 + return "dtype not supported"
196 + if kv_cache_dtype not in (
197 + None,
198 + "auto",
199 + "fp8",
200 + "fp8_e4m3",
201 + "fp8_ds_mla",
202 + ):
203 + return "kv_cache_dtype not supported"
204 + vllm_config = get_current_vllm_config()
205 + if vllm_config.model_config is not None:
206 + hf_text_config = vllm_config.model_config.hf_text_config
207 + index_topk = getattr(hf_text_config, "index_topk", None)
208 + if index_topk is None:
209 + return (
210 + "FLASHINFER_MLA_SPARSE_SM120 requires a model with "
211 + "index_topk config"
212 + )
213 + if int(index_topk) != 2048:
214 + return (
215 + "FLASHINFER_MLA_SPARSE_SM120 requires index_topk=2048; "
216 + f"got {index_topk}"
217 + )
218 + return None
219 +
220 + @staticmethod
221 + def get_kv_cache_shape(
222 + num_blocks: int,
223 + block_size: int,
224 + num_kv_heads: int, # assumed to be 1 for MLA
225 + head_size: int,
226 + cache_dtype_str: str = "auto",
227 + ) -> tuple[int, ...]:
228 + if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"):
229 + # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE.
230 + return (num_blocks, block_size, 656)
231 + return (num_blocks, block_size, head_size)
232 +
233 + @classmethod
234 + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
235 + return None
236 +
237 +
238 @dataclass
239 class FlashInferMLASparseMetadata(AttentionMetadata):
240 """Attention metadata for FlashInfer MLA Sparse backend."""
@@ -353,6 +442,8 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
442 if is_quantized_kv_cache(self.kv_cache_dtype):
443 self.bmm2_scale *= layer._k_scale_float
444
445 + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
446 +
447 o = trtllm_batch_decode_with_kv_cache_mla(
448 query=q.unsqueeze(1),
449 kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py new
+155
@@ -0,0 +1,155 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""SM120 implementation variant for ``FLASHINFER_MLA_SPARSE_SM120``."""
4 +
5 +from typing import TYPE_CHECKING, cast
6 +
7 +import torch
8 +
9 +from vllm.v1.attention.backend import (
10 + AttentionLayer,
11 + AttentionType,
12 + SparseMLAAttentionImpl,
13 +)
14 +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
15 + FlashInferMLASparseMetadata,
16 + _get_workspace_buffer,
17 +)
18 +from vllm.v1.attention.backends.mla.sparse_utils import (
19 + triton_convert_req_index_to_global_index,
20 +)
21 +
22 +if TYPE_CHECKING:
23 + from vllm.model_executor.models.deepseek_v2 import Indexer
24 +
25 +
26 +def _kv_scale_format_for_model(model_type: str | None) -> str:
27 + if model_type is not None and model_type.startswith("glm"):
28 + return "arbitrary_fp32"
29 + return "pow2_fp32"
30 +
31 +
32 +class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
33 + """SM120 FlashInfer sparse-MLA implementation."""
34 +
35 + def __init__(
36 + self,
37 + num_heads: int,
38 + head_size: int,
39 + scale: float,
40 + num_kv_heads: int,
41 + alibi_slopes: list[float] | None,
42 + sliding_window: int | None,
43 + kv_cache_dtype: str,
44 + logits_soft_cap: float | None,
45 + attn_type: str,
46 + kv_sharing_target_layer_name: str | None,
47 + indexer: "Indexer | None" = None,
48 + **mla_args,
49 + ) -> None:
50 + if any([alibi_slopes, sliding_window, logits_soft_cap]):
51 + raise NotImplementedError(
52 + "FLASHINFER_MLA_SPARSE_SM120 does not support alibi_slopes / "
53 + "sliding_window / logits_soft_cap"
54 + )
55 + if attn_type != AttentionType.DECODER:
56 + raise NotImplementedError(
57 + "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention"
58 + )
59 +
60 + self.num_heads = num_heads
61 + self.head_size = head_size
62 + self.scale = float(scale)
63 + self.num_kv_heads = num_kv_heads
64 + self.kv_cache_dtype = kv_cache_dtype
65 + if self.kv_cache_dtype != "fp8_ds_mla":
66 + raise NotImplementedError(
67 + "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla "
68 + f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}."
69 + )
70 +
71 + self.kv_lora_rank: int = mla_args["kv_lora_rank"]
72 + self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
73 + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
74 + from vllm.config import get_current_vllm_config
75 +
76 + vllm_config = get_current_vllm_config()
77 + model_type = None
78 + if vllm_config.model_config is not None:
79 + model_type = getattr(
80 + vllm_config.model_config.hf_text_config, "model_type", None
81 + )
82 + self.kv_scale_format = _kv_scale_format_for_model(model_type)
83 +
84 + assert indexer is not None, (
85 + "FLASHINFER_MLA_SPARSE_SM120 requires a sparse-MLA indexer "
86 + "(model with index_topk in its config)."
87 + )
88 + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer
89 + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120
90 +
91 + if not has_flashinfer_sparse_mla_sm120():
92 + raise RuntimeError(
93 + "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's "
94 + "sparse MLA decode API."
95 + )
96 + assert self.topk_indices_buffer is not None
97 +
98 + self.supports_quant_query_input = False
99 + self._workspace_buffer: torch.Tensor | None = None
100 +
101 + def forward_mqa(
102 + self,
103 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
104 + kv_c_and_k_pe_cache: torch.Tensor,
105 + attn_metadata: FlashInferMLASparseMetadata,
106 + layer: AttentionLayer,
107 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
108 + if isinstance(q, tuple):
109 + q = torch.cat(q, dim=-1)
110 +
111 + num_actual_toks = q.shape[0]
112 +
113 + assert self.topk_indices_buffer is not None
114 + topk_indices = self.topk_indices_buffer[:num_actual_toks]
115 +
116 + topk_indices_physical = cast(
117 + torch.Tensor,
118 + triton_convert_req_index_to_global_index(
119 + attn_metadata.req_id_per_token[:num_actual_toks],
120 + attn_metadata.block_table,
121 + topk_indices,
122 + BLOCK_SIZE=attn_metadata.block_size,
123 + NUM_TOPK_TOKENS=topk_indices.shape[1],
124 + ),
125 + )
126 +
127 + output = q.new_empty(
128 + (num_actual_toks, self.num_heads, self.kv_lora_rank),
129 + dtype=q.dtype,
130 + )
131 +
132 + if self._workspace_buffer is None:
133 + self._workspace_buffer = _get_workspace_buffer(q.device)
134 +
135 + from vllm.utils.flashinfer import (
136 + flashinfer_trtllm_batch_decode_with_kv_cache_mla,
137 + )
138 +
139 + out = flashinfer_trtllm_batch_decode_with_kv_cache_mla(
140 + query=q.unsqueeze(1),
141 + kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1),
142 + workspace_buffer=self._workspace_buffer,
143 + qk_nope_head_dim=self.qk_nope_head_dim,
144 + kv_lora_rank=self.kv_lora_rank,
145 + qk_rope_head_dim=self.qk_rope_head_dim,
146 + block_tables=topk_indices_physical.unsqueeze(1),
147 + seq_lens=None,
148 + max_seq_len=attn_metadata.topk_tokens,
149 + out=output.unsqueeze(1),
150 + bmm1_scale=self.scale,
151 + bmm2_scale=1.0,
152 + sparse_mla_top_k=attn_metadata.topk_tokens,
153 + kv_scale_format=self.kv_scale_format,
154 + )
155 + return out.squeeze(1), None
vllm/v1/attention/backends/mla/sparse_swa.py
+74 -15
@@ -74,14 +74,14 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
74 # determines the SWA block size of 64 tokens per block.
75 # TODO(yifan): make SWA block size automatically determined and configurable.
76 self.block_size = 64
77 - # uint8: legacy FlashMLA UE8M0 paged layout. bfloat16 / float8_e4m3fn:
78 - # FlashInfer contiguous full-cache layout.
77 + # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn:
78 + # contiguous full-cache layout.
79 assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
80
81 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
82 - # FlashMLA's UE8M0 paged layout needs 576B alignment; FlashInfer's
83 - # contiguous bf16/fp8 cache uses the natural element-size page.
84 - is_flashmla = self.cache_config.cache_dtype == "fp8_ds_mla"
82 + # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
83 + # bf16/fp8 cache uses the natural element-size page.
84 + uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
85 return SlidingWindowMLASpec(
86 block_size=self.block_size,
87 num_kv_heads=1,
@@ -89,7 +89,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
89 dtype=self.dtype,
90 sliding_window=self.window_size,
91 cache_dtype_str=self.cache_config.cache_dtype,
92 - alignment=576 if is_flashmla else None,
92 + alignment=576 if uses_fp8_ds_mla_layout else None,
93 model_version="deepseek_v4",
94 )
95
@@ -164,6 +164,11 @@ class DeepseekSparseSWAMetadata:
164 token_to_req_indices: torch.Tensor | None = None # [num_tokens]
165 decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, window_size]
166 decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens]
167 + # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill).
168 + prefill_swa_indices: torch.Tensor | None = (
169 + None # [num_prefill_tokens, 1, window_size]
170 + )
171 + prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens]
172
173 # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
174 num_decodes: int = 0
@@ -343,6 +348,20 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
348 dtype=torch.int32,
349 device=self.device,
350 )
351 + # Allocated unconditionally — consumer picks paged-direct vs dequant
352 + # at call time.
353 + self.prefill_swa_indices = torch.zeros(
354 + max_tokens,
355 + 1,
356 + self.window_size,
357 + dtype=torch.int32,
358 + device=self.device,
359 + )
360 + self.prefill_swa_lens = torch.zeros(
361 + max_tokens,
362 + dtype=torch.int32,
363 + device=self.device,
364 + )
365 self.is_valid_token = torch.zeros(
366 max_tokens,
367 dtype=torch.bool,
@@ -402,6 +421,29 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
421 block_table,
422 block_table.stride(0),
423 self.block_size,
424 + token_offset=0,
425 + TRITON_BLOCK_SIZE=1024,
426 + )
427 +
428 + # Prefill SWA indices live in paged coordinates. `token_offset` lets
429 + # the kernel read is_valid_token / token_to_req_indices at absolute
430 + # prefill positions while writing output starting at index 0.
431 + if num_prefill_tokens > 0:
432 + prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
433 + prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
434 + _compute_swa_indices_and_lens_kernel[(num_prefill_tokens,)](
435 + prefill_swa_indices,
436 + prefill_swa_indices.stride(0),
437 + prefill_swa_lens,
438 + self.window_size,
439 + query_start_loc,
440 + seq_lens,
441 + token_to_req_indices,
442 + is_valid_token,
443 + block_table,
444 + block_table.stride(0),
445 + self.block_size,
446 + token_offset=num_decode_tokens,
447 TRITON_BLOCK_SIZE=1024,
448 )
449
@@ -431,6 +473,16 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
473 token_to_req_indices=token_to_req_indices,
474 decode_swa_indices=self.decode_swa_indices[:num_decode_tokens],
475 decode_swa_lens=self.decode_swa_lens[:num_decode_tokens],
476 + prefill_swa_indices=(
477 + self.prefill_swa_indices[:num_prefill_tokens]
478 + if num_prefill_tokens > 0
479 + else None
480 + ),
481 + prefill_swa_lens=(
482 + self.prefill_swa_lens[:num_prefill_tokens]
483 + if num_prefill_tokens > 0
484 + else None
485 + ),
486 block_size=self.block_size,
487 num_decodes=num_decodes,
488 num_prefills=num_prefills,
@@ -465,6 +517,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
517 num_decode_tokens == 0
518 or current_platform.is_rocm()
519 or current_platform.is_xpu()
520 + or current_platform.is_device_capability_family(120)
521 ):
522 return out
523 for layer_type in self._layer_types:
@@ -489,7 +542,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
542 Returns a dict of keyword arguments to pass to the
543 DeepseekSparseSWAMetadata constructor.
544
492 - Note: C128A topk indices are computed by the FlashMLASparse builder
545 + Note: C128A sparse metadata is computed by the FlashMLASparse builder
546 (which owns the C128A block_table), not here.
547 """
548 result: dict[str, torch.Tensor | int | None] = {}
@@ -539,10 +592,14 @@ def _compute_prefill_metadata_kernel(
592 """Compute prefill gather_lens in a single pass."""
593 offset = tl.arange(0, BLOCK_SIZE)
594 mask = offset < num_prefills
595 + # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
596 + # masked-off lanes even though the load mask gates the actual read, so
597 + # clamp the offset. Caller guarantees num_prefills > 0.
598 + safe_offset = tl.minimum(offset, num_prefills - 1)
599
543 - seq_len = tl.load(seq_lens_ptr + num_decodes + offset, mask=mask)
544 - qsl_start = tl.load(query_start_loc_ptr + num_decodes + offset, mask=mask)
545 - qsl_end = tl.load(query_start_loc_ptr + num_decodes + offset + 1, mask=mask)
600 + seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
601 + qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
602 + qsl_end = tl.load(query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask)
603
604 query_len = qsl_end - qsl_start
605 prefix_len = seq_len - query_len
@@ -551,7 +608,7 @@ def _compute_prefill_metadata_kernel(
608 tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
609
610
554 -@triton.jit
611 +@triton.jit(do_not_specialize=["token_offset"])
612 def _compute_swa_indices_and_lens_kernel(
613 swa_indices_ptr,
614 swa_indices_stride,
@@ -564,12 +621,14 @@ def _compute_swa_indices_and_lens_kernel(
621 block_table_ptr,
622 block_table_stride,
623 block_size,
624 + token_offset,
625 TRITON_BLOCK_SIZE: tl.constexpr,
626 ):
569 - token_idx = tl.program_id(0)
627 + pid = tl.program_id(0)
628 + token_idx = pid + token_offset
629 is_valid = tl.load(is_valid_token_ptr + token_idx)
630 if not is_valid:
572 - tl.store(swa_lens_ptr + token_idx, 0)
631 + tl.store(swa_lens_ptr + pid, 0)
632 return
633
634 req_idx = tl.load(token_to_req_indices_ptr + token_idx)
@@ -586,7 +645,7 @@ def _compute_swa_indices_and_lens_kernel(
645 end_pos = pos + 1
646
647 swa_len = end_pos - start_pos
589 - tl.store(swa_lens_ptr + token_idx, swa_len)
648 + tl.store(swa_lens_ptr + pid, swa_len)
649
650 for i in range(0, window_size, TRITON_BLOCK_SIZE):
651 offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
@@ -602,7 +661,7 @@ def _compute_swa_indices_and_lens_kernel(
661
662 slot_ids = tl.where(offset < swa_len, slot_ids, -1)
663 tl.store(
605 - swa_indices_ptr + token_idx * swa_indices_stride + offset,
664 + swa_indices_ptr + pid * swa_indices_stride + offset,
665 slot_ids,
666 mask=offset < window_size,
667 )
vllm/v1/attention/backends/registry.py
+5 -1
@@ -71,7 +71,11 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
71 )
72 FLASHINFER_MLA_SPARSE = (
73 "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
74 - "FlashInferMLASparseBackend"
74 + "FlashInferMLASparseTRTLLMBackend"
75 + )
76 + FLASHINFER_MLA_SPARSE_SM120 = (
77 + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."
78 + "FlashInferMLASparseSM120Backend"
79 )
80 TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
81 CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"
vllm/v1/attention/ops/flashmla.py
+1 -1
@@ -73,7 +73,7 @@ def is_flashmla_sparse_supported() -> tuple[bool, str | None]:
73 ):
74 return (
75 False,
76 - "FlashMLA Sparse is only supported on Hopper and Blackwell devices.",
76 + "FlashMLA Sparse is only supported on Hopper and Blackwell DC devices.",
77 )
78 return True, None
79
vllm/v1/worker/gpu/warmup.py
+130
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 +from contextlib import AbstractContextManager, nullcontext
6 from typing import Any
7
8 import numpy as np
9 import torch
10
11 from vllm import PoolingParams, SamplingParams
12 +from vllm.logger import init_logger
13 from vllm.utils.math_utils import cdiv
14 from vllm.v1.core.sched.output import (
15 CachedRequestData,
@@ -18,6 +20,134 @@ from vllm.v1.core.sched.output import (
20 from vllm.v1.request import Request
21 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
22
23 +logger = init_logger(__name__)
24 +
25 +
26 +def run_mixed_prefill_decode_warmup(
27 + model_runner: GPUModelRunner,
28 + worker_execute_model: Callable[[SchedulerOutput], Any],
29 + worker_sample_tokens: Callable[[GrammarOutput | None], Any],
30 + num_tokens: int,
31 + *,
32 + mixed_step_context: AbstractContextManager[object] | None = None,
33 + req_id_prefix: str = "_v2_mixed_warmup",
34 +) -> bool:
35 + """Run a V2 mixed prefill+decode step through normal scheduler inputs."""
36 + if model_runner.is_pooling_model or num_tokens < 3:
37 + return False
38 +
39 + decode_req_id = f"{req_id_prefix}_decode_"
40 + prefill_req_id = f"{req_id_prefix}_prefill_"
41 + decode_prompt_len = 2
42 + decode_scheduled_tokens = 1
43 + prefill_len = num_tokens - decode_scheduled_tokens
44 + decode_token_ids = list(range(decode_prompt_len))
45 + prefill_token_ids = list(range(prefill_len))
46 +
47 + kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
48 + num_kv_cache_groups = len(kv_cache_groups)
49 + group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups]
50 + decode_prefill_block_counts = [
51 + cdiv(decode_prompt_len, block_size) for block_size in group_block_sizes
52 + ]
53 + decode_block_counts = [
54 + cdiv(decode_prompt_len + decode_scheduled_tokens, block_size)
55 + for block_size in group_block_sizes
56 + ]
57 + decode_block_deltas = [
58 + decode - prefill
59 + for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts)
60 + ]
61 + prefill_block_counts = [
62 + cdiv(prefill_len, block_size) for block_size in group_block_sizes
63 + ]
64 + required_blocks = sum(decode_block_counts) + sum(prefill_block_counts)
65 + if model_runner.kv_cache_config.num_blocks <= required_blocks:
66 + logger.warning(
67 + "Skipping V2 mixed prefill+decode warmup because only %d KV blocks "
68 + "are available for %d required warmup blocks.",
69 + model_runner.kv_cache_config.num_blocks,
70 + required_blocks,
71 + )
72 + return False
73 +
74 + next_block_id = 1
75 +
76 + def _alloc_blocks(num_blocks: int) -> list[int]:
77 + nonlocal next_block_id
78 + block_ids = list(range(next_block_id, next_block_id + num_blocks))
79 + next_block_id += num_blocks
80 + return block_ids
81 +
82 + sampling_params = SamplingParams(max_tokens=2, temperature=0.0)
83 +
84 + decode_prefill_output = SchedulerOutput.make_empty()
85 + decode_prefill_output.scheduled_new_reqs = [
86 + NewRequestData(
87 + req_id=decode_req_id,
88 + prompt_token_ids=decode_token_ids,
89 + mm_features=[],
90 + sampling_params=sampling_params,
91 + pooling_params=None,
92 + block_ids=tuple(_alloc_blocks(n) for n in decode_prefill_block_counts),
93 + num_computed_tokens=0,
94 + lora_request=None,
95 + prefill_token_ids=decode_token_ids,
96 + ),
97 + ]
98 + decode_prefill_output.num_scheduled_tokens = {
99 + decode_req_id: decode_prompt_len,
100 + }
101 + decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len
102 + decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
103 +
104 + decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas)
105 + cached_decode_req = CachedRequestData.make_empty()
106 + cached_decode_req.req_ids = [decode_req_id]
107 + cached_decode_req.num_computed_tokens = [decode_prompt_len]
108 + cached_decode_req.num_output_tokens = [1]
109 + cached_decode_req.new_block_ids = [
110 + decode_new_blocks if any(decode_block_deltas) else None
111 + ]
112 +
113 + mixed_output = SchedulerOutput.make_empty()
114 + mixed_output.scheduled_cached_reqs = cached_decode_req
115 + mixed_output.scheduled_new_reqs = [
116 + NewRequestData(
117 + req_id=prefill_req_id,
118 + prompt_token_ids=prefill_token_ids,
119 + mm_features=[],
120 + sampling_params=sampling_params,
121 + pooling_params=None,
122 + block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
123 + num_computed_tokens=0,
124 + lora_request=None,
125 + prefill_token_ids=prefill_token_ids,
126 + ),
127 + ]
128 + mixed_output.num_scheduled_tokens = {
129 + decode_req_id: decode_scheduled_tokens,
130 + prefill_req_id: prefill_len,
131 + }
132 + mixed_output.total_num_scheduled_tokens = num_tokens
133 + mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
134 +
135 + cleanup_output = SchedulerOutput.make_empty()
136 + cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id}
137 +
138 + context = mixed_step_context or nullcontext()
139 + model_runner.kv_connector.set_disabled(True)
140 + try:
141 + worker_execute_model(decode_prefill_output)
142 + worker_sample_tokens(None)
143 + with context:
144 + worker_execute_model(mixed_output)
145 + worker_sample_tokens(None)
146 + worker_execute_model(cleanup_output)
147 + finally:
148 + model_runner.kv_connector.set_disabled(False)
149 + return True
150 +
151
152 @torch.inference_mode()
153 def warmup_kernels(