diff --git a/.buildkite/intel_jobs/expert_parallelism_intel.yaml b/.buildkite/intel_jobs/expert_parallelism_intel.yaml index 6c81fb04c0..24dfb07f5f 100644 --- a/.buildkite/intel_jobs/expert_parallelism_intel.yaml +++ b/.buildkite/intel_jobs/expert_parallelism_intel.yaml @@ -1,5 +1,5 @@ group: Expert Parallelism -depends_on: +depends_on: - image-build-xpu steps: - label: EPLB Algorithm diff --git a/.buildkite/intel_jobs/models_multimodal_intel.yaml b/.buildkite/intel_jobs/models_multimodal_intel.yaml index f29f142516..0e12690604 100644 --- a/.buildkite/intel_jobs/models_multimodal_intel.yaml +++ b/.buildkite/intel_jobs/models_multimodal_intel.yaml @@ -1,5 +1,5 @@ group: Models - Multimodal -depends_on: +depends_on: - image-build-xpu steps: - label: "Multi-Modal Models (Standard) 1: qwen2" diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 183c42dc79..dc2a61bc2b 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -8,43 +8,73 @@ if (DEFINED ENV{DEEPGEMM_SRC_DIR}) set(DEEPGEMM_SRC_DIR $ENV{DEEPGEMM_SRC_DIR}) endif() +# Local tree: set deepgemm_SOURCE_DIR directly (no FetchContent download). +# Upstream git: use FetchContent_Populate with explicit options (CMP0169 NEW +# disallows one-argument Populate(dep) after Declare; MakeAvailable would run +# DeepGEMM's top-level CMakeLists.txt, which vLLM must not load). if(DEEPGEMM_SRC_DIR) - FetchContent_Declare( - deepgemm - SOURCE_DIR ${DEEPGEMM_SRC_DIR} - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) + # cmake_path(ABSOLUTE_PATH ...) reads the path from ; NORMALIZE is a + # flag (no trailing path argument). Resolve relative paths against vLLM root. + set(_deepgemm_user_src "${DEEPGEMM_SRC_DIR}") + cmake_path(ABSOLUTE_PATH _deepgemm_user_src + BASE_DIRECTORY "${CMAKE_SOURCE_DIR}" + NORMALIZE) + set(DEEPGEMM_SRC_DIR "${_deepgemm_user_src}") + if(NOT IS_DIRECTORY "${DEEPGEMM_SRC_DIR}") + message(FATAL_ERROR + "DEEPGEMM_SRC_DIR is not an existing directory: '${DEEPGEMM_SRC_DIR}'") + endif() + set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}") + message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() - # This ref should be kept in sync with tools/install_deepgemm.sh - FetchContent_Declare( - deepgemm - GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git - GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f - GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" - GIT_PROGRESS TRUE - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) + # Keep in sync with tools/install_deepgemm.sh + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") + set(_DEEPGEMM_UPSTREAM_TAG "891d57b4db1071624b5c8fa0d1e51cb317fa709f") + + set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") + if(NOT _deepgemm_fc_root) + set(_deepgemm_fc_root "${CMAKE_BINARY_DIR}/_deps") + endif() + set(_deepgemm_src "${_deepgemm_fc_root}/deepgemm-src") + set(_deepgemm_bin "${_deepgemm_fc_root}/deepgemm-build") + set(_deepgemm_sub "${_deepgemm_fc_root}/deepgemm-subbuild") + + if(EXISTS "${_deepgemm_src}/csrc/python_api.cpp") + set(deepgemm_SOURCE_DIR "${_deepgemm_src}") + set(deepgemm_BINARY_DIR "${_deepgemm_bin}") + else() + FetchContent_Populate( + deepgemm + SUBBUILD_DIR "${_deepgemm_sub}" + SOURCE_DIR "${_deepgemm_src}" + BINARY_DIR "${_deepgemm_bin}" + GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}" + GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}" + GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_PROGRESS TRUE + ) + endif() + message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}") endif() -# Use FetchContent_Populate (not MakeAvailable) to avoid processing -# DeepGEMM's own CMakeLists.txt which has incompatible find_package calls. -FetchContent_GetProperties(deepgemm) -if(NOT deepgemm_POPULATED) - FetchContent_Populate(deepgemm) -endif() -message(STATUS "DeepGEMM is available at ${deepgemm_SOURCE_DIR}") - -# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 +# DeepGEMM requires CUDA 12.3+ for SM90, 12.9+ for SM100 (official upstream), +# and 12.8+ for SM120 / SM12x. CUDA 13+ can use the family-specific SM12x +# arch; CUDA 12.x builds the arch-specific SM120/SM121 variants. set(DEEPGEMM_SUPPORT_ARCHS) if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3) list(APPEND DEEPGEMM_SUPPORT_ARCHS "9.0a") endif() -if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) - list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") -elseif(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) - list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8) + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0f") + else() + list(APPEND DEEPGEMM_SUPPORT_ARCHS "10.0a") + endif() + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0) + list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0f") + else() + list(APPEND DEEPGEMM_SUPPORT_ARCHS "12.0a" "12.1a") + endif() endif() cuda_archs_loose_intersection(DEEPGEMM_ARCHS diff --git a/cmake/external_projects/qutlass.cmake b/cmake/external_projects/qutlass.cmake index b653bbfce7..29c5c6528b 100644 --- a/cmake/external_projects/qutlass.cmake +++ b/cmake/external_projects/qutlass.cmake @@ -6,25 +6,47 @@ if(DEFINED ENV{QUTLASS_SRC_DIR}) set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR}) endif() +# CMP0169 NEW: one-argument FetchContent_Populate(name) after Declare is invalid. +# Use explicit Populate(...) for git, or set SOURCE_DIR for local trees. if(QUTLASS_SRC_DIR) - FetchContent_Declare( - qutlass - SOURCE_DIR ${QUTLASS_SRC_DIR} - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) + set(_qutlass_user_src "${QUTLASS_SRC_DIR}") + cmake_path(ABSOLUTE_PATH _qutlass_user_src + BASE_DIRECTORY "${CMAKE_SOURCE_DIR}" + NORMALIZE) + set(QUTLASS_SRC_DIR "${_qutlass_user_src}") + if(NOT IS_DIRECTORY "${QUTLASS_SRC_DIR}") + message(FATAL_ERROR + "[QUTLASS] QUTLASS_SRC_DIR is not an existing directory: '${QUTLASS_SRC_DIR}'") + endif() + set(qutlass_SOURCE_DIR "${QUTLASS_SRC_DIR}") + set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused") else() - FetchContent_Declare( - qutlass - GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git - GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65 - GIT_PROGRESS TRUE - CONFIGURE_COMMAND "" - BUILD_COMMAND "" - ) -endif() + set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git") + set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65") -FetchContent_Populate(qutlass) + set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}") + if(NOT _qutlass_fc_root) + set(_qutlass_fc_root "${CMAKE_BINARY_DIR}/_deps") + endif() + set(_qutlass_src "${_qutlass_fc_root}/qutlass-src") + set(_qutlass_bin "${_qutlass_fc_root}/qutlass-build") + set(_qutlass_sub "${_qutlass_fc_root}/qutlass-subbuild") + + if(EXISTS "${_qutlass_src}/qutlass/csrc/bindings.cpp") + set(qutlass_SOURCE_DIR "${_qutlass_src}") + set(qutlass_BINARY_DIR "${_qutlass_bin}") + else() + FetchContent_Populate( + qutlass + SUBBUILD_DIR "${_qutlass_sub}" + SOURCE_DIR "${_qutlass_src}" + BINARY_DIR "${_qutlass_bin}" + GIT_REPOSITORY "${_QUTLASS_UPSTREAM_REPO}" + GIT_TAG "${_QUTLASS_UPSTREAM_TAG}" + GIT_PROGRESS TRUE + ) + endif() +endif() if(NOT qutlass_SOURCE_DIR) message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.") diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index 2dff668a01..6ac2a2c863 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -133,16 +133,6 @@ Priority is **1 = highest** (tried first). | 7 | `FLASHINFER_MLA_SPARSE`**\*** | | 8 | `FLASHMLA_SPARSE` | -**Ampere/Hopper (SM 8.x-9.x):** - -| Priority | Backend | -| -------- | ------- | -| 1 | `FLASH_ATTN_MLA` | -| 2 | `FLASHMLA` | -| 3 | `FLASHINFER_MLA` | -| 4 | `TRITON_MLA` | -| 5 | `FLASHMLA_SPARSE` | - > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise. > > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details. @@ -231,7 +221,8 @@ MLA decode backends are selected using the standard | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE_SM120` | bf16 | `auto`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 64, 256 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 12.x | | `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | @@ -248,10 +239,11 @@ DeepSeek V4 sparse MLA uses its own decode backends, selected via `--attention-backend=` (e.g., `FLASHMLA_SPARSE_DSV4`, `FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index pipeline (compressor + SWA + indexer, 256-token blocks, head 512); -default on NVIDIA is `FLASHMLA_SPARSE_DSV4`. +default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and +`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures. | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ | -| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | Any | -| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | +| `FLASHINFER_MLA_SPARSE_DSV4` | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_ds_mla` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x, 12.x | +| `FLASHMLA_SPARSE_DSV4` | bf16 | `auto`, `fp8_ds_mla`, `fp8` | 256 | 512 | ✅ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | | `ROCM_FLASHMLA_SPARSE_DSV4` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | diff --git a/tests/model_executor/test_flashinfer_autotune_cache.py b/tests/model_executor/test_flashinfer_autotune_cache.py deleted file mode 100644 index 7e6a83bb4d..0000000000 --- a/tests/model_executor/test_flashinfer_autotune_cache.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import sys -from hashlib import sha256 -from pathlib import Path -from types import SimpleNamespace - -from vllm.model_executor.warmup import kernel_warmup - - -def test_resolve_flashinfer_autotune_file_default_layout( - monkeypatch, tmp_path: Path -) -> None: - fake_jit = SimpleNamespace( - env=SimpleNamespace( - FLASHINFER_WORKSPACE_DIR=Path("/flashinfer-cache/0.6.11.post2/103a") - ) - ) - fake_flashinfer = SimpleNamespace(jit=fake_jit) - monkeypatch.setitem(sys.modules, "flashinfer", fake_flashinfer) - monkeypatch.setitem(sys.modules, "flashinfer.jit", fake_jit) - monkeypatch.setattr( - kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"] - ) - monkeypatch.setattr(kernel_warmup.envs, "VLLM_CACHE_ROOT", str(tmp_path)) - monkeypatch.setattr(kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", None) - - runner = SimpleNamespace(vllm_config=SimpleNamespace()) - cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest() - - path = kernel_warmup._resolve_flashinfer_autotune_file(runner) - - assert path == ( - tmp_path - / "flashinfer_autotune_cache" - / "0.6.11.post2" - / "103a" - / cache_hash - / "autotune_configs.json" - ) - assert path.parent.is_dir() - - -def test_resolve_flashinfer_autotune_file_uses_override_dir( - monkeypatch, tmp_path: Path -) -> None: - monkeypatch.setattr( - kernel_warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR", str(tmp_path) - ) - monkeypatch.setattr( - kernel_warmup, "aot_compile_hash_factors", lambda _: ["env-hash", "config-hash"] - ) - - runner = SimpleNamespace(vllm_config=SimpleNamespace()) - cache_hash = sha256(str(["env-hash", "config-hash"]).encode()).hexdigest() - - path = kernel_warmup._resolve_flashinfer_autotune_file(runner) - - assert path == tmp_path / cache_hash / "autotune_configs.json" diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py new file mode 100644 index 0000000000..3a7677e751 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Behavior checks for FlashInfer SM120 sparse MLA backend selection.""" + +from types import SimpleNamespace + +import torch + +from vllm.config import set_current_vllm_config +from vllm.platforms.interface import DeviceCapability +from vllm.utils import flashinfer as fi_utils +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseSM120Backend, +) +from vllm.v1.attention.backends.registry import AttentionBackendEnum + + +def _fake_vllm_config(model_type: str) -> SimpleNamespace: + return SimpleNamespace( + model_config=SimpleNamespace( + hf_text_config=SimpleNamespace(model_type=model_type, index_topk=2048), + ), + ) + + +def test_sm120_backend_uses_dedicated_backend_name() -> None: + assert FlashInferMLASparseSM120Backend.get_name() == "FLASHINFER_MLA_SPARSE_SM120" + assert ( + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120.get_class() + is FlashInferMLASparseSM120Backend + ) + + +def test_v32_glm_sm120_backend_accepts_glm_block_size( + monkeypatch, +) -> None: + monkeypatch.setattr(fi_utils, "has_flashinfer_sparse_mla_sm120", lambda: True) + + with set_current_vllm_config(_fake_vllm_config("glm4_moe")): + invalid_reasons = FlashInferMLASparseSM120Backend.validate_configuration( + head_size=576, + dtype=torch.bfloat16, + kv_cache_dtype="fp8", + block_size=256, + use_mla=True, + has_sink=False, + use_sparse=True, + use_mm_prefix=False, + use_per_head_quant_scales=False, + device_capability=DeviceCapability(12, 0), + attn_type="decoder", + ) + + assert invalid_reasons == [] diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 22acc748d2..6e389604c8 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -36,7 +36,7 @@ if not current_platform.is_cuda(): from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( - FlashInferMLASparseBackend, + FlashInferMLASparseTRTLLMBackend, ) from vllm.v1.attention.backends.mla.flashmla_sparse import ( FlashMLASparseBackend, @@ -174,8 +174,8 @@ def _quantize_dequantize_fp8_ds_mla( @pytest.mark.parametrize( "backend_cls", - [FlashMLASparseBackend, FlashInferMLASparseBackend], - ids=["FlashMLA", "FlashInfer"], + [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend], + ids=["FlashMLA", "FlashInferTRTLLM"], ) @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys())) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"]) @@ -217,9 +217,12 @@ def test_sparse_backend_decode_correctness( ok, reason = flashmla.is_flashmla_sparse_supported() if not ok: pytest.skip(reason) - elif backend_cls == FlashInferMLASparseBackend: - if not current_platform.has_device_capability(100): - pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher") + elif backend_cls == FlashInferMLASparseTRTLLMBackend: + device_capability = current_platform.get_device_capability() + if device_capability is None or not backend_cls.supports_compute_capability( + device_capability + ): + pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability") batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name] use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla" diff --git a/tests/v1/spec_decode/test_acceptance_length.py b/tests/v1/spec_decode/test_acceptance_length.py index 90e3821e2f..4d093f6055 100644 --- a/tests/v1/spec_decode/test_acceptance_length.py +++ b/tests/v1/spec_decode/test_acceptance_length.py @@ -132,9 +132,9 @@ def get_available_attention_backends() -> list[str]: ) return [ - backend.name - for backend, _ in valid_backends - if backend not in EXCLUDED_BACKENDS + candidate.backend.name + for candidate in valid_backends + if candidate.backend not in EXCLUDED_BACKENDS ] diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 91720911a6..f8b87fa860 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -690,7 +690,9 @@ def parse_compute_capability(node: ast.ClassDef) -> str: major_list.sort() if len(major_list) == 1: return f"{major_list[0]}.x" - return f"{major_list[0]}.x-{major_list[-1]}.x" + if major_list == list(range(major_list[0], major_list[-1] + 1)): + return f"{major_list[0]}.x-{major_list[-1]}.x" + return ", ".join(f"{major}.x" for major in major_list) if min_cap: if max_cap: @@ -1668,7 +1670,8 @@ def generate_mla_section( "`--attention-backend=` (e.g., `FLASHMLA_SPARSE_DSV4`,", "`FLASHINFER_MLA_SPARSE_DSV4`). They share the V4 sparse-index", "pipeline (compressor + SWA + indexer, 256-token blocks, head 512);", - "default on NVIDIA is `FLASHMLA_SPARSE_DSV4`.", + "default on NVIDIA is `FLASHINFER_MLA_SPARSE_DSV4` on SM12x and", + "`FLASHMLA_SPARSE_DSV4` on other supported CUDA architectures.", "", ] ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 3344f6d008..051468ed14 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -208,6 +208,7 @@ from vllm.config import ( get_current_vllm_config, get_current_vllm_config_or_none, ) +from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import ( get_dcp_group, is_global_first_rank, @@ -319,6 +320,22 @@ def _detect_output_quant_key( return kFp8StaticTensorSym +def _canonicalize_sparse_mla_kv_cache_dtype( + attn_backend: type[AttentionBackend], + kv_cache_dtype: CacheDType, +) -> CacheDType: + backend_name = attn_backend.get_name() + if backend_name == "FLASHMLA_SPARSE" and is_quantized_kv_cache(kv_cache_dtype): + return "fp8_ds_mla" + if backend_name == "FLASHINFER_MLA_SPARSE_SM120" and kv_cache_dtype in ( + "auto", + "fp8", + "fp8_e4m3", + ): + return "fp8_ds_mla" + return kv_cache_dtype + + class MLAAttention(nn.Module, AttentionLayerBase): """Multi-Head Latent Attention layer. @@ -369,7 +386,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim if cache_config is not None: - kv_cache_dtype = cache_config.cache_dtype + kv_cache_dtype: CacheDType = cache_config.cache_dtype calculate_kv_scales = cache_config.calculate_kv_scales else: kv_cache_dtype = "auto" @@ -393,24 +410,22 @@ class MLAAttention(nn.Module, AttentionLayerBase): num_heads=self.num_heads, ) - # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format - # Automatically convert fp8 kv-cache format to "fp8_ds_mla" - if ( - self.attn_backend.get_name() == "FLASHMLA_SPARSE" - and is_quantized_kv_cache(kv_cache_dtype) - and kv_cache_dtype != "fp8_ds_mla" - ): - assert cache_config is not None - cache_config.cache_dtype = "fp8_ds_mla" - kv_cache_dtype = "fp8_ds_mla" + normalized_kv_cache_dtype = _canonicalize_sparse_mla_kv_cache_dtype( + self.attn_backend, kv_cache_dtype + ) + if normalized_kv_cache_dtype != kv_cache_dtype: + if cache_config is not None: + cache_config.cache_dtype = normalized_kv_cache_dtype + kv_cache_dtype = normalized_kv_cache_dtype logger.info_once( - "Using DeepSeek's fp8_ds_mla KV cache format. To use standard " - "fp8 kv-cache format, please set `--attention-backend " - "FLASHINFER_MLA_SPARSE`" + "Using %s KV cache format for %s backend.", + kv_cache_dtype, + self.attn_backend.get_name(), ) if ( self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE" + and kv_cache_dtype != "fp8_ds_mla" and is_quantized_kv_cache(kv_cache_dtype) ): logger.info_once( diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index c74cb2d9a7..f4319b99c8 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -23,24 +23,82 @@ def expert_num_tokens_round_up_and_sum( return torch.sum(ent).item() +def compute_aligned_M_and_alignment( + M: int, + num_topk: int, + local_num_experts: int, + alignment: int, + expert_tokens_meta: mk.ExpertTokensMetadata | None, +) -> tuple[int, int]: + """Return (M_sum, alignment_used). + + `alignment_used` may be smaller than the caller-supplied `alignment` on + SM100/SM120 when DeepGEMM can JIT a smaller BLOCK_M for the per-call + expected_m. Callers that index by block size (e.g. ``M_sum // block_m``) + or assert workspace alignment must use the returned `alignment_used`, + not their original `alignment` argument. + + Prefer this over the int-returning :func:`compute_aligned_M` when the + GEMM call site needs to wrap itself in ``mk_alignment_scope`` or + otherwise reason about the actual per-expert padding. + """ + if (expert_tokens_meta is not None) and ( + expert_tokens_meta.expert_num_tokens_cpu is not None + ): + return ( + expert_num_tokens_round_up_and_sum( + expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment + ), + alignment, + ) + + # expert_num_tokens not on cpu. Cap padding by min(M*num_topk, + # local_num_experts) — at batch=1 decode only `num_topk` experts can be + # active, so the worst-case `local_num_experts*(align-1)` is too loose. + # Also shrink `alignment` to DeepGEMM's per-call theoretical BLOCK_M on + # SM100/SM120 when smaller. + expected_m = M * num_topk + try: + from vllm.utils.deep_gemm import ( + get_theoretical_mk_alignment_for_contiguous_layout, + ) + + # num_groups=local_num_experts so the helper recovers per-expert em; + # omitting it over-picks BLOCK_M on SM120 (heuristic assumes em is + # already per-expert). + per_call_align = get_theoretical_mk_alignment_for_contiguous_layout( + expected_m=expected_m, + num_groups=local_num_experts, + ) + if per_call_align and per_call_align <= alignment: + alignment = per_call_align + except Exception: + pass + + max_active_experts = min(M * num_topk, local_num_experts) + M_sum = (M * num_topk) + max_active_experts * (alignment - 1) + M_sum = round_up(M_sum, alignment) + return M_sum, alignment + + def compute_aligned_M( M: int, num_topk: int, local_num_experts: int, alignment: int, expert_tokens_meta: mk.ExpertTokensMetadata | None, -): - if (expert_tokens_meta is not None) and ( - expert_tokens_meta.expert_num_tokens_cpu is not None - ): - return expert_num_tokens_round_up_and_sum( - expert_tokens_meta.expert_num_tokens_cpu, alignment=alignment - ) +) -> int: + """Return ``M_sum`` only (backward-compat wrapper). - # expert_num_tokens information is not available on the cpu. - # compute the max required size. - M_sum = (M * num_topk) + local_num_experts * (alignment - 1) - M_sum = round_up(M_sum, alignment) + Equivalent to :func:`compute_aligned_M_and_alignment`'s first return + value. Existing downstream callers and the warmup path that only size + a workspace use this. Call sites that need the actual per-expert + alignment (to wrap GEMMs in ``mk_alignment_scope``) should use + :func:`compute_aligned_M_and_alignment` instead. + """ + M_sum, _ = compute_aligned_M_and_alignment( + M, num_topk, local_num_experts, alignment, expert_tokens_meta + ) return M_sum @@ -51,12 +109,6 @@ def apply_expert_map(expert_id, expert_map): return expert_id -@triton.jit -def round_up_128(x: int) -> int: - y = 128 - return ((x + y - 1) // y) * y - - @triton.jit def _fwd_kernel_ep_scatter_1( num_recv_tokens_per_expert, @@ -65,6 +117,7 @@ def _fwd_kernel_ep_scatter_1( num_experts: tl.constexpr, BLOCK_E: tl.constexpr, BLOCK_EXPERT_NUM: tl.constexpr, + ALIGN_M: tl.constexpr, ): cur_expert = tl.program_id(0) @@ -74,7 +127,8 @@ def _fwd_kernel_ep_scatter_1( mask=offset_cumsum < num_experts, other=0, ) - tokens_per_expert = round_up_128(tokens_per_expert) + # Round up to ALIGN_M so cumsum matches the workspace's per-expert slices. + tokens_per_expert = ((tokens_per_expert + ALIGN_M - 1) // ALIGN_M) * ALIGN_M cumsum = tl.cumsum(tokens_per_expert) - tokens_per_expert # Extract this block's offset from the register vector (warp shuffle, @@ -227,10 +281,12 @@ def ep_scatter( output_tensor_scale: torch.Tensor, m_indices: torch.Tensor, output_index: torch.Tensor, + align_m: int = 128, block_size: int = 128, pack_ue8m0: bool = False, ): - BLOCK_E = 128 # token num of per expert is aligned to 128 + # BLOCK_E is the m_indices fill-loop tile (masked), independent of align_m. + BLOCK_E = 128 BLOCK_D = block_size # block size of activation-scale quantization num_warps = 8 num_experts = num_recv_tokens_per_expert.shape[0] @@ -238,7 +294,7 @@ def ep_scatter( # grid = (triton.cdiv(hidden_size, BLOCK_D), num_experts) grid = num_experts - assert m_indices.shape[0] % BLOCK_E == 0 + assert m_indices.shape[0] % align_m == 0 assert expert_start_loc.shape[0] == num_experts # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is. @@ -253,6 +309,7 @@ def ep_scatter( num_warps=num_warps, BLOCK_E=BLOCK_E, BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts), + ALIGN_M=align_m, ) grid = min(recv_topk.shape[0], 1024 * 8) @@ -418,7 +475,7 @@ def deepgemm_moe_permute( if block_size is not None: block_k = block_size - M_sum = compute_aligned_M( + M_sum, align_used = compute_aligned_M_and_alignment( M=topk_ids.size(0), num_topk=topk_ids.size(1), local_num_experts=local_num_experts, @@ -482,11 +539,12 @@ def deepgemm_moe_permute( output_tensor_scale=aq_scale_out, m_indices=expert_ids, output_index=inv_perm, + align_m=align_used, block_size=block_k, pack_ue8m0=pack_ue8m0, ) - return aq_out, aq_scale_out, expert_ids, inv_perm + return aq_out, aq_scale_out, expert_ids, inv_perm, align_used def deepgemm_unpermute_and_reduce( diff --git a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py index c8611217a1..275f80c68f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py @@ -318,11 +318,12 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular): def supports_packed_ue8m0_act_scales(self) -> bool: """ - DeepGemm supports packed ue8m0 activation scales format in devices == sm100 + DeepGemm supports packed ue8m0 activation scales on Blackwell-family + GPUs (SM100 datacenter and SM120 consumer). """ - return ( - is_deep_gemm_e8m0_used() - and current_platform.is_device_capability_family(100) + return is_deep_gemm_e8m0_used() and ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) ) def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 5681d12554..160cff1eeb 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( - compute_aligned_M, + compute_aligned_M_and_alignment, deepgemm_moe_permute, deepgemm_unpermute_and_reduce, ) @@ -43,6 +43,7 @@ from vllm.utils.deep_gemm import ( is_deep_gemm_supported, m_grouped_fp8_fp4_gemm_nt_contiguous, m_grouped_fp8_gemm_nt_contiguous, + mk_alignment_scope, ) from vllm.utils.import_utils import has_deep_gemm @@ -210,10 +211,10 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): # Use the contiguous-layout M alignment (matches apply()); block_shape[0] # is the quant block (1 for MXFP8) and would under-size the workspace. block_m = get_mk_alignment_for_contiguous_layout()[0] - M_sum = compute_aligned_M( + M_sum, align_used = compute_aligned_M_and_alignment( M, topk, local_num_experts, block_m, expert_tokens_meta ) - assert M_sum % block_m == 0 + assert M_sum % align_used == 0 activation_out_dim = self.adjust_N_for_activation(N, activation) workspace1 = (M_sum, max(activation_out_dim, K)) @@ -316,7 +317,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): assert w2.size(1) == K - M_sum = compute_aligned_M( + M_sum, _ = compute_aligned_M_and_alignment( M=topk_ids.size(0), num_topk=topk_ids.size(1), local_num_experts=local_num_experts, @@ -327,7 +328,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): a1q_perm = _resize_cache( workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K) ) - a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute( + a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute( aq=a1q, aq_scale=a1q_scale, topk_ids=topk_ids, @@ -349,23 +350,35 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): else {} ) - mm1_out = _resize_cache(workspace2, (M_sum, N)) - m_grouped_fp8_gemm_nt_contiguous( - (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs - ) + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment; + # otherwise the scheduler can pick the wrong expert id from m_indices + # under cudagraph replay. + with mk_alignment_scope(align_used): + mm1_out = _resize_cache(workspace2, (M_sum, N)) + m_grouped_fp8_gemm_nt_contiguous( + (a1q, a1q_scale), + (w1, self.w1_scale), + mm1_out, + expert_ids, + **gemm_kwargs, + ) - activation_out_dim = self.adjust_N_for_activation(N, activation) - quant_out = _resize_cache( - workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim) - ) - a2q, a2q_scale = self._act_mul_quant( - input=mm1_out.view(-1, N), output=quant_out, activation=activation - ) + activation_out_dim = self.adjust_N_for_activation(N, activation) + quant_out = _resize_cache( + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim) + ) + a2q, a2q_scale = self._act_mul_quant( + input=mm1_out.view(-1, N), output=quant_out, activation=activation + ) - mm2_out = _resize_cache(workspace2, (M_sum, K)) - m_grouped_fp8_gemm_nt_contiguous( - (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, **gemm_kwargs - ) + mm2_out = _resize_cache(workspace2, (M_sum, K)) + m_grouped_fp8_gemm_nt_contiguous( + (a2q, a2q_scale), + (w2, self.w2_scale), + mm2_out, + expert_ids, + **gemm_kwargs, + ) if apply_router_weight_on_input: topk_weights = torch.ones_like(topk_weights) @@ -384,7 +397,8 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): """DeepGemm-based fused MoE expert implementation for FP4 weights. Uses m_grouped_fp8_fp4_gemm_nt_contiguous with FP8 activations and - MXFP4 (FP4 E2M1 packed as uint8) weights. Requires SM100+ (Blackwell). + MXFP4 (FP4 E2M1 packed as uint8) weights. Requires Blackwell-family + GPUs (SM100 datacenter or SM120 consumer). """ # FP8 activation block size (hardcoded since mxfp4_w4a8 quant config @@ -409,9 +423,9 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): def _supports_current_device() -> bool: from vllm.platforms import current_platform - return ( - is_deep_gemm_supported() - and current_platform.is_device_capability_family(100) + return is_deep_gemm_supported() and ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) ) @staticmethod @@ -454,10 +468,10 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): activation: MoEActivation, ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: block_m = get_mk_alignment_for_contiguous_layout()[0] - M_sum = compute_aligned_M( + M_sum, align_used = compute_aligned_M_and_alignment( M, topk, local_num_experts, block_m, expert_tokens_meta ) - assert M_sum % block_m == 0 + assert M_sum % align_used == 0 activation_out_dim = self.adjust_N_for_activation(N, activation) workspace1 = (M_sum, max(activation_out_dim, K)) @@ -533,7 +547,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): if global_num_experts == -1: global_num_experts = local_num_experts - M_sum = compute_aligned_M( + M_sum, _ = compute_aligned_M_and_alignment( M=topk_ids.size(0), num_topk=topk_ids.size(1), local_num_experts=local_num_experts, @@ -544,7 +558,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): a1q_perm = _resize_cache( workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K) ) - a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute( + a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute( aq=a1q, aq_scale=a1q_scale, topk_ids=topk_ids, @@ -555,37 +569,40 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): ) assert a1q.size(0) == M_sum - # FC1: FP8 activations x FP4 weights - # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4). - mm1_out = _resize_cache(workspace2, (M_sum, N)) - m_grouped_fp8_fp4_gemm_nt_contiguous( - (a1q, a1q_scale), - (w1.view(torch.int8), self.w1_scale), - mm1_out, - expert_ids, - recipe_a=(1, self._ACT_BLOCK_K), - recipe_b=(1, self._WEIGHT_BLOCK_K), - ) + # Cap DG's BLOCK_M heuristic at the workspace's per-expert alignment; + # see DeepGemmExperts.apply for rationale. + with mk_alignment_scope(align_used): + # FC1: FP8 activations x FP4 weights + # DeepGEMM 2.4.2 requires FP4-packed weights as int8 (kPackedFP4). + mm1_out = _resize_cache(workspace2, (M_sum, N)) + m_grouped_fp8_fp4_gemm_nt_contiguous( + (a1q, a1q_scale), + (w1.view(torch.int8), self.w1_scale), + mm1_out, + expert_ids, + recipe_a=(1, self._ACT_BLOCK_K), + recipe_b=(1, self._WEIGHT_BLOCK_K), + ) - # SwiGLU activation + FP8 requant - activation_out_dim = self.adjust_N_for_activation(N, activation) - quant_out = _resize_cache( - workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim) - ) - a2q, a2q_scale = self._act_mul_quant( - input=mm1_out.view(-1, N), output=quant_out, activation=activation - ) + # SwiGLU activation + FP8 requant + activation_out_dim = self.adjust_N_for_activation(N, activation) + quant_out = _resize_cache( + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim) + ) + a2q, a2q_scale = self._act_mul_quant( + input=mm1_out.view(-1, N), output=quant_out, activation=activation + ) - # FC2: FP8 activations x FP4 weights - mm2_out = _resize_cache(workspace2, (M_sum, K)) - m_grouped_fp8_fp4_gemm_nt_contiguous( - (a2q, a2q_scale), - (w2.view(torch.int8), self.w2_scale), - mm2_out, - expert_ids, - recipe_a=(1, self._ACT_BLOCK_K), - recipe_b=(1, self._WEIGHT_BLOCK_K), - ) + # FC2: FP8 activations x FP4 weights + mm2_out = _resize_cache(workspace2, (M_sum, K)) + m_grouped_fp8_fp4_gemm_nt_contiguous( + (a2q, a2q_scale), + (w2.view(torch.int8), self.w2_scale), + mm2_out, + expert_ids, + recipe_a=(1, self._ACT_BLOCK_K), + recipe_b=(1, self._WEIGHT_BLOCK_K), + ) if apply_router_weight_on_input: topk_weights = torch.ones_like(topk_weights) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index ab76cea132..5d94d82c01 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -57,6 +57,40 @@ if has_triton_kernels(): ) +def _pack_deepgemm_mxfp4_scales( + w13_weight: torch.Tensor, + w2_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + deepgemm_post_process_weight_scale_block, + ) + + num_experts = w13_weight.shape[0] + intermediate_size_2 = w13_weight.shape[1] # = intermediate*2 + hidden_size = w13_weight.shape[2] * 2 # weight is FP4-packed + intermediate_size = w2_weight.shape[2] * 2 # weight is FP4-packed + block_shape = (1, 32) # MXFP4 block (per-row, K=32) + + return ( + deepgemm_post_process_weight_scale_block( + ws=w13_weight_scale.data, + mn=intermediate_size_2, + k=hidden_size, + quant_block_shape=block_shape, + num_groups=num_experts, + ), + deepgemm_post_process_weight_scale_block( + ws=w2_weight_scale.data, + mn=hidden_size, + k=intermediate_size, + quant_block_shape=block_shape, + num_groups=num_experts, + ), + ) + + class Mxfp4MoeBackend(Enum): NONE = "None" # DeepGEMM FP8xFP4 backend (SM100+) @@ -652,15 +686,18 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( """Convert loaded weights into backend-specific kernel format.""" if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: - from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, + w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( + w13_weight, + w2_weight, + w13_weight_scale, + w2_weight_scale, ) return ( w13_weight.data, w2_weight.data, - _upcast_e8m0_to_fp32(w13_weight_scale.data), - _upcast_e8m0_to_fp32(w2_weight_scale.data), + w13_weight_scale, + w2_weight_scale, w13_bias, w2_bias, ) @@ -1195,17 +1232,18 @@ def convert_weight_to_mxfp4_moe_kernel_format( """ if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: - from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, + w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( + w13_weight, + w2_weight, + w13_weight_scale, + w2_weight_scale, ) - # Weights stay as uint8 packed FP4 — no layout change needed. - # Convert E8M0 uint8 scales to float32. return ( w13_weight.data, w2_weight.data, - _upcast_e8m0_to_fp32(w13_weight_scale.data), - _upcast_e8m0_to_fp32(w2_weight_scale.data), + w13_weight_scale, + w2_weight_scale, w13_bias, w2_bias, ) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index be1167332e..32a2d86899 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -1058,6 +1058,34 @@ def _upcast_e8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor: return fp32_bits.view(torch.float32) +def deepgemm_post_process_weight_scale_block( + ws: torch.Tensor, + mn: int, + k: int, + quant_block_shape: tuple[int, ...], + num_groups: int, + is_sfa: bool = False, +) -> torch.Tensor: + if ws.dtype in (torch.float8_e8m0fnu, torch.uint8): + # Scales already in E8M0 from checkpoint; upcast to fp32 and let + # DeepGEMM pack the layout expected by the target architecture. + ws = _upcast_e8m0_to_fp32(ws) + else: + assert ws.dtype == torch.float32, ( + f"Expected tensor scales dtype to be torch.float32 or " + f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead" + ) + + return transform_sf_into_required_layout( + sf=ws, + mn=mn, + k=k, + recipe=(1, quant_block_shape[0], quant_block_shape[1]), + num_groups=num_groups, + is_sfa=is_sfa, + ) + + def deepgemm_post_process_fp8_weight_block( wq: torch.Tensor, ws: torch.Tensor, @@ -1073,13 +1101,13 @@ def deepgemm_post_process_fp8_weight_block( if ws.dtype in (torch.float8_e8m0fnu, torch.uint8): # Scales already in E8M0 from checkpoint (float8_e8m0fnu, or raw E8M0 - # bits as uint8 for MXFP8) — upcast to fp32 and skip requantization + # bits as uint8 for MXFP8) - upcast to fp32 and skip requantization # (weights already have power-of-two scales). ws = _upcast_e8m0_to_fp32(ws) else: assert ws.dtype == torch.float32, ( f"Expected tensor scales dtype to be torch.float32 or " - f"torch.float8_e8m0fnu, got {ws.dtype} instead" + f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead" ) if use_e8m0: requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape) @@ -1094,16 +1122,12 @@ def deepgemm_post_process_fp8_weight_block( r = wq.size(0) // g wq = wq.view(g, r, d) ws = ws.view(g, r // quant_block_shape[0], d // quant_block_shape[1]) - # Pre-transform scale with recipe=(1, 128, 128) to broadcast + pack - # into TMA-aligned UE8M0 (INT32) layout. At runtime fp8_einsum uses - # recipe=(1, 1, 128) which sees INT dtype and skips re-transform. - dg_ws = transform_sf_into_required_layout( - sf=ws, + dg_ws = deepgemm_post_process_weight_scale_block( + ws=ws, mn=r, k=d, - recipe=(1, quant_block_shape[0], quant_block_shape[1]), + quant_block_shape=quant_block_shape, num_groups=g, - is_sfa=False, ) return wq, dg_ws @@ -1113,22 +1137,12 @@ def deepgemm_post_process_fp8_weight_block( wq = wq.unsqueeze(0) ws = ws.unsqueeze(0) - # From https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/utils/layout.hpp#L46 - # (1, block_n, block_k): (1, 128, 128) for FP8 block, (1, 1, 32) for MXFP8. - recipe = (1, quant_block_shape[0], quant_block_shape[1]) - - # Ref : https://github.com/deepseek-ai/DeepGEMM/blob/c9f8b34dcdacc20aa746b786f983492c51072870/csrc/apis/gemm.hpp - # DeepGemm uses the `transform_sf_into_required_layout` function to - # represent scales in the correct format. - dg_ws = transform_sf_into_required_layout( - sf=ws, + dg_ws = deepgemm_post_process_weight_scale_block( + ws=ws, mn=wq.size(1), k=wq.size(2), - recipe=recipe, + quant_block_shape=quant_block_shape, num_groups=wq.size(0), - # is the scale factors for A in (Refers to the argument A in A @ B). - # Weights are B. - is_sfa=False, ) if original_ndim == 2: diff --git a/vllm/model_executor/warmup/deep_gemm_warmup.py b/vllm/model_executor/warmup/deep_gemm_warmup.py index 78fa68a376..d41604fc7a 100644 --- a/vllm/model_executor/warmup/deep_gemm_warmup.py +++ b/vllm/model_executor/warmup/deep_gemm_warmup.py @@ -12,7 +12,9 @@ from tqdm import tqdm import vllm.envs as envs from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank from vllm.model_executor.layers.fused_moe import MoERunner -from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import ( + compute_aligned_M_and_alignment, +) from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( TritonOrDeepGemmExperts, @@ -25,6 +27,7 @@ from vllm.utils.deep_gemm import ( fp8_gemm_nt, get_mk_alignment_for_contiguous_layout, m_grouped_fp8_gemm_nt_contiguous, + mk_alignment_scope, ) from vllm.utils.math_utils import cdiv from vllm.utils.platform_utils import num_compute_units @@ -238,7 +241,7 @@ def _get_grouped_gemm_params( w2: torch.Tensor, num_topk: int, max_tokens: int, -) -> tuple[int, int, torch.Tensor]: +) -> tuple[int, int, list[tuple[int, int, torch.Tensor]]]: assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts" block_m = get_mk_alignment_for_contiguous_layout()[0] @@ -248,19 +251,46 @@ def _get_grouped_gemm_params( # Assumes all ranks have the same max_num_batched_tokens max_tokens = get_dp_group().world_size * max_tokens - # This is the maximum GroupedGemm M size that we expect to run - # the grouped_gemm with. - MAX_M = compute_aligned_M( - max_tokens, num_topk, num_experts, block_m, expert_tokens_meta=None + request_m_values = _generate_optimal_warmup_m_values( + max_tokens, + max(w1.size(1), w2.size(1)), + device, ) - # Distribute expert-ids evenly. - MAX_BLOCKS = MAX_M // block_m - expert_ids_block = torch.randint( - low=0, high=num_experts, size=(MAX_BLOCKS,), device=device, dtype=torch.int32 - ) - expert_ids = torch.repeat_interleave(expert_ids_block, block_m, dim=0) + request_m_values = sorted({m for m in (*request_m_values, max_tokens) if m > 0}) + if not request_m_values: + return 0, block_m, [] - return MAX_M, block_m, expert_ids + cases_by_shape: dict[tuple[int, int], torch.Tensor] = {} + for request_m in request_m_values: + M_sum, align_used = compute_aligned_M_and_alignment( + M=request_m, + num_topk=num_topk, + local_num_experts=num_experts, + alignment=block_m, + expert_tokens_meta=None, + ) + if (M_sum, align_used) in cases_by_shape: + continue + + num_blocks = M_sum // align_used + expert_ids_block = torch.randint( + low=0, + high=num_experts, + size=(num_blocks,), + device=device, + dtype=torch.int32, + ) + cases_by_shape[(M_sum, align_used)] = torch.repeat_interleave( + expert_ids_block, align_used, dim=0 + ) + + max_m = max(M_sum for M_sum, _ in cases_by_shape) + warmup_cases = [ + (M_sum, align_used, expert_ids) + for (M_sum, align_used), expert_ids in sorted(cases_by_shape.items()) + ] + + return max_m, block_m, warmup_cases def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( @@ -278,7 +308,11 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( ): return - MAX_M, block_m, expert_ids = _get_grouped_gemm_params(w1, w2, num_topk, max_tokens) + MAX_M, block_m, warmup_cases = _get_grouped_gemm_params( + w1, w2, num_topk, max_tokens + ) + if not warmup_cases: + return device = w1.device def _warmup(w: torch.Tensor, w_scale: torch.Tensor): @@ -289,15 +323,14 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup( ) out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16) - m_values = list(range(block_m, MAX_M + 1, block_m)) - - for num_tokens in m_values: - m_grouped_fp8_gemm_nt_contiguous( - (a1q[:num_tokens], a1q_scales[:num_tokens]), - (w, w_scale), - out[:num_tokens], - expert_ids[:num_tokens], - ) + for num_tokens, align_used, expert_ids in warmup_cases: + with mk_alignment_scope(align_used): + m_grouped_fp8_gemm_nt_contiguous( + (a1q[:num_tokens], a1q_scales[:num_tokens]), + (w, w_scale), + out[:num_tokens], + expert_ids, + ) if pbar is not None: pbar.update(1) @@ -350,8 +383,8 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int: w13, _, w2, _, num_topk = _extract_data_from_fused_moe_module(m) if w13.size() in seen_grouped_sizes and w2.size() in seen_grouped_sizes: continue - MAX_M, block_m, _ = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens) - n_values = (MAX_M - block_m) // block_m + 1 + _, _, warmup_cases = _get_grouped_gemm_params(w13, w2, num_topk, max_tokens) + n_values = len(warmup_cases) if w13.size() not in seen_grouped_sizes: total += n_values seen_grouped_sizes.add(w13.size()) diff --git a/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py new file mode 100644 index 0000000000..5b4900b50c --- /dev/null +++ b/vllm/model_executor/warmup/deepseek_v4_mhc_warmup.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warm up DeepSeek V4 mHC TileLang kernels before serving requests. + +Ported from lucifer1004/vllm-jasl with the two env-var knobs removed +(`VLLM_ENABLE_DEEPSEEK_V4_MHC_WARMUP`, `VLLM_DEEPSEEK_V4_MHC_WARMUP_TOKEN_SIZES`). +Gating is intrinsic: non-DSv4 models and layers without hc_* attributes +return early, so the warmup is a no-op except where it's needed. +""" + +import time +from collections.abc import Iterable + +import torch + +from vllm.logger import init_logger +from vllm.tracing import instrument +from vllm.utils.math_utils import cdiv + +logger = init_logger(__name__) + +_AUTO_WARMUP_MAX_TOKENS = 16_384 +_DEFAULT_TOKEN_SIZE_CANDIDATES = ( + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16_384, +) + + +def _compute_mhc_pre_num_split( + *, + num_tokens: int, + hidden_size: int, + hc_mult: int, + num_sms: int, +) -> int: + block_k = 64 + block_m = 64 + k = hc_mult * hidden_size + grid_size = cdiv(num_tokens, block_m) + split_k = num_sms // grid_size + num_block_k = cdiv(k, block_k) + split_k = min(split_k, num_block_k // 4) + return max(split_k, 1) + + +def _normalize_token_sizes( + token_sizes: Iterable[int], + *, + max_tokens: int, +) -> list[int]: + return sorted({size for size in token_sizes if 1 <= size <= max_tokens}) + + +def _select_mhc_warmup_token_sizes( + *, + max_tokens: int, + cudagraph_capture_sizes: list[int], +) -> list[int]: + if max_tokens <= 0: + return [] + + max_auto_tokens = min(max_tokens, _AUTO_WARMUP_MAX_TOKENS) + candidates = list(_DEFAULT_TOKEN_SIZE_CANDIDATES) + candidates.extend(cudagraph_capture_sizes) + candidates.append(max_auto_tokens) + return _normalize_token_sizes(candidates, max_tokens=max_auto_tokens) + + +def _find_first_mhc_layer(model: torch.nn.Module) -> torch.nn.Module | None: + for module in model.modules(): + if module.__class__.__name__ != "DeepseekV4DecoderLayer": + continue + if all( + hasattr(module, attr) + for attr in ( + "hc_pre", + "hc_post", + "hc_attn_fn", + "hc_attn_scale", + "hc_attn_base", + "hc_ffn_fn", + "hc_ffn_scale", + "hc_ffn_base", + ) + ): + return module + return None + + +def _find_deepseek_v4_model(model: torch.nn.Module) -> torch.nn.Module | None: + for module in model.modules(): + if module.__class__.__name__ != "DeepseekV4Model": + continue + if all( + hasattr(module, attr) + for attr in ("hc_head_fn", "hc_head_scale", "hc_head_base") + ): + return module + return None + + +def _warmup_layer_mhc( + layer: torch.nn.Module, + token_sizes: list[int], +) -> None: + max_tokens = max(token_sizes) + hidden_size = int(layer.hidden_size) + hc_mult = int(layer.hc_mult) + device = layer.hc_attn_fn.device + residual = torch.zeros( + max_tokens, + hc_mult, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + + for size in token_sizes: + residual_slice = residual[:size] + for fn, scale, base in ( + (layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base), + (layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base), + ): + layer_input, post_mix, comb_mix = layer.hc_pre( + residual_slice, + fn, + scale, + base, + ) + layer.hc_post(layer_input, residual_slice, post_mix, comb_mix) + + +def _warmup_hc_head( + model: torch.nn.Module, + token_sizes: list[int], +) -> None: + # Upstream a8887c208 ("[DSV4] aiter mhc support (ROCm)") refactored + # ``hc_head`` from a free function into the ``HCHeadOp`` CustomOp + # instance attached to the model as ``hc_head_op``. We call through + # that instance so the warmup exercises the same dispatched + # implementation as the inference path. + hc_head_op = getattr(model, "hc_head_op", None) + if hc_head_op is None: + return + + max_tokens = max(token_sizes) + hidden_size = int(model.config.hidden_size) + hc_mult = int(model.hc_mult) + device = model.hc_head_fn.device + hidden_states = torch.zeros( + max_tokens, + hc_mult, + hidden_size, + dtype=torch.bfloat16, + device=device, + ) + + for size in token_sizes: + hc_head_op( + hidden_states[:size], + model.hc_head_fn, + model.hc_head_scale, + model.hc_head_base, + model.rms_norm_eps, + model.hc_eps, + ) + + +@instrument(span_name="DeepSeek V4 mHC warmup") +def deepseek_v4_mhc_warmup( + model: torch.nn.Module, + *, + max_tokens: int, + cudagraph_capture_sizes: list[int] | None = None, +) -> None: + # Cheap model-type gate before walking ``model.modules()``. The class + # walk below is O(num_layers) and shows up in startup time on very + # large checkpoints; bail out for any model that is not DeepSeek V4. + config = getattr(model, "config", None) + model_type = getattr(config, "model_type", None) if config is not None else None + if model_type is not None and model_type != "deepseek_v4": + return + + layer = _find_first_mhc_layer(model) + if layer is None: + return + + device = layer.hc_attn_fn.device + if device.type != "cuda": + return + + deepseek_model = _find_deepseek_v4_model(model) + token_sizes = _select_mhc_warmup_token_sizes( + max_tokens=max_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes or [], + ) + if not token_sizes: + return + + started = time.perf_counter() + logger.info( + "Warming up DeepSeek V4 mHC TileLang kernels for token sizes: %s", + token_sizes, + ) + with torch.inference_mode(): + _warmup_layer_mhc(layer, token_sizes) + if deepseek_model is not None: + _warmup_hc_head(deepseek_model, token_sizes) + torch.accelerator.synchronize() + logger.info( + "DeepSeek V4 mHC TileLang warmup finished in %.2f seconds.", + time.perf_counter() - started, + ) diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py new file mode 100644 index 0000000000..8c5bbb0fd4 --- /dev/null +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer autotune cache helpers.""" + +import hashlib +import os +import tempfile +from contextlib import suppress +from pathlib import Path +from typing import TYPE_CHECKING + +import vllm.envs as envs +from vllm.compilation.caching import aot_compile_hash_factors + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + +def flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: + factors = aot_compile_hash_factors(runner.vllm_config) + return hashlib.sha256(str(factors).encode()).hexdigest() + + +def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: + override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR + if override_dir: + root = Path(override_dir).expanduser() + else: + from flashinfer.jit import env as flashinfer_jit_env + + flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR + root = ( + Path(envs.VLLM_CACHE_ROOT) + / "flashinfer_autotune_cache" + / flashinfer_workspace.parent.name + / flashinfer_workspace.name + ) + + output_dir = root / flashinfer_autotune_cache_hash(runner) + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir / "autotune_configs.json" + + +def write_flashinfer_autotune_cache(cache_path: Path, contents: bytes) -> None: + cache_path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_path = tempfile.mkstemp( + dir=cache_path.parent, suffix=".tmp", prefix=f".{cache_path.name}." + ) + try: + with os.fdopen(fd, "wb") as f: + f.write(contents) + os.replace(tmp_path, cache_path) + except BaseException: + with suppress(OSError): + os.unlink(tmp_path) + raise diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py new file mode 100644 index 0000000000..44be769e24 --- /dev/null +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Warmup and autotune helpers for FlashInfer sparse MLA backends.""" + +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.logger import init_logger +from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + resolve_flashinfer_autotune_file, + write_flashinfer_autotune_cache, +) +from vllm.platforms import current_platform +from vllm.utils.flashinfer import autotune as flashinfer_autotune +from vllm.utils.flashinfer import has_flashinfer +from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2GPUModelRunner + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + +_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset( + { + "FLASHMLA_SPARSE_DSV4", + "FLASHINFER_MLA_SPARSE_DSV4", + "ROCM_FLASHMLA_SPARSE_DSV4", + "DEEPSEEK_SPARSE_SWA", + } +) +_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_SM120"}) +_DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS = frozenset({"FLASHINFER_MLA_SPARSE_DSV4"}) + +_FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS = { + "FLASHINFER_MLA_SPARSE_SM120": "DSv3.2", + "FLASHINFER_MLA_SPARSE_DSV4": "DSv4", +} + +_SPARSE_MLA_MIXED_WARMUP_TOKENS = 16 + + +def _attention_backend_name(backend: object) -> str | None: + get_name = getattr(backend, "get_name", None) + if get_name is None: + return None + try: + return get_name() + except NotImplementedError: + return None + + +def _has_deepseek_v4_sparse_mla_backend(runner: "GPUModelRunner") -> bool: + for groups in getattr(runner, "attn_groups", []) or (): + for group in groups: + name = _attention_backend_name(getattr(group, "backend", None)) + if name in _DEEPSEEK_V4_SPARSE_MLA_BACKENDS: + return True + return False + + +def _flashinfer_sparse_mla_decode_label( + runner: "GPUModelRunner", + allowed_backends: frozenset[str], +) -> str | None: + for groups in getattr(runner, "attn_groups", []) or (): + for group in groups: + name = _attention_backend_name(getattr(group, "backend", None)) + if name in allowed_backends: + return _FLASHINFER_SM120_SPARSE_MLA_DECODE_LABELS.get(name) + return None + + +def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int: + return max(0, min(num_tokens, max_tokens)) + + +def _uses_v2_model_runner(runner: "GPUModelRunner") -> bool: + vllm_config = getattr(runner, "vllm_config", None) + return bool(getattr(vllm_config, "use_v2_model_runner", False)) + + +def _run_flashinfer_sparse_mla_decode_autotune( + worker: "Worker", + num_tokens: int, + allowed_backends: frozenset[str], +) -> bool: + """Autotune FlashInfer's SM120 sparse-MLA decode path.""" + runner = worker.model_runner + log_label = _flashinfer_sparse_mla_decode_label(runner, allowed_backends) + if log_label is None: + return False + if worker.vllm_config.kernel_config.enable_flashinfer_autotune is not True: + return False + if not has_flashinfer() or not current_platform.is_device_capability_family(120): + return False + + try: + from flashinfer.autotuner import AutoTuner + except ImportError: + logger.warning( + "Skipping FlashInfer SM120 sparse MLA decode autotune because " + "FlashInfer autotuner is unavailable." + ) + return False + + from vllm.distributed.parallel_state import get_world_group + + world = get_world_group() + is_leader = world.rank_in_group == 0 + cache_path = resolve_flashinfer_autotune_file(runner) + + dummy_run_kwargs = dict( + num_tokens=num_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) + + if is_leader: + logger.info( + "Autotuning FlashInfer SM120 sparse MLA %s decode with cache: %s", + log_label, + cache_path, + ) + + with torch.inference_mode(): + warmup_executed = True + if is_leader: + if _uses_v2_model_runner(runner): + v2_runner = cast("V2GPUModelRunner", runner) + warmup_executed = run_mixed_prefill_decode_warmup( + v2_runner, + worker.execute_model, + worker.sample_tokens, + num_tokens, + mixed_step_context=flashinfer_autotune(True, cache=str(cache_path)), + req_id_prefix="_sparse_mla_v2_warmup", + ) + else: + with flashinfer_autotune(True, cache=str(cache_path)): + runner._dummy_run(**dummy_run_kwargs) + else: + if _uses_v2_model_runner(runner): + v2_runner = cast("V2GPUModelRunner", runner) + warmup_executed = run_mixed_prefill_decode_warmup( + v2_runner, + worker.execute_model, + worker.sample_tokens, + num_tokens, + req_id_prefix="_sparse_mla_v2_warmup", + ) + else: + runner._dummy_run(**dummy_run_kwargs) + + if not warmup_executed: + return False + + tune_results: bytes | None = None + if is_leader and cache_path.exists(): + with open(cache_path, "rb") as f: + tune_results = f.read() + + tune_results = world.broadcast_object(tune_results, src=0) + if tune_results is None: + logger.warning( + "No FlashInfer SM120 sparse MLA %s decode autotune cache entries found. " + "Falling back to FlashInfer's default tactic heuristic.", + log_label, + ) + world.barrier() + return True + + write_flashinfer_autotune_cache(cache_path, tune_results) + world.barrier() + + AutoTuner.get().load_configs(str(cache_path)) + logger.info( + "FlashInfer SM120 sparse MLA %s decode autotune cache loaded on rank %d " + "from %s.", + log_label, + world.rank_in_group, + cache_path, + ) + return True + + +def _flashinfer_sparse_mla_decode_autotune( + worker: "Worker", + num_tokens: int, +) -> bool: + return _run_flashinfer_sparse_mla_decode_autotune( + worker, num_tokens, _FLASHINFER_MLA_SPARSE_BACKENDS + ) + + +def _deepseek_v4_sparse_mla_decode_autotune( + worker: "Worker", + num_tokens: int, +) -> bool: + return _run_flashinfer_sparse_mla_decode_autotune( + worker, num_tokens, _DEEPSEEK_V4_FLASHINFER_MLA_SPARSE_BACKENDS + ) + + +def flashinfer_sparse_mla_decode_autotune_warmup(worker: "Worker") -> None: + """Autotune generic FlashInfer sparse MLA decode when selected.""" + runner = worker.model_runner + if runner.is_pooling_model: + return + + max_tokens = worker.scheduler_config.max_num_batched_tokens + mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens) + if mixed_tokens <= 0: + return + _flashinfer_sparse_mla_decode_autotune(worker, mixed_tokens) + + +def deepseek_v4_sparse_mla_attention_warmup(worker: "Worker") -> None: + """Warm DSv4 sparse-MLA mixed prefill+decode attention.""" + runner = worker.model_runner + if runner.is_pooling_model or not _has_deepseek_v4_sparse_mla_backend(runner): + return + + max_tokens = worker.scheduler_config.max_num_batched_tokens + mixed_tokens = _clamp_warmup_tokens(_SPARSE_MLA_MIXED_WARMUP_TOKENS, max_tokens) + if mixed_tokens <= 0: + return + + logger.info( + "Warming up DeepSeek V4 sparse MLA attention for mixed tokens=%s.", + mixed_tokens, + ) + mixed_warmup_done = _deepseek_v4_sparse_mla_decode_autotune(worker, mixed_tokens) + if not mixed_warmup_done: + if _uses_v2_model_runner(runner): + v2_runner = cast("V2GPUModelRunner", runner) + run_mixed_prefill_decode_warmup( + v2_runner, + worker.execute_model, + worker.sample_tokens, + mixed_tokens, + req_id_prefix="_sparse_mla_v2_warmup", + ) + else: + runner._dummy_run( + num_tokens=mixed_tokens, + skip_eplb=True, + is_profile=True, + force_attention=True, + create_mixed_batch=True, + ) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 61d2376abb..754270e652 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -6,16 +6,24 @@ This is useful specifically for JIT'ed kernels as we don't want JIT'ing to happen during model execution. """ -import hashlib -from pathlib import Path from typing import TYPE_CHECKING import torch import vllm.envs as envs -from vllm.compilation.caching import aot_compile_hash_factors from vllm.logger import init_logger from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup +from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import ( + deepseek_v4_mhc_warmup, +) +from vllm.model_executor.warmup.flashinfer_autotune_cache import ( + resolve_flashinfer_autotune_file, + write_flashinfer_autotune_cache, +) +from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( + deepseek_v4_sparse_mla_attention_warmup, + flashinfer_sparse_mla_decode_autotune_warmup, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -27,36 +35,26 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def _flashinfer_autotune_cache_hash(runner: "GPUModelRunner") -> str: - factors = aot_compile_hash_factors(runner.vllm_config) - return hashlib.sha256(str(factors).encode()).hexdigest() - - -def _resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: - override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR - if override_dir: - root = Path(override_dir).expanduser() - else: - from flashinfer.jit import env as flashinfer_jit_env - - flashinfer_workspace = flashinfer_jit_env.FLASHINFER_WORKSPACE_DIR - root = ( - Path(envs.VLLM_CACHE_ROOT) - / "flashinfer_autotune_cache" - / flashinfer_workspace.parent.name - / flashinfer_workspace.name - ) - - output_dir = root / _flashinfer_autotune_cache_hash(runner) - output_dir.mkdir(parents=True, exist_ok=True) - return output_dir / "autotune_configs.json" - - def kernel_warmup(worker: "Worker"): from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( minimax_m3_msa_warmup, ) + # DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder + # layer per token; warm them across token sizes first so the first real + # request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside). + deepseek_v4_mhc_warmup( + worker.get_model(), + max_tokens=worker.scheduler_config.max_num_batched_tokens, + cudagraph_capture_sizes=( + worker.vllm_config.compilation_config.cudagraph_capture_sizes or [] + ), + ) + + # Run next so input-prep kernels JIT against pristine runner state. + flashinfer_sparse_mla_decode_autotune_warmup(worker) + deepseek_v4_sparse_mla_attention_warmup(worker) + # Deep GEMM warmup do_deep_gemm_warmup = ( envs.VLLM_USE_DEEP_GEMM @@ -147,7 +145,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: world = get_world_group() is_leader = world.rank_in_group == 0 - cache_path = _resolve_flashinfer_autotune_file(runner) + cache_path = resolve_flashinfer_autotune_file(runner) if is_leader: logger.info("Using FlashInfer autotune cache file: %s", cache_path) @@ -183,9 +181,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: "Falling back to default tactics." ) else: - if not is_leader and world.local_rank == 0: - with open(cache_path, "wb") as f: - f.write(tune_results) + write_flashinfer_autotune_cache(cache_path, tune_results) world.barrier() from flashinfer.autotuner import AutoTuner diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 2930258488..29a19d9026 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -62,23 +62,22 @@ logger = init_logger(__name__) def _resolve_dsv4_kv_cache_dtype( - use_flashmla_fp8_layout: bool, + use_fp8_ds_mla_layout: bool, kv_cache_dtype: str, cache_config: CacheConfig | None, ) -> tuple[str, torch.dtype]: """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``. Both layouts are paged; they differ in the per-token block format. The - FlashMLA fp8 layout (FlashMLA / ROCm Aiter) is the ``fp8_ds_mla`` format: - UE8M0 block-scaled fp8 packed as ``uint8`` (the canonical ``fp8_ds_mla`` - string is written back onto ``cache_config`` so the page-size specs pick - the 576B per-token slot). Otherwise (FlashInfer) each token's KV row is - stored in its plain element dtype — bf16 or per-tensor FP8 E4M3. + ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the + canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the + page-size specs pick the 576B per-token slot). Plain-row backends store each + token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3. """ - if use_flashmla_fp8_layout: + if use_fp8_ds_mla_layout: # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. assert kv_cache_dtype.startswith("fp8"), ( - f"DeepseekV4 FlashMLA fp8 layout only supports fp8 kv-cache, " + f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, " f"got {kv_cache_dtype}" ) if kv_cache_dtype != "fp8_ds_mla": @@ -100,18 +99,20 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): The platform-specific sparse-MLA forward (``forward_mqa`` / ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a - subclass — ``DeepseekV4FlashMLAAttention`` / ``DeepseekV4FlashInferMLAAttention`` - (CUDA) or ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the - platform-specific deepseek_v4 model module. The base is never instantiated - directly. + subclass — ``DeepseekV4FlashMLAAttention`` / + ``DeepseekV4FlashInferSM120Attention`` / + ``DeepseekV4FlashInferMLAAttention`` (CUDA) or + ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific + deepseek_v4 model module. The base is never instantiated directly. """ # Provided by the platform subclass. backend_cls: ClassVar[type[AttentionBackend]] # KV-cache per-token block format (both layouts are paged). True (default) - # = FlashMLA / ROCm fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); - # False = FlashInfer plain bf16 / per-tensor fp8 KV row. - use_flashmla_fp8_layout: ClassVar[bool] = True + # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain + # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when + # a single attention class dispatches across arch-specific layouts. + use_fp8_ds_mla_layout: ClassVar[bool] = True # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather # workspace allocated in _forward_prefill and is also read by the dummy-run # path to pre-reserve that workspace. @@ -145,6 +146,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" raise NotImplementedError + def _uses_fp8_ds_mla_layout(self) -> bool: + """Return whether this instance stores fp8 KV in fp8_ds_mla layout.""" + return self.use_fp8_ds_mla_layout + def __init__( self, vllm_config: VllmConfig, @@ -276,13 +281,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): ) self.max_model_len = vllm_config.model_config.max_model_len - # Resolve the kv-cache dtype from this backend's block format (a - # ClassVar set by the subclass): fp8_ds_mla (UE8M0 block-scaled fp8 as - # uint8) for FlashMLA / ROCm, vs a plain bf16 / per-tensor fp8 row for - # FlashInfer. The same resolution drives the SWA cache tensor dtype - # below. + # Resolve the kv-cache dtype from this backend's block format. The same + # resolution drives the SWA cache tensor dtype below. self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( - self.use_flashmla_fp8_layout, cache_config.cache_dtype, cache_config + self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config ) self.swa_cache_layer = DeepseekV4SWACache( @@ -539,7 +541,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): # kv is unchanged; attention reads kv solely via swa_kv_cache. if cache_dtype == torch.uint8: - # Legacy FlashMLA UE8M0 paged path. Horizontally fused: + # fp8_ds_mla UE8M0 paged path. Horizontally fused: # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling # the padding head slots; the kernel allocates and returns # the padded q tensor. @@ -557,10 +559,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): swa_metadata.block_size, ) - # FlashInfer full-cache path: the [num_blocks, block_size, 512] cache - # stores the KV row in its plain dtype (no Q padding). bf16 rewrites q - # in place; per-tensor fp8 writes a separately-allocated fp8 q and - # quantizes the KV row. + # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV + # row in its element dtype (no Q padding). bf16 rewrites q in place; + # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the + # KV row. block_size = swa_metadata.block_size swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim) if cache_dtype == torch.bfloat16: @@ -601,18 +603,18 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.compress_ratio <= 1 ): # SWA part. Allocated separately as DeepseekV4SWACache. return None - # FlashMLA uses the fp8_ds_mla block format (UE8M0 block-scaled fp8 as - # uint8, 576B aligned); FlashInfer stores a plain bf16 / per-tensor fp8 - # row with no extra alignment. - is_flashmla = self.kv_cache_dtype == "fp8_ds_mla" + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B + # alignment; plain bf16 / per-tensor fp8 rows use natural element-size + # pages. + uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, - dtype=torch.uint8 if is_flashmla else self.kv_cache_torch_dtype, + dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, compress_ratio=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if is_flashmla else None, # FlashMLA needs 576B + alignment=576 if uses_fp8_ds_mla_layout else None, model_version="deepseek_v4", ) diff --git a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py index 000bb51b20..b667b87679 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py +++ b/vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py @@ -40,9 +40,18 @@ def _fused_inv_rope_fp8_quant_per_head( USE_GDC: tl.constexpr, launch_pdl: tl.constexpr, # triton metadata ): - # int64: stride multiply overflows int32 past num_tokens=32768 (IMA). + # Cast every stride to int64 — without this, Python-int strides are + # inferred as int32 and `pid_token(int64) × stride(int32)` can lower to + # int32 arithmetic, wrapping past 2³¹ for large prefill batches → IMA. pid_token = tl.program_id(0).to(tl.int64) pid_gh = tl.program_id(1).to(tl.int64) + o_stride_token = o_stride_token.to(tl.int64) + o_stride_head = o_stride_head.to(tl.int64) + cache_stride_pos = cache_stride_pos.to(tl.int64) + fp8_stride_group = fp8_stride_group.to(tl.int64) + fp8_stride_token = fp8_stride_token.to(tl.int64) + scale_stride_group = scale_stride_group.to(tl.int64) + scale_stride_k = scale_stride_k.to(tl.int64) g = pid_gh // heads_per_group head_in_group = pid_gh % heads_per_group diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 20be18e336..1efa987fe7 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -155,17 +155,17 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase): raise ValueError(f"Invalid compress ratio: {compress_ratio}") def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # FlashMLA's UE8M0 paged layout needs 576B alignment; the FlashInfer - # full-cache path shares state pages with contiguous KV pages, so - # padding would break page matching. - is_flashmla = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain + # full-cache rows share state pages with contiguous KV pages, so padding + # would break page matching. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( # only has one vector instead of K + V block_size=self.block_size, num_kv_heads=1, head_size=self.state_dim, dtype=self.dtype, sliding_window=self.sliding_window, - alignment=576 if is_flashmla else None, + alignment=576 if uses_fp8_ds_mla_layout else None, ) def forward(self): ... @@ -340,8 +340,8 @@ class DeepseekCompressor(nn.Module): k_cache_layer = self._static_forward_context[self.k_cache_prefix] kv_cache = k_cache_layer.kv_cache - # FlashInfer V4 reads a contiguous bf16 / per-tensor fp8 cache row; the - # legacy FlashMLA path uses the UE8M0 paged uint8 layout. + # Plain-row V4 reads a contiguous bf16 / per-tensor fp8 cache row; the + # fp8_ds_mla path uses the UE8M0 paged uint8 layout. store_full_kv = self.head_dim == 512 and kv_cache.dtype != torch.uint8 store_full_fp8 = kv_cache.dtype == torch.float8_e4m3fn fp8_scale = ( @@ -358,8 +358,8 @@ class DeepseekCompressor(nn.Module): compress_norm_rope_store_cutedsl, ) - # head=512 on CUDA always uses cutedsl, for both the legacy UE8M0 - # layout and the FlashInfer full-cache layout. The full-cache flags + # head=512 on CUDA always uses cutedsl, for both the fp8_ds_mla + # layout and the plain full-cache layout. The full-cache flags # are consumed only here. compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl extra_kwargs: dict[str, Any] = dict( diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 9b2542450b..f35fa03252 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -1,13 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""DeepSeek V4 FlashInfer TRTLLM-gen sparse MLA backend. - -Uses FlashInfer's public ``trtllm_batch_decode_sparse_mla_dsv4`` launcher with a -plain bf16 / per-tensor FP8 KV row (vs FlashMLA's packed ``fp8_ds_mla`` block -format). Shares the V4 sparse-index pipeline (SWA cache + compressor + indexer, -256-token blocks, head_size 512) with the FlashMLA V4 backend; only the -attention forward differs. -""" +"""DeepSeek V4 FlashInfer sparse MLA backend.""" from typing import TYPE_CHECKING, ClassVar, cast @@ -18,6 +11,7 @@ from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.models.deepseek_v4.common.ops import ( build_flashinfer_mixed_sparse_indices, + compute_global_topk_indices_and_lens, ) from vllm.models.deepseek_v4.nvidia.ops.o_proj import ( compute_fp8_einsum_recipe, @@ -27,13 +21,14 @@ from vllm.models.deepseek_v4.sparse_mla import ( DeepseekV4FlashMLABackend, DeepseekV4FlashMLAMetadata, ) +from vllm.platforms import current_platform +from vllm.platforms.interface import DeviceCapability from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4 +from vllm.v1.attention.backend import MultipleOf if TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata -# 128 MB TRTLLM-gen workspace, allocated once per device and zero-initialized -# (required for first use). Reused across all FlashInfer V4 layers. _FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 _flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {} @@ -51,34 +46,113 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend): - """Shares the FlashMLA V4 metadata/cache pipeline; swaps the attention impl. + """FlashInfer backend using the DSv4 sparse metadata/cache layout. - Inheriting from the FlashMLA V4 backend reuses its ``DeepseekV4FlashMLAMetadata`` - builder. + Inheriting from the FlashMLA V4 backend reuses its + ``DeepseekV4FlashMLAMetadata`` builder. """ supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16", "fp8"] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "bfloat16", + "fp8", + "fp8_e4m3", + "fp8_ds_mla", + ] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [256] @staticmethod def get_name() -> str: return "FLASHINFER_MLA_SPARSE_DSV4" + @classmethod + def get_supported_head_sizes(cls) -> list[int]: + return [512] + + @classmethod + def supports_sink(cls) -> bool: + return True + + @classmethod + def is_sparse(cls) -> bool: + return True + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return capability.major in [10, 12] + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + if device_capability.major == 10: + if kv_cache_dtype == "fp8_ds_mla": + return ( + "FLASHINFER_MLA_SPARSE_DSV4 SM10x uses the plain " + "per-tensor FP8 KV layout, not fp8_ds_mla" + ) + if kv_cache_dtype not in (None, "auto", "bfloat16", "fp8", "fp8_e4m3"): + return "kv_cache_dtype not supported" + return None + if device_capability.major == 12: + if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"): + return "kv_cache_dtype not supported" + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 + + if not has_flashinfer_sparse_mla_sm120(): + return ( + "FLASHINFER_MLA_SPARSE_DSV4 SM120 requires FlashInfer's " + "sparse MLA decode API" + ) + return None + return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x" + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + device_capability = current_platform.get_device_capability() + if device_capability is not None and device_capability.major == 12: + return DeepseekV4FlashMLABackend.get_kv_cache_shape( + num_blocks, + block_size, + num_kv_heads, + head_size, + cache_dtype_str, + ) + assert num_kv_heads == 1 + return (num_blocks, block_size, head_size) + class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): - """FlashInfer TRTLLM-gen sparse MLA attention layer for DeepSeek V4.""" + """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4.""" backend_cls = DeepseekV4FlashInferMLASparseBackend - # FlashInfer stores a plain bf16 / per-tensor fp8 KV row, not the FlashMLA - # packed fp8_ds_mla block format (UE8M0 block-scaled fp8 as uint8). - use_flashmla_fp8_layout: ClassVar[bool] = False + use_fp8_ds_mla_layout: ClassVar[bool] = False @classmethod def get_padded_num_q_heads(cls, num_heads: int) -> int: # FP8 decode kernel only supports h_q = 64 or 128. if num_heads > 128: raise ValueError( - f"DeepseekV4 Flashinfer MLA Sparse does not support {num_heads} heads " + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " "(FP8 decode kernel requires h_q in {64, 128})." ) return 64 if num_heads <= 64 else 128 @@ -106,8 +180,6 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): # per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``. if self.kv_cache_torch_dtype != torch.float8_e4m3fn: return - # TODO: load real per-tensor Q/KV scales from the checkpoint; unit - # scales until the scale tensor names are wired. fp8_q_scale = 1.0 fp8_kv_scale = 1.0 self.register_buffer( @@ -125,9 +197,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): torch.tensor([fp8_kv_scale], dtype=torch.float32), persistent=False, ) - # TRTLLM-gen takes scalar scale args on a distinct (correct) C++ path - # vs 1-elem tensors, so these are Python floats. bmm1 folds the softmax - # scale and the Q/KV per-tensor scales; bmm2 is the KV scale. + # TRTLLM-gen takes scalar scale args on a distinct C++ path vs + # one-element tensors, so these are Python floats. self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale self._flashinfer_fp8_bmm2_scale = fp8_kv_scale @@ -387,9 +458,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): query_start_loc_cpu = swa_metadata.query_start_loc_cpu assert query_start_loc is not None and query_start_loc_cpu is not None - # Keep Perkz's two-call decode/prefill split: the TRTLLM-gen launcher is - # tuned for uniform-q batches, and collapsing the mixed batch into a - # single call is the suspected source of the prior IMA. + # Keep the TRTLLM-gen decode/prefill split: the launcher is tuned for + # uniform-q batches, and this avoids flattening mixed batches into one call. if num_decode_tokens > 0: decode_cu = query_start_loc[: num_decodes + 1] decode_cu_cpu = query_start_loc_cpu[: num_decodes + 1] @@ -434,3 +504,379 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): cum_seq_lens_q=prefill_cu, max_q_len=int(prefill_lens_cpu.max().item()), ) + + +class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention): + """DeepSeek V4 sparse MLA attention through FlashInfer's SM120 kernels.""" + + backend_cls = DeepseekV4FlashInferMLASparseBackend + use_fp8_ds_mla_layout: ClassVar[bool] = True + + @staticmethod + def _get_workspace(device: torch.device) -> torch.Tensor: + return _get_flashinfer_dsv4_workspace(device) + + @staticmethod + def _as_sparse_cache(kv_cache: torch.Tensor) -> torch.Tensor: + if kv_cache.dtype == torch.float8_e4m3fn: + kv_cache = kv_cache.view(torch.uint8) + if kv_cache.dim() == 4: + return kv_cache + return kv_cache.unsqueeze(-2) + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + if num_heads <= 16: + return 16 + if num_heads <= 32: + return 32 + if num_heads <= 64: + return 64 + if num_heads <= 128: + return 128 + raise ValueError( + f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads " + "(SM120 kernel requires h_q in {16, 32, 64, 128})." + ) + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + return deep_gemm_fp8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + tma_aligned_scales=self._tma_aligned_scales, + ) + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 + + if not has_flashinfer_sparse_mla_sm120(): + raise RuntimeError( + "FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires FlashInfer's " + "sparse MLA decode API." + ) + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe() + # Per-tensor FP8 cache path scales. + if self.kv_cache_torch_dtype != torch.float8_e4m3fn: + return + fp8_q_scale = 1.0 + fp8_kv_scale = 1.0 + self.register_buffer( + "_flashinfer_fp8_q_scale", + torch.tensor([fp8_q_scale], dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_flashinfer_fp8_q_scale_inv", + torch.tensor([1.0 / fp8_q_scale], dtype=torch.float32), + persistent=False, + ) + self.register_buffer( + "_flashinfer_fp8_kv_scale", + torch.tensor([fp8_kv_scale], dtype=torch.float32), + persistent=False, + ) + # FlashInfer expects scalar scale arguments for this path. + self._flashinfer_fp8_bmm1_scale = self.scale * fp8_q_scale * fp8_kv_scale + self._flashinfer_fp8_bmm2_scale = fp8_kv_scale + + def _reserve_empty_forward_workspace(self) -> None: + self._get_workspace( + torch.device("cuda", torch.accelerator.current_device_index()) + ) + + def _forward_sparse_impl( + self, + q: torch.Tensor, + output: torch.Tensor, + flashmla_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + self_kv_cache: torch.Tensor | None, + swa_kv_cache: torch.Tensor, + swa_only: bool, + ) -> None: + num_decode_tokens = swa_metadata.num_decode_tokens + if swa_metadata.num_prefills > 0: + self._forward_prefill( + q=q[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if swa_metadata.num_decodes > 0: + self._forward_decode( + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + # Output may be padded to backend-supported head counts. + assert output.shape[0] == q.shape[0] and output.shape[-1] == q.shape[-1], ( + f"output buffer shape {output.shape} incompatible with q shape {q.shape}" + ) + assert output.shape[1] >= q.shape[1], ( + f"output heads {output.shape[1]} must be >= q heads {q.shape[1]}" + ) + # Per-tensor FP8 q produces a bf16 attention output. + expected_output_dtype = ( + torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype + ) + assert output.dtype == expected_output_dtype, ( + f"output dtype {output.dtype} must match expected {expected_output_dtype} " + f"for q dtype {q.dtype}" + ) + + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + if attn_metadata is None: + self._reserve_empty_forward_workspace() + output.zero_() + return + + assert isinstance(attn_metadata, dict) + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix) + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = self.compress_ratio <= 1 + # SWA-only layers don't allocate their own compressed KV cache. + self_kv_cache = self.kv_cache if not swa_only else None + swa_kv_cache = self.swa_cache_layer.kv_cache + + self._forward_sparse_impl( + q=q, + output=output, + flashmla_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + self_kv_cache=self_kv_cache, + swa_kv_cache=swa_kv_cache, + swa_only=swa_only, + ) + + def _prepare_query(self, q: torch.Tensor, output: torch.Tensor) -> torch.Tensor: + if self.kv_cache_torch_dtype == torch.float8_e4m3fn: + assert q.dtype == torch.float8_e4m3fn + q = q.to(torch.bfloat16) + else: + assert q.dtype == torch.bfloat16 + padded_heads = output.shape[1] + if q.shape[1] < padded_heads: + padded_query = q.new_zeros((q.shape[0], padded_heads, q.shape[2])) + padded_query[:, : q.shape[1], :] = q + q = padded_query + return q.contiguous() + + def _forward_decode( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + extra_sparse_indices = None + extra_sparse_lengths = None + if not swa_only: + if attn_metadata is None: + raise RuntimeError( + "Sparse MLA metadata is required for compressed layers." + ) + if swa_metadata.is_valid_token is None: + raise RuntimeError( + "SWA validity metadata is required for compressed layers." + ) + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if self.compress_ratio == 4: + if self.topk_indices_buffer is None: + raise RuntimeError( + "C4A decode requires top-k indices from the indexer." + ) + block_size = attn_metadata.block_size // self.compress_ratio + global_indices, extra_sparse_lengths = ( + compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + ) + extra_sparse_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + extra_sparse_indices = attn_metadata.c128a_global_decode_topk_indices + extra_sparse_lengths = attn_metadata.c128a_decode_topk_lens + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + assert swa_indices is not None + assert swa_lens is not None + q = self._prepare_query(q, output) + swa_cache = self._as_sparse_cache(self.swa_cache_layer.kv_cache) + extra_cache = self._as_sparse_cache(kv_cache) if kv_cache is not None else None + if extra_cache is not None and extra_sparse_indices is None: + raise RuntimeError( + "Compressed sparse MLA decode requires compressed sparse indices." + ) + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=q, + swa_kv_cache=swa_cache, + workspace_buffer=self._get_workspace(q.device), + sparse_indices=swa_indices, + compressed_kv_cache=extra_cache, + out=output, + bmm1_scale=self.scale, + sinks=self.attn_sink, + kv_layout="NHD", + swa_topk_lens=swa_lens, + extra_sparse_indices=extra_sparse_indices, + extra_sparse_topk_lens=extra_sparse_lengths, + ) + + def _forward_prefill( + self, + q: torch.Tensor, + compressed_k_cache: torch.Tensor | None, + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + swa_only = self.compress_ratio <= 1 + + num_prefills = swa_metadata.num_prefills + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + num_prefill_tokens = swa_metadata.num_prefill_tokens + + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + assert query_start_loc_cpu is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + local_topk_indices: torch.Tensor | None + if swa_only: + local_topk_indices = None + elif self.compress_ratio == 4: + if self.topk_indices_buffer is None: + raise RuntimeError( + "C4A prefill requires top-k indices from the indexer." + ) + local_topk_indices = self.topk_indices_buffer[ + num_decode_tokens : num_decode_tokens + num_prefill_tokens + ] + else: + if attn_metadata is None: + raise RuntimeError("C128A prefill metadata is missing.") + local_topk_indices = attn_metadata.c128a_prefill_topk_indices + + extra_sparse_indices: torch.Tensor | None = None + extra_sparse_lengths: torch.Tensor | None = None + if local_topk_indices is not None: + if attn_metadata is None: + raise RuntimeError("C4A prefill metadata is missing.") + if swa_metadata.token_to_req_indices is None: + raise RuntimeError("C4A prefill request mapping is missing.") + if swa_metadata.is_valid_token is None: + raise RuntimeError("C4A prefill validity metadata is missing.") + prefill_token_slice = slice( + num_decode_tokens, num_decode_tokens + num_prefill_tokens + ) + block_size = attn_metadata.block_size // self.compress_ratio + extra_sparse_indices, extra_sparse_lengths = ( + compute_global_topk_indices_and_lens( + local_topk_indices, + swa_metadata.token_to_req_indices[prefill_token_slice], + attn_metadata.block_table, + block_size, + swa_metadata.is_valid_token[prefill_token_slice], + ) + ) + + assert swa_metadata.prefill_swa_indices is not None + assert swa_metadata.prefill_swa_lens is not None + + q = self._prepare_query(q, output) + swa_kv_paged = self._as_sparse_cache(swa_k_cache) + if swa_only: + extra_kv_paged = None + else: + if compressed_k_cache is None: + raise RuntimeError( + "Compressed sparse MLA layers require their compressed KV cache." + ) + extra_kv_paged = self._as_sparse_cache(compressed_k_cache) + + num_chunks = ( + num_prefills + self.PREFILL_CHUNK_SIZE - 1 + ) // self.PREFILL_CHUNK_SIZE + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * self.PREFILL_CHUNK_SIZE + chunk_end = min(chunk_start + self.PREFILL_CHUNK_SIZE, num_prefills) + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + extra_sparse_indices_chunk = ( + extra_sparse_indices[query_start:query_end] + if extra_sparse_indices is not None + else None + ) + extra_sparse_lengths_chunk = ( + extra_sparse_lengths[query_start:query_end] + if extra_sparse_lengths is not None + else None + ) + + q_chunk = q[query_start:query_end] + swa_indices_chunk = swa_metadata.prefill_swa_indices[query_start:query_end] + swa_lens_chunk = swa_metadata.prefill_swa_lens[query_start:query_end] + if extra_kv_paged is not None and extra_sparse_indices_chunk is None: + raise RuntimeError( + "Compressed sparse MLA prefill requires compressed sparse indices." + ) + flashinfer_trtllm_batch_decode_sparse_mla_dsv4( + query=q_chunk, + swa_kv_cache=swa_kv_paged, + workspace_buffer=self._get_workspace(q.device), + sparse_indices=swa_indices_chunk, + compressed_kv_cache=extra_kv_paged, + out=output[query_start:query_end], + bmm1_scale=self.scale, + sinks=self.attn_sink, + kv_layout="NHD", + swa_topk_lens=swa_lens_chunk, + extra_sparse_indices=extra_sparse_indices_chunk, + extra_sparse_topk_lens=extra_sparse_lengths_chunk, + ) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 868fc3f5fd..aa60ad34ce 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -60,9 +60,11 @@ from vllm.model_executor.utils import set_weight_attrs from vllm.models.deepseek_v4.attention import DeepseekV4Attention from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, + DeepseekV4FlashInferSM120Attention, ) from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum @@ -736,14 +738,34 @@ class DeepseekV4MoE(nn.Module): def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: """Pick the CUDA sparse-MLA attention class for the configured backend. - An explicit ``--attention-backend FLASHINFER_MLA_SPARSE_DSV4`` selects the - FlashInfer TRTLLM-gen path; otherwise the FlashMLA path is used. + The generic CUDA backend selector does not instantiate DSv4 layers directly, + so map generic sparse-MLA choices to the DSv4-specialized attention class. + Without an explicit backend, SM12 defaults to FlashInfer while the other + CUDA arches keep the FlashMLA path. """ - if ( - vllm_config.attention_config.backend - == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4 + backend = vllm_config.attention_config.backend + device_capability = current_platform.get_device_capability() + if backend in ( + AttentionBackendEnum.FLASHINFER_MLA_SPARSE, + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120, ): + raise ValueError( + f"{backend.name} is not a DeepSeek V4 attention backend. " + "Use FLASHINFER_MLA_SPARSE_DSV4 for DeepSeek V4 FlashInfer " + "sparse MLA." + ) + if backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE_DSV4: + if device_capability is not None and device_capability.major == 12: + return DeepseekV4FlashInferSM120Attention return DeepseekV4FlashInferMLAAttention + if backend in ( + AttentionBackendEnum.FLASHMLA_SPARSE, + AttentionBackendEnum.FLASHMLA_SPARSE_DSV4, + ): + return DeepseekV4FlashMLAAttention + + if device_capability is not None and device_capability.major == 12: + return DeepseekV4FlashInferSM120Attention return DeepseekV4FlashMLAAttention diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index 2d2a42824c..136a96a45d 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -86,6 +86,10 @@ class DeepseekV4FlashMLABackend(AttentionBackend): def is_sparse(cls) -> bool: return True + @classmethod + def supports_sink(cls) -> bool: + return True + @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: return capability.major in [9, 10] diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 259077da35..dabe6058e4 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -11,7 +11,7 @@ import platform from collections.abc import Callable from datetime import timedelta from functools import cache, lru_cache, wraps -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, NamedTuple, TypeVar import torch from torch.distributed import PrefixStore, ProcessGroup @@ -31,6 +31,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.config.cache import CacheDType from vllm.config.kernel import IrOpPriorityConfig + from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.selector import AttentionSelectorConfig else: VllmConfig = None @@ -126,6 +127,11 @@ def _get_backend_priorities( AttentionBackendEnum.TRITON_MLA, *sparse_backends, ] + elif device_capability.major == 12: + return [ + AttentionBackendEnum.TRITON_MLA, + AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120, + ] else: return [ AttentionBackendEnum.FLASH_ATTN_MLA, @@ -153,6 +159,21 @@ def _get_backend_priorities( ] +def _backend_cls_path(backend_cls: type[AttentionBackend]) -> str: + module, qualname = backend_cls.full_cls_name() + return f"{module}.{qualname}" + + +def _get_attn_backend_class(backend: AttentionBackendEnum) -> type[AttentionBackend]: + return backend.get_class() + + +class _BackendCandidate(NamedTuple): + backend_class: type[AttentionBackend] + backend: AttentionBackendEnum + priority: int + + def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]: @wraps(fn) def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: @@ -340,7 +361,7 @@ class CudaPlatformBase(Platform): attn_selector_config: AttentionSelectorConfig, num_heads: int | None = None, ) -> tuple[ - list[tuple[AttentionBackendEnum, int]], + list[_BackendCandidate], dict[AttentionBackendEnum, tuple[int, list[str]]], ]: valid_backends_priorities = [] @@ -354,7 +375,7 @@ class CudaPlatformBase(Platform): ) for priority, backend in enumerate(backend_priorities): try: - backend_class = backend.get_class() + backend_class = _get_attn_backend_class(backend) invalid_reasons_i = backend_class.validate_configuration( device_capability=device_capability, **attn_selector_config._asdict(), @@ -364,7 +385,9 @@ class CudaPlatformBase(Platform): if invalid_reasons_i: invalid_reasons[backend] = (priority, invalid_reasons_i) else: - valid_backends_priorities.append((backend, priority)) + valid_backends_priorities.append( + _BackendCandidate(backend_class, backend, priority) + ) return valid_backends_priorities, invalid_reasons @@ -381,7 +404,7 @@ class CudaPlatformBase(Platform): # First try checking just the selected backend, if there is one. if selected_backend is not None: try: - backend_class = selected_backend.get_class() + backend_class = _get_attn_backend_class(selected_backend) invalid_reasons = backend_class.validate_configuration( device_capability=device_capability, **attn_selector_config._asdict(), @@ -395,7 +418,7 @@ class CudaPlatformBase(Platform): ) else: logger.info("Using %s backend.", selected_backend) - return selected_backend.get_path() + return _backend_cls_path(backend_class) # No selected backend or the selected backend is invalid, # so we try finding a valid backend. @@ -425,13 +448,13 @@ class CudaPlatformBase(Platform): # We have found some valid backends. Select the one with the # highest priority. - sorted_indices = sorted( - range(len(valid_backends_priorities)), - key=lambda i: valid_backends_priorities[i][1], + selected_candidate = min( + valid_backends_priorities, + key=lambda candidate: candidate.priority, ) - selected_index = sorted_indices[0] - selected_backend = valid_backends_priorities[selected_index][0] - selected_priority = valid_backends_priorities[selected_index][1] + selected_backend_class = selected_candidate.backend_class + selected_backend = selected_candidate.backend + selected_priority = selected_candidate.priority # If the user specified --block-size (but not --attention-backend), # check whether that constraint precluded any higher-priority backends. @@ -457,10 +480,14 @@ class CudaPlatformBase(Platform): logger.info_once( "Using %s attention backend out of potential backends: %s.", selected_backend.name, - "[" + ", ".join(f"'{b[0].name}'" for b in valid_backends_priorities) + "]", + "[" + + ", ".join( + f"'{candidate.backend.name}'" for candidate in valid_backends_priorities + ) + + "]", ) - return selected_backend.get_path() + return _backend_cls_path(selected_backend_class) @classmethod def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]: @@ -635,7 +662,11 @@ class CudaPlatformBase(Platform): @classmethod def support_deep_gemm(cls) -> bool: """Currently, only Hopper and Blackwell GPUs are supported.""" - return cls.is_device_capability(90) or cls.is_device_capability_family(100) + return ( + cls.is_device_capability(90) + or cls.is_device_capability_family(100) + or cls.is_device_capability_family(120) + ) @classmethod def is_integrated_gpu(cls, device_id: int = 0) -> bool: diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 1ddc93ff5e..0a1644bcc5 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -5,6 +5,7 @@ Users of vLLM should always import **only** these wrappers. """ +import contextlib import functools import importlib import os @@ -37,7 +38,10 @@ def should_auto_disable_deep_gemm(model_type: str | None) -> bool: """ if model_type is None: return False - if not current_platform.is_device_capability_family(100): + if not ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ): return False return model_type in _DEEPGEMM_BLACKWELL_EXCLUDED_MODEL_TYPES @@ -71,7 +75,10 @@ class DeepGemmQuantScaleFMT(Enum): cls._oracle_cache = ( # type: ignore cls.UE8M0 - if current_platform.is_device_capability_family(100) + if ( + current_platform.is_device_capability_family(100) + or current_platform.is_device_capability_family(120) + ) else cls.FLOAT32_CEIL_UE8M0 ) @@ -138,7 +145,15 @@ _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None +_get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( + None +) _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None +_pack_ue8m0_to_int_impl: Callable[..., Any] | None = None +_get_mn_major_tma_aligned_packed_ue8m0_tensor_impl: Callable[..., Any] | None = None +_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl: ( + Callable[..., Any] | None +) = None @functools.cache @@ -203,7 +218,11 @@ def _lazy_init() -> None: global _tf32_hc_prenorm_gemm_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl + global _get_theoretical_mk_alignment_for_contiguous_layout_impl global _transform_sf_into_required_layout_impl + global _pack_ue8m0_to_int_impl + global _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl + global _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl # fast path if ( _cublaslt_gemm_nt_impl is not None @@ -218,6 +237,9 @@ def _lazy_init() -> None: or _tf32_hc_prenorm_gemm_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None + or _pack_ue8m0_to_int_impl is not None + or _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None + or _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is not None ): return @@ -258,9 +280,19 @@ def _lazy_init() -> None: _get_mk_alignment_for_contiguous_layout_impl = getattr( _dg, "get_mk_alignment_for_contiguous_layout", None ) + _get_theoretical_mk_alignment_for_contiguous_layout_impl = getattr( + _dg, "get_theoretical_mk_alignment_for_contiguous_layout", None + ) _transform_sf_into_required_layout_impl = getattr( _dg, "transform_sf_into_required_layout", None ) + _pack_ue8m0_to_int_impl = getattr(_dg, "pack_ue8m0_to_int", None) + _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr( + _dg, "get_mn_major_tma_aligned_packed_ue8m0_tensor", None + ) + _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl = getattr( + _dg, "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", None + ) DeepGemmQuantScaleFMT.init_oracle_cache() @@ -280,7 +312,6 @@ def set_num_sms(num_sms: int) -> None: dg.set_num_sms(num_sms) -@functools.cache def get_mk_alignment_for_contiguous_layout() -> list[int]: _lazy_init() if _get_mk_alignment_for_contiguous_layout_impl is None: @@ -289,6 +320,70 @@ def get_mk_alignment_for_contiguous_layout() -> list[int]: return [mk_align_size, mk_align_size] +def get_theoretical_mk_alignment_for_contiguous_layout( + expected_m: int | None = None, + num_groups: int | None = None, +) -> int: + """Per-call optimal M alignment for grouped contiguous GEMMs. + + `expected_m` is the TOTAL routed tokens (sum across experts, typically + M × num_topk). `num_groups` is the number of experts on this rank. + The helper divides to recover per-expert em and picks an alignment based + on data-driven thresholds (see deep_gemm runtime.hpp comments). + + Older callers that omit `num_groups` are interpreted as passing already + per-expert em (legacy behaviour preserved for backward compat). + """ + _lazy_init() + if _get_theoretical_mk_alignment_for_contiguous_layout_impl is None: + return _missing() + if num_groups is None: + return _get_theoretical_mk_alignment_for_contiguous_layout_impl(expected_m) + if num_groups <= 0: + raise ValueError(f"num_groups must be positive, got {num_groups}") + try: + return _get_theoretical_mk_alignment_for_contiguous_layout_impl( + expected_m, num_groups + ) + except TypeError: + per_group_m = None if expected_m is None else cdiv(expected_m, num_groups) + return _get_theoretical_mk_alignment_for_contiguous_layout_impl(per_group_m) + + +def set_mk_alignment_for_contiguous_layout(value: int) -> None: + """Set DeepGEMM's BLOCK_M cap for grouped contiguous GEMMs. + + The DG heuristic constrains BLOCK_M ≤ this value when picking a kernel + layout. Use this in concert with `compute_aligned_M_and_alignment`'s + per-call alignment so the workspace's per-expert padding matches the + kernel's BLOCK_M; a mismatch leads to the scheduler reading the wrong + expert_id from `m_indices` at `m_block_idx * BLOCK_M` stride and + OOB-indexing the B-weights tensor (manifests as IMA under CUDA-graph + replay). + """ + _lazy_init() + dg = _import_deep_gemm() + if dg is None: + raise RuntimeError("DeepGEMM is not available") + dg.set_mk_alignment_for_contiguous_layout(value) + + +@contextlib.contextmanager +def mk_alignment_scope(value: int): + """Temporarily set DeepGEMM's BLOCK_M cap, restoring on exit. + + Use around a sequence of grouped-contiguous GEMM calls whose workspace + is padded to `value` (typically the per_call_align returned by + `compute_aligned_M_and_alignment`). + """ + prev = get_mk_alignment_for_contiguous_layout()[0] + set_mk_alignment_for_contiguous_layout(value) + try: + yield + finally: + set_mk_alignment_for_contiguous_layout(prev) + + def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor: """Wrapper for DeepGEMM's get_mn_major_tma_aligned_tensor""" _lazy_init() @@ -297,6 +392,48 @@ def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor: return _get_mn_major_tma_aligned_tensor_impl(x) +def pack_ue8m0_to_int(x: torch.Tensor) -> torch.Tensor: + """Pack 4 UE8M0 (uint8) scales into one int32. + + DeepGEMM's SM100/SM120 FP8/FP4 kernels accept either ``float32`` scales + (legacy format, 4 B/scale) or ``int32`` packed UE8M0 scales (1 B/scale + after 4:1 packing — 4× smaller than the legacy fp32 representation). + """ + _lazy_init() + if _pack_ue8m0_to_int_impl is None: + return _missing() + return _pack_ue8m0_to_int_impl(x) + + +def get_mn_major_tma_aligned_packed_ue8m0_tensor(x: torch.Tensor) -> torch.Tensor: + """Pack UE8M0 (uint8) → int32 with the MN-major TMA-aligned layout the + DeepGEMM kernels consume directly. 16× smaller than the fp32 legacy SF + format. Use for non-grouped 2D scale tensors. + """ + _lazy_init() + if _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None: + return _missing() + return _get_mn_major_tma_aligned_packed_ue8m0_tensor_impl(x) + + +def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + sf: torch.Tensor, + ks_tensor: torch.Tensor, + ks: list[int], + gran_k: int, +) -> torch.Tensor: + """Grouped (3D, expert-batched) variant of + ``get_mn_major_tma_aligned_packed_ue8m0_tensor``. Use for MoE weight + scale tensors of shape ``(num_experts, mn, k_scale)``. + """ + _lazy_init() + if _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl is None: + return _missing() + return _get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor_impl( + sf, ks_tensor, ks, gran_k + ) + + def cublaslt_gemm_nt(*args, **kwargs): _lazy_init() if _cublaslt_gemm_nt_impl is None: @@ -601,4 +738,8 @@ __all__ = [ "should_use_deepgemm_for_fp8_linear", "get_col_major_tma_aligned_tensor", "get_mk_alignment_for_contiguous_layout", + "get_theoretical_mk_alignment_for_contiguous_layout", + "pack_ue8m0_to_int", + "get_mn_major_tma_aligned_packed_ue8m0_tensor", + "get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", ] diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index e051827786..c8e6a8419b 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -72,11 +72,10 @@ def _missing(*_: Any, **__: Any) -> NoReturn: ) -def _missing_dsv4_sparse_mla(*_: Any, **__: Any) -> NoReturn: +def _missing_sparse_mla(*_: Any, **__: Any) -> NoReturn: raise RuntimeError( - "flashinfer.mla.trtllm_batch_decode_sparse_mla_dsv4 is not available. " - "Install a FlashInfer build that includes DeepSeek V4 sparse MLA " - "TRTLLM-GEN support." + "FlashInfer sparse MLA decode APIs are not available. " + "Install a FlashInfer build that includes sparse MLA decode support." ) @@ -149,14 +148,18 @@ flashinfer_b12x_fused_moe = _lazy_import_wrapper( trtllm_fp4_block_scale_moe = _lazy_import_wrapper( "flashinfer", "trtllm_fp4_block_scale_moe" ) -# DeepSeek V4 sparse MLA TRTLLM-GEN decode launcher (public wrapper). Handles -# the SWA + compressed KV pools, the concatenated sparse-index matrix, and -# per-tensor FP8 / BF16 inputs with BF16 output. -flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper( - "flashinfer.mla", - "trtllm_batch_decode_sparse_mla_dsv4", - fallback_fn=_missing_dsv4_sparse_mla, +flashinfer_trtllm_batch_decode_with_kv_cache_mla = _lazy_import_wrapper( + "flashinfer.decode", + "trtllm_batch_decode_with_kv_cache_mla", + fallback_fn=_missing_sparse_mla, ) +flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper( + "flashinfer.decode", + "trtllm_batch_decode_sparse_mla_dsv4", + fallback_fn=_missing_sparse_mla, +) + + # Special case for autotune since it returns a context manager autotune = _lazy_import_wrapper( "flashinfer.autotuner", @@ -209,6 +212,26 @@ def has_flashinfer_moe() -> bool: ) +@functools.cache +def has_flashinfer_sparse_mla_sm120() -> bool: + """Return ``True`` if FlashInfer sparse MLA decode support is available.""" + if not has_flashinfer(): + return False + try: + from flashinfer.autotuner import autotune + from flashinfer.decode import ( + trtllm_batch_decode_sparse_mla_dsv4, + trtllm_batch_decode_with_kv_cache_mla, + ) + except ImportError: + return False + return ( + callable(trtllm_batch_decode_sparse_mla_dsv4) + and callable(trtllm_batch_decode_with_kv_cache_mla) + and callable(autotune) + ) + + @functools.cache def has_flashinfer_cutedsl() -> bool: """Return ``True`` if FlashInfer cutedsl module is available.""" @@ -988,6 +1011,7 @@ __all__ = [ "flashinfer_b12x_fused_moe", "flashinfer_convert_sf_to_mma_layout", "trtllm_fp4_block_scale_moe", + "flashinfer_trtllm_batch_decode_with_kv_cache_mla", "flashinfer_trtllm_batch_decode_sparse_mla_dsv4", "autotune", "has_flashinfer_moe", diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index ebf607b65a..ccd70c6ca3 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -436,7 +436,7 @@ class CommonAttentionMetadata: positions: torch.Tensor | None = None """(num_actual_tokens,) token positions. Optional; set when the caller has positions available so that builders can pre-compute position-dependent - metadata (e.g. C128A topk indices for DeepSeek V4).""" + sparse metadata for DeepSeek V4 C128A layers.""" is_prefilling: torch.Tensor | None = None """(batch_size,) bool tensor: True if request is still in prefill phase diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 5547c62649..2a944d0618 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -1,23 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""FlashInfer MLA Sparse Attention Backend. - -This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k -for models like DeepSeek-V3.2 that use index-based sparse attention. - -For sparse MLA: -- block_tables shape changes from [batch_size, max_num_blocks] (dense) - to [batch_size, q_len_per_request, sparse_mla_top_k] (sparse) -- The sparse indices represent physical cache slot positions to attend to -- sparse_mla_top_k parameter must be set to the topk value -""" +"""FlashInfer sparse MLA attention backend.""" from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar import numpy as np import torch -from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla from vllm.config import VllmConfig from vllm.config.cache import CacheDType @@ -52,34 +41,13 @@ logger = init_logger(__name__) FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024 -class FlashInferMLASparseBackend(AttentionBackend): - """FlashInfer MLA backend with sparse attention support. - - This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k - for models like DeepSeek-V3.2 that use index-based sparse attention. - """ - - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] - supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ - "auto", - "float16", - "bfloat16", - "fp8", - "fp8_e4m3", - ] - - @staticmethod - def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: - return [32, 64] +class _FlashInferMLASparseBackendBase(AttentionBackend): + """Common metadata for concrete FlashInfer sparse MLA backends.""" @staticmethod def get_name() -> str: return "FLASHINFER_MLA_SPARSE" - @staticmethod - def get_impl_cls() -> type["FlashInferMLASparseImpl"]: - return FlashInferMLASparseImpl - @staticmethod def get_builder_cls() -> type["FlashInferMLASparseMetadataBuilder"]: return FlashInferMLASparseMetadataBuilder @@ -96,9 +64,29 @@ class FlashInferMLASparseBackend(AttentionBackend): def is_sparse(cls) -> bool: return True + +class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase): + """FlashInfer sparse MLA backend using the TRTLLM-gen launcher.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "float16", + "bfloat16", + "fp8", + "fp8_e4m3", + ] + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [32, 64] + + @staticmethod + def get_impl_cls() -> type[SparseMLAAttentionImpl]: + return FlashInferMLASparseImpl + @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: - # FlashInfer sparse MLA targets Blackwell (SM 10.x) return capability.major == 10 @classmethod @@ -114,10 +102,15 @@ class FlashInferMLASparseBackend(AttentionBackend): use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: - # FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192] from vllm.config import get_current_vllm_config vllm_config = get_current_vllm_config() + if kv_cache_dtype == "fp8_ds_mla": + return ( + "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype" + ) + + # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192]. if vllm_config.model_config is not None: hf_text_config = vllm_config.model_config.hf_text_config qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1) @@ -146,6 +139,102 @@ class FlashInferMLASparseBackend(AttentionBackend): return "HND" +class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase): + """FlashInfer sparse MLA backend for SM120.""" + + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16] + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ + "auto", + "fp8", + "fp8_e4m3", + "fp8_ds_mla", + ] + + @staticmethod + def get_name() -> str: + return "FLASHINFER_MLA_SPARSE_SM120" + + @staticmethod + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: + return [64, 256] + + @staticmethod + def get_impl_cls() -> type[SparseMLAAttentionImpl]: + from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import ( + FlashInferMLASparseSM120Impl, + ) + + return FlashInferMLASparseSM120Impl + + @classmethod + def supports_compute_capability(cls, capability: DeviceCapability) -> bool: + return capability.major == 12 + + @classmethod + def supports_combination( + cls, + head_size: int, + dtype: torch.dtype, + kv_cache_dtype: CacheDType | None, + block_size: int | None, + use_mla: bool, + has_sink: bool, + use_sparse: bool, + use_mm_prefix: bool, + device_capability: DeviceCapability, + ) -> str | None: + from vllm.config import get_current_vllm_config + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 + + if not has_flashinfer_sparse_mla_sm120(): + return ( + "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's " + "sparse MLA decode API" + ) + if dtype != torch.bfloat16: + return "dtype not supported" + if kv_cache_dtype not in ( + None, + "auto", + "fp8", + "fp8_e4m3", + "fp8_ds_mla", + ): + return "kv_cache_dtype not supported" + vllm_config = get_current_vllm_config() + if vllm_config.model_config is not None: + hf_text_config = vllm_config.model_config.hf_text_config + index_topk = getattr(hf_text_config, "index_topk", None) + if index_topk is None: + return ( + "FLASHINFER_MLA_SPARSE_SM120 requires a model with " + "index_topk config" + ) + if int(index_topk) != 2048: + return ( + "FLASHINFER_MLA_SPARSE_SM120 requires index_topk=2048; " + f"got {index_topk}" + ) + return None + + @staticmethod + def get_kv_cache_shape( + num_blocks: int, + block_size: int, + num_kv_heads: int, # assumed to be 1 for MLA + head_size: int, + cache_dtype_str: str = "auto", + ) -> tuple[int, ...]: + if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"): + # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE. + return (num_blocks, block_size, 656) + return (num_blocks, block_size, head_size) + + @classmethod + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None": + return None + + @dataclass class FlashInferMLASparseMetadata(AttentionMetadata): """Attention metadata for FlashInfer MLA Sparse backend.""" @@ -353,6 +442,8 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata if is_quantized_kv_cache(self.kv_cache_dtype): self.bmm2_scale *= layer._k_scale_float + from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla + o = trtllm_batch_decode_with_kv_cache_mla( query=q.unsqueeze(1), kv_cache=kv_c_and_k_pe_cache.unsqueeze(1), diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py new file mode 100644 index 0000000000..35b57b9c2b --- /dev/null +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SM120 implementation variant for ``FLASHINFER_MLA_SPARSE_SM120``.""" + +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.v1.attention.backend import ( + AttentionLayer, + AttentionType, + SparseMLAAttentionImpl, +) +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import ( + FlashInferMLASparseMetadata, + _get_workspace_buffer, +) +from vllm.v1.attention.backends.mla.sparse_utils import ( + triton_convert_req_index_to_global_index, +) + +if TYPE_CHECKING: + from vllm.model_executor.models.deepseek_v2 import Indexer + + +def _kv_scale_format_for_model(model_type: str | None) -> str: + if model_type is not None and model_type.startswith("glm"): + return "arbitrary_fp32" + return "pow2_fp32" + + +class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]): + """SM120 FlashInfer sparse-MLA implementation.""" + + def __init__( + self, + num_heads: int, + head_size: int, + scale: float, + num_kv_heads: int, + alibi_slopes: list[float] | None, + sliding_window: int | None, + kv_cache_dtype: str, + logits_soft_cap: float | None, + attn_type: str, + kv_sharing_target_layer_name: str | None, + indexer: "Indexer | None" = None, + **mla_args, + ) -> None: + if any([alibi_slopes, sliding_window, logits_soft_cap]): + raise NotImplementedError( + "FLASHINFER_MLA_SPARSE_SM120 does not support alibi_slopes / " + "sliding_window / logits_soft_cap" + ) + if attn_type != AttentionType.DECODER: + raise NotImplementedError( + "FLASHINFER_MLA_SPARSE_SM120 only supports decoder self-attention" + ) + + self.num_heads = num_heads + self.head_size = head_size + self.scale = float(scale) + self.num_kv_heads = num_kv_heads + self.kv_cache_dtype = kv_cache_dtype + if self.kv_cache_dtype != "fp8_ds_mla": + raise NotImplementedError( + "FLASHINFER_MLA_SPARSE_SM120 requires the packed fp8_ds_mla " + f"KV cache layout; got kv_cache_dtype={kv_cache_dtype!r}." + ) + + self.kv_lora_rank: int = mla_args["kv_lora_rank"] + self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"] + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"] + from vllm.config import get_current_vllm_config + + vllm_config = get_current_vllm_config() + model_type = None + if vllm_config.model_config is not None: + model_type = getattr( + vllm_config.model_config.hf_text_config, "model_type", None + ) + self.kv_scale_format = _kv_scale_format_for_model(model_type) + + assert indexer is not None, ( + "FLASHINFER_MLA_SPARSE_SM120 requires a sparse-MLA indexer " + "(model with index_topk in its config)." + ) + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer + from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120 + + if not has_flashinfer_sparse_mla_sm120(): + raise RuntimeError( + "FLASHINFER_MLA_SPARSE_SM120 requires FlashInfer's " + "sparse MLA decode API." + ) + assert self.topk_indices_buffer is not None + + self.supports_quant_query_input = False + self._workspace_buffer: torch.Tensor | None = None + + def forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: FlashInferMLASparseMetadata, + layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if isinstance(q, tuple): + q = torch.cat(q, dim=-1) + + num_actual_toks = q.shape[0] + + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[:num_actual_toks] + + topk_indices_physical = cast( + torch.Tensor, + triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[:num_actual_toks], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + ), + ) + + output = q.new_empty( + (num_actual_toks, self.num_heads, self.kv_lora_rank), + dtype=q.dtype, + ) + + if self._workspace_buffer is None: + self._workspace_buffer = _get_workspace_buffer(q.device) + + from vllm.utils.flashinfer import ( + flashinfer_trtllm_batch_decode_with_kv_cache_mla, + ) + + out = flashinfer_trtllm_batch_decode_with_kv_cache_mla( + query=q.unsqueeze(1), + kv_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(1), + workspace_buffer=self._workspace_buffer, + qk_nope_head_dim=self.qk_nope_head_dim, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=topk_indices_physical.unsqueeze(1), + seq_lens=None, + max_seq_len=attn_metadata.topk_tokens, + out=output.unsqueeze(1), + bmm1_scale=self.scale, + bmm2_scale=1.0, + sparse_mla_top_k=attn_metadata.topk_tokens, + kv_scale_format=self.kv_scale_format, + ) + return out.squeeze(1), None diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index a3fd39bed7..df23f34378 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -74,14 +74,14 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): # determines the SWA block size of 64 tokens per block. # TODO(yifan): make SWA block size automatically determined and configurable. self.block_size = 64 - # uint8: legacy FlashMLA UE8M0 paged layout. bfloat16 / float8_e4m3fn: - # FlashInfer contiguous full-cache layout. + # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn: + # contiguous full-cache layout. assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: - # FlashMLA's UE8M0 paged layout needs 576B alignment; FlashInfer's - # contiguous bf16/fp8 cache uses the natural element-size page. - is_flashmla = self.cache_config.cache_dtype == "fp8_ds_mla" + # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous + # bf16/fp8 cache uses the natural element-size page. + uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( block_size=self.block_size, num_kv_heads=1, @@ -89,7 +89,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - alignment=576 if is_flashmla else None, + alignment=576 if uses_fp8_ds_mla_layout else None, model_version="deepseek_v4", ) @@ -164,6 +164,11 @@ class DeepseekSparseSWAMetadata: token_to_req_indices: torch.Tensor | None = None # [num_tokens] decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, window_size] decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens] + # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill). + prefill_swa_indices: torch.Tensor | None = ( + None # [num_prefill_tokens, 1, window_size] + ) + prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens] # Number of decode/prefill requests/tokens (batch is reordered: decodes first) num_decodes: int = 0 @@ -343,6 +348,20 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): dtype=torch.int32, device=self.device, ) + # Allocated unconditionally — consumer picks paged-direct vs dequant + # at call time. + self.prefill_swa_indices = torch.zeros( + max_tokens, + 1, + self.window_size, + dtype=torch.int32, + device=self.device, + ) + self.prefill_swa_lens = torch.zeros( + max_tokens, + dtype=torch.int32, + device=self.device, + ) self.is_valid_token = torch.zeros( max_tokens, dtype=torch.bool, @@ -402,6 +421,29 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): block_table, block_table.stride(0), self.block_size, + token_offset=0, + TRITON_BLOCK_SIZE=1024, + ) + + # Prefill SWA indices live in paged coordinates. `token_offset` lets + # the kernel read is_valid_token / token_to_req_indices at absolute + # prefill positions while writing output starting at index 0. + if num_prefill_tokens > 0: + prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens] + prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens] + _compute_swa_indices_and_lens_kernel[(num_prefill_tokens,)]( + prefill_swa_indices, + prefill_swa_indices.stride(0), + prefill_swa_lens, + self.window_size, + query_start_loc, + seq_lens, + token_to_req_indices, + is_valid_token, + block_table, + block_table.stride(0), + self.block_size, + token_offset=num_decode_tokens, TRITON_BLOCK_SIZE=1024, ) @@ -431,6 +473,16 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): token_to_req_indices=token_to_req_indices, decode_swa_indices=self.decode_swa_indices[:num_decode_tokens], decode_swa_lens=self.decode_swa_lens[:num_decode_tokens], + prefill_swa_indices=( + self.prefill_swa_indices[:num_prefill_tokens] + if num_prefill_tokens > 0 + else None + ), + prefill_swa_lens=( + self.prefill_swa_lens[:num_prefill_tokens] + if num_prefill_tokens > 0 + else None + ), block_size=self.block_size, num_decodes=num_decodes, num_prefills=num_prefills, @@ -465,6 +517,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): num_decode_tokens == 0 or current_platform.is_rocm() or current_platform.is_xpu() + or current_platform.is_device_capability_family(120) ): return out for layer_type in self._layer_types: @@ -489,7 +542,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): Returns a dict of keyword arguments to pass to the DeepseekSparseSWAMetadata constructor. - Note: C128A topk indices are computed by the FlashMLASparse builder + Note: C128A sparse metadata is computed by the FlashMLASparse builder (which owns the C128A block_table), not here. """ result: dict[str, torch.Tensor | int | None] = {} @@ -539,10 +592,14 @@ def _compute_prefill_metadata_kernel( """Compute prefill gather_lens in a single pass.""" offset = tl.arange(0, BLOCK_SIZE) mask = offset < num_prefills + # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for + # masked-off lanes even though the load mask gates the actual read, so + # clamp the offset. Caller guarantees num_prefills > 0. + safe_offset = tl.minimum(offset, num_prefills - 1) - seq_len = tl.load(seq_lens_ptr + num_decodes + offset, mask=mask) - qsl_start = tl.load(query_start_loc_ptr + num_decodes + offset, mask=mask) - qsl_end = tl.load(query_start_loc_ptr + num_decodes + offset + 1, mask=mask) + seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask) + qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask) + qsl_end = tl.load(query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask) query_len = qsl_end - qsl_start prefix_len = seq_len - query_len @@ -551,7 +608,7 @@ def _compute_prefill_metadata_kernel( tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) -@triton.jit +@triton.jit(do_not_specialize=["token_offset"]) def _compute_swa_indices_and_lens_kernel( swa_indices_ptr, swa_indices_stride, @@ -564,12 +621,14 @@ def _compute_swa_indices_and_lens_kernel( block_table_ptr, block_table_stride, block_size, + token_offset, TRITON_BLOCK_SIZE: tl.constexpr, ): - token_idx = tl.program_id(0) + pid = tl.program_id(0) + token_idx = pid + token_offset is_valid = tl.load(is_valid_token_ptr + token_idx) if not is_valid: - tl.store(swa_lens_ptr + token_idx, 0) + tl.store(swa_lens_ptr + pid, 0) return req_idx = tl.load(token_to_req_indices_ptr + token_idx) @@ -586,7 +645,7 @@ def _compute_swa_indices_and_lens_kernel( end_pos = pos + 1 swa_len = end_pos - start_pos - tl.store(swa_lens_ptr + token_idx, swa_len) + tl.store(swa_lens_ptr + pid, swa_len) for i in range(0, window_size, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) @@ -602,7 +661,7 @@ def _compute_swa_indices_and_lens_kernel( slot_ids = tl.where(offset < swa_len, slot_ids, -1) tl.store( - swa_indices_ptr + token_idx * swa_indices_stride + offset, + swa_indices_ptr + pid * swa_indices_stride + offset, slot_ids, mask=offset < window_size, ) diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index bdaa752a60..5bbabc13dd 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -71,7 +71,11 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): ) FLASHINFER_MLA_SPARSE = ( "vllm.v1.attention.backends.mla.flashinfer_mla_sparse." - "FlashInferMLASparseBackend" + "FlashInferMLASparseTRTLLMBackend" + ) + FLASHINFER_MLA_SPARSE_SM120 = ( + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse." + "FlashInferMLASparseSM120Backend" ) TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend" CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend" diff --git a/vllm/v1/attention/ops/flashmla.py b/vllm/v1/attention/ops/flashmla.py index df04f5bf22..c84a495859 100644 --- a/vllm/v1/attention/ops/flashmla.py +++ b/vllm/v1/attention/ops/flashmla.py @@ -73,7 +73,7 @@ def is_flashmla_sparse_supported() -> tuple[bool, str | None]: ): return ( False, - "FlashMLA Sparse is only supported on Hopper and Blackwell devices.", + "FlashMLA Sparse is only supported on Hopper and Blackwell DC devices.", ) return True, None diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 0da845a067..3192b9aeaa 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from typing import Any import numpy as np import torch from vllm import PoolingParams, SamplingParams +from vllm.logger import init_logger from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import ( CachedRequestData, @@ -18,6 +20,134 @@ from vllm.v1.core.sched.output import ( from vllm.v1.request import Request from vllm.v1.worker.gpu.model_runner import GPUModelRunner +logger = init_logger(__name__) + + +def run_mixed_prefill_decode_warmup( + model_runner: GPUModelRunner, + worker_execute_model: Callable[[SchedulerOutput], Any], + worker_sample_tokens: Callable[[GrammarOutput | None], Any], + num_tokens: int, + *, + mixed_step_context: AbstractContextManager[object] | None = None, + req_id_prefix: str = "_v2_mixed_warmup", +) -> bool: + """Run a V2 mixed prefill+decode step through normal scheduler inputs.""" + if model_runner.is_pooling_model or num_tokens < 3: + return False + + decode_req_id = f"{req_id_prefix}_decode_" + prefill_req_id = f"{req_id_prefix}_prefill_" + decode_prompt_len = 2 + decode_scheduled_tokens = 1 + prefill_len = num_tokens - decode_scheduled_tokens + decode_token_ids = list(range(decode_prompt_len)) + prefill_token_ids = list(range(prefill_len)) + + kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups + num_kv_cache_groups = len(kv_cache_groups) + group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups] + decode_prefill_block_counts = [ + cdiv(decode_prompt_len, block_size) for block_size in group_block_sizes + ] + decode_block_counts = [ + cdiv(decode_prompt_len + decode_scheduled_tokens, block_size) + for block_size in group_block_sizes + ] + decode_block_deltas = [ + decode - prefill + for decode, prefill in zip(decode_block_counts, decode_prefill_block_counts) + ] + prefill_block_counts = [ + cdiv(prefill_len, block_size) for block_size in group_block_sizes + ] + required_blocks = sum(decode_block_counts) + sum(prefill_block_counts) + if model_runner.kv_cache_config.num_blocks <= required_blocks: + logger.warning( + "Skipping V2 mixed prefill+decode warmup because only %d KV blocks " + "are available for %d required warmup blocks.", + model_runner.kv_cache_config.num_blocks, + required_blocks, + ) + return False + + next_block_id = 1 + + def _alloc_blocks(num_blocks: int) -> list[int]: + nonlocal next_block_id + block_ids = list(range(next_block_id, next_block_id + num_blocks)) + next_block_id += num_blocks + return block_ids + + sampling_params = SamplingParams(max_tokens=2, temperature=0.0) + + decode_prefill_output = SchedulerOutput.make_empty() + decode_prefill_output.scheduled_new_reqs = [ + NewRequestData( + req_id=decode_req_id, + prompt_token_ids=decode_token_ids, + mm_features=[], + sampling_params=sampling_params, + pooling_params=None, + block_ids=tuple(_alloc_blocks(n) for n in decode_prefill_block_counts), + num_computed_tokens=0, + lora_request=None, + prefill_token_ids=decode_token_ids, + ), + ] + decode_prefill_output.num_scheduled_tokens = { + decode_req_id: decode_prompt_len, + } + decode_prefill_output.total_num_scheduled_tokens = decode_prompt_len + decode_prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + + decode_new_blocks = tuple(_alloc_blocks(n) for n in decode_block_deltas) + cached_decode_req = CachedRequestData.make_empty() + cached_decode_req.req_ids = [decode_req_id] + cached_decode_req.num_computed_tokens = [decode_prompt_len] + cached_decode_req.num_output_tokens = [1] + cached_decode_req.new_block_ids = [ + decode_new_blocks if any(decode_block_deltas) else None + ] + + mixed_output = SchedulerOutput.make_empty() + mixed_output.scheduled_cached_reqs = cached_decode_req + mixed_output.scheduled_new_reqs = [ + NewRequestData( + req_id=prefill_req_id, + prompt_token_ids=prefill_token_ids, + mm_features=[], + sampling_params=sampling_params, + pooling_params=None, + block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), + num_computed_tokens=0, + lora_request=None, + prefill_token_ids=prefill_token_ids, + ), + ] + mixed_output.num_scheduled_tokens = { + decode_req_id: decode_scheduled_tokens, + prefill_req_id: prefill_len, + } + mixed_output.total_num_scheduled_tokens = num_tokens + mixed_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + + cleanup_output = SchedulerOutput.make_empty() + cleanup_output.finished_req_ids = {decode_req_id, prefill_req_id} + + context = mixed_step_context or nullcontext() + model_runner.kv_connector.set_disabled(True) + try: + worker_execute_model(decode_prefill_output) + worker_sample_tokens(None) + with context: + worker_execute_model(mixed_output) + worker_sample_tokens(None) + worker_execute_model(cleanup_output) + finally: + model_runner.kv_connector.set_disabled(False) + return True + @torch.inference_mode() def warmup_kernels(