iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Attention][MLA] Support fp8_ds_mla KV cache for NoPE-512 models on SM90 (#59246)

Browse Source 
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Co-authored-by: Leoyzen <leoyzen@gmail.com>
This commit is contained in:
Nicolò Lucchesi
2026-10-05 13:12:00 +00:00
committed by GitHub
co-authored by Leoyzen
parent d0d6e5f3a2
commit ff53f32409
5 changed files with 82 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py
+25 -20
View File
@@ -7,9 +7,9 @@ v_head_dim=256, kv_lora_rank=512 -> head_size 512. These tests pin the
selection contract only (no GPU required):
* ``get_supported_head_sizes`` advertises 512 so the backend is a candidate;
* ``supports_combination`` admits exactly two shapes for 512 -- a quantized
DS-MLA cache format (zero-padded 576/656B envelope) and a rope-free bf16
cache on SM90 -- and rejects everything else;
* ``supports_combination`` admits 512 only for rope-free models on SM90, with
a bf16 cache or an fp8_ds_mla cache (zero RoPE slot), and rejects
everything else;
* adding 512 must NOT reorder the SM100 priority list for bf16, which is why
the bf16 arm is restricted to SM90.
"""
@@ -19,6 +19,7 @@ from types import SimpleNamespace
import pytest
import torch
from vllm.platforms import current_platform
from vllm.platforms.interface import DeviceCapability
from vllm.v1.attention.backends.mla.flashmla_sparse import (
QUANTIZED_DS_MLA_CACHE_FORMATS,
@@ -29,7 +30,7 @@ SM90 = DeviceCapability(major=9, minor=0)
SM100 = DeviceCapability(major=10, minor=0)
BF16_CACHE_DTYPES = (None, "auto", "bfloat16", "float16")
QUANTIZED_CACHE_DTYPES = ("fp8_ds_mla", "nvfp4_ds_mla")
NOPE_SM90_CACHE_DTYPES = (*BF16_CACHE_DTYPES, "fp8_ds_mla")
def _supports_combination(head_size, kv_cache_dtype, capability):
@@ -86,34 +87,26 @@ def test_quantized_ds_mla_formats_are_the_envelope_set():
assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS
@pytest.mark.parametrize("kv_cache_dtype", QUANTIZED_CACHE_DTYPES)
def test_nope_512_quantized_rejected_here(kv_cache_dtype):
# Quantized DS-MLA NoPE-512 is served by the zero-padded 576/656B
# envelope, wired up in a separate change; this backend must reject it.
reason = _supports_combination(512, kv_cache_dtype, SM90)
assert reason is not None
@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
def test_nope_512_bf16_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model):
@pytest.mark.parametrize("kv_cache_dtype", NOPE_SM90_CACHE_DTYPES)
def test_nope_512_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model):
assert _supports_combination(512, kv_cache_dtype, SM90) is None
@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
def test_nope_512_bf16_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model):
@pytest.mark.parametrize("kv_cache_dtype", NOPE_SM90_CACHE_DTYPES)
def test_nope_512_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model):
reason = _supports_combination(512, kv_cache_dtype, SM90)
assert reason is not None and "rope-free" in reason
@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
def test_nope_512_bf16_sm100_rejected(kv_cache_dtype, rope_free_model):
@pytest.mark.parametrize("kv_cache_dtype", NOPE_SM90_CACHE_DTYPES)
def test_nope_512_sm100_rejected(kv_cache_dtype, rope_free_model):
# The SM90 restriction is what keeps the SM100 priority order unchanged.
reason = _supports_combination(512, kv_cache_dtype, SM100)
assert reason is not None
@pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3"))
def test_nope_512_plain_fp8_rejected(kv_cache_dtype, rope_free_model):
@pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3", "nvfp4_ds_mla"))
def test_nope_512_other_quantized_rejected(kv_cache_dtype, rope_free_model):
reason = _supports_combination(512, kv_cache_dtype, SM90)
assert reason is not None
@@ -174,3 +167,15 @@ def test_sm100_bf16_512_priority_unchanged(num_heads):
def test_sm100_576_priority_unchanged():
order = _sparse_order("auto", 576, num_heads=32)
assert order[0] == "FLASHMLA_SPARSE", order
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA backends")
def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla():
"""The sparse backends ranked ahead of FLASHMLA_SPARSE for SM90 NoPE-512
reject fp8_ds_mla, so the auto-selector lands on FlashMLA."""
from vllm.v1.attention.backends.registry import AttentionBackendEnum
order = _sparse_order("fp8_ds_mla", 512, num_heads=64, capability=SM90)
for name in order[: order.index("FLASHMLA_SPARSE")]:
backend_cls = AttentionBackendEnum[name].get_class()
assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name
tests/v1/attention/test_mla_backends.py
+2 -3
View File
@@ -538,10 +538,9 @@ def create_and_prepopulate_kv_cache(
if fp8_attention:
if use_fp8_ds_mla:
kv_lora_rank = kv_c_contexts[0].shape[-1]
rope_dim = k_pe_contexts[0].shape[-1]
# 4 * 4: 4 float32 scale values for 128-element tiles
# 2 * rope_dim: 16-bit RoPE values
kv_entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
# 2 * 64: 16-bit RoPE values (zero-filled for NoPE models)
kv_entry_size = kv_lora_rank + 4 * 4 + 2 * 64
elif use_nvfp4_ds_mla:
kv_lora_rank = kv_c_contexts[0].shape[-1]
rope_dim = k_pe_contexts[0].shape[-1]
tests/v1/attention/test_sparse_mla_backends.py
+20 -6
View File
@@ -295,7 +295,8 @@ def _quantize_dequantize_fp8_ds_mla(
rope_dim = k_pe.shape[-1]
num_tokens = kv_c.shape[0]
num_blocks = max(1, math.ceil(num_tokens / block_size))
entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
# The 64-dim RoPE slot is always there; NoPE rows zero-fill it.
entry_size = kv_lora_rank + 4 * 4 + 2 * 64
tmp_cache = torch.zeros(
num_blocks, block_size, entry_size, dtype=torch.uint8, device=kv_c.device
@@ -413,6 +414,13 @@ def _quantize_dequantize_nvfp4_ds_mla(
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
@pytest.mark.parametrize("block_size", [32, 64])
@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
@pytest.mark.parametrize(
("total_num_heads", "qk_nope_head_dim", "qk_rope_head_dim", "v_head_dim"),
[
pytest.param(128, 128, 64, 128, id="deepseek"),
pytest.param(64, 256, 0, 256, id="glm53_flash_nope"),
],
)
def test_sparse_backend_decode_correctness(
default_vllm_config,
dist_init,
@@ -424,8 +432,18 @@ def test_sparse_backend_decode_correctness(
workspace_init,
q_scale: float,
k_scale: float,
total_num_heads: int,
qk_nope_head_dim: int,
qk_rope_head_dim: int,
v_head_dim: int,
monkeypatch,
):
if qk_rope_head_dim == 0 and (
backend_cls != FlashMLASparseBackend
or kv_cache_dtype not in ("auto", "fp8_ds_mla")
or torch.cuda.get_device_capability()[0] != 9
):
pytest.skip("NoPE sparse MLA is served by FlashMLA on SM90 only")
if (
batch_name == "large_q_pure_prefill"
and backend_cls == FlashMLASparseBackend
@@ -484,15 +502,10 @@ def test_sparse_backend_decode_correctness(
device = torch.device(DEVICE_TYPE)
dtype = torch.bfloat16
# Model hyper-parameters (kept intentionally small for the unit test)
total_num_heads = 128
# Compute per-rank heads for simulated TP
num_heads = max(1, total_num_heads // tensor_parallel_size)
kv_lora_rank = 512
qk_nope_head_dim = 128
qk_rope_head_dim = 64
v_head_dim = 128
head_size = kv_lora_rank + qk_rope_head_dim
topk_tokens = 128
@@ -3809,6 +3822,7 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b
topk_indices_buffer=topk_indices,
num_heads=2,
kv_lora_rank=1,
q_head_size=3,
index_group=None,
index_group_index=0,
dcp_world_size=1,
vllm/v1/attention/backend.py
+3 -2
View File
@@ -816,8 +816,9 @@ class AttentionImplBase(ABC, Generic[T]):
is_sparse: ClassVar[bool] = False
# Whether this impl provides a dense-MHA prefill path (forward_mha). Sparse
# impls without one run the top-k MQA path for all requests.
supports_dense_mha_prefill: ClassVar[bool] = True
# impls without one run the top-k MQA path for all requests. Impls may
# override it per instance, depending on the kv-cache layout.
supports_dense_mha_prefill: bool = True
# Required attributes that all impls should have
num_heads: int
vllm/v1/attention/backends/mla/flashmla_sparse.py
+32 -13
View File
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar
import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
from vllm import _custom_ops as ops
from vllm.config import VllmConfig, get_current_vllm_config
@@ -127,6 +128,9 @@ QUANTIZED_DS_MLA_CACHE_FORMATS: frozenset[str] = frozenset(
{"fp8_ds_mla", "nvfp4_ds_mla"}
)
# RoPE dims in an fp8_ds_mla row (the bf16 tail of the 656B layout).
FP8_DS_MLA_ROPE_DIM = 64
class FlashMLASparseBackend(AttentionBackend):
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
@@ -186,16 +190,16 @@ class FlashMLASparseBackend(AttentionBackend):
) -> str | None:
if head_size == 512:
# GLM5Next NoPE (qk_rope_head_dim == 0, kv_lora_rank == 512) has
# head_size 512 and is served here by the direct 512-wide bf16
# cache on SM90. Quantized DS-MLA caches need the zero-padded
# 576/656B envelope, which is wired up separately; plain fp8,
# SM100 bf16, and rope-carrying 512 models must fall through to
# head_size 512 and is served here on SM90, either by the direct
# 512-wide bf16 cache or by the 576 fp8_ds_mla layout with a zero
# RoPE slot (see FlashMLASparseImpl.q_head_size). Plain fp8,
# SM100, and rope-carrying 512 models must fall through to
# FlashInfer/TRITON.
if (
kv_cache_dtype in (None, "auto", "bfloat16", "float16")
kv_cache_dtype in (None, "auto", "bfloat16", "float16", "fp8_ds_mla")
and device_capability.major == 9
):
# Direct bf16 NoPE-512 is only correct for rope-free models.
# NoPE-512 is only correct for rope-free models.
# Precedent for reading hf_text_config in supports_combination:
# flashinfer_mla_sparse.py.
from vllm.config import get_current_vllm_config
@@ -210,8 +214,8 @@ class FlashMLASparseBackend(AttentionBackend):
)
else:
return (
"FLASHMLA_SPARSE supports head_size 512 only with bf16 "
"kv-cache on SM90 (NoPE), got "
"FLASHMLA_SPARSE supports head_size 512 only with bf16 or "
"fp8_ds_mla kv-cache on SM90 (NoPE), got "
f"kv_cache_dtype={kv_cache_dtype}, "
f"capability={device_capability}"
)
@@ -732,6 +736,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
)
self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads)
# The kernels read queries as wide as the cache rows. A NoPE model on
# fp8_ds_mla keeps the DeepSeek 576 layout: the cache writer zero-fills
# the RoPE slot and forward_mqa zero-pads q_pe to match (q_pe . 0 = 0).
self.q_head_size = head_size
if kv_cache_dtype == "fp8_ds_mla" and self.qk_rope_head_dim == 0:
self.q_head_size = self.kv_lora_rank + FP8_DS_MLA_ROPE_DIM
# Dense-MHA prefill would gather the cache into a
# kv_lora_rank-wide workspace, too narrow for the 576 rows.
self.supports_dense_mha_prefill = False
vllm_config = get_current_vllm_config()
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
q_concat_heads = num_heads
@@ -741,7 +755,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
// self.prefill_padding
* self.prefill_padding
)
q_concat_shape = (max_tokens, q_concat_heads, head_size)
q_concat_shape = (max_tokens, q_concat_heads, self.q_head_size)
if is_quantized_kv_cache(kv_cache_dtype):
assert kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS, (
"FlashMLA Sparse Attention backend only supports the "
@@ -777,11 +791,11 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
# PCP+DCP upconverts this rank's KV shard, then all-gathers the
# shards into a workspace of the full prefill size.
shard_rows //= parallel_config.decode_context_parallel_size
self.prefill_workspace_shape = (shard_rows, head_size)
self.prefill_workspace_shape = (shard_rows, self.q_head_size)
self.workspace_specs.append((self.prefill_workspace_shape, torch.bfloat16))
if self.pcp_dcp_kv_gather:
self.workspace_specs.append(
((prefill_workspace_size, head_size), torch.bfloat16)
((prefill_workspace_size, self.q_head_size), torch.bfloat16)
)
prefill_query_heads = num_heads
if self.pcp_dcp_kv_gather and self.dcp_world_size > self.pcp_world_size:
@@ -792,7 +806,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
self.workspace_specs.extend(
(shape, torch.bfloat16)
for shape in (
(max_tokens, padded_prefill_query_heads, head_size),
(max_tokens, padded_prefill_query_heads, self.q_head_size),
(max_tokens, padded_prefill_query_heads, self.kv_lora_rank),
(max_tokens, prefill_query_heads, self.kv_lora_rank),
)
@@ -1370,11 +1384,16 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
if q_pe.size(-1) == 0:
# NoPE (GLM5Next): concat_mla_q requires rope_dim == 64,
# copy directly into the head-padded buffer instead.
q[:, : ql_nope.shape[1]].copy_(ql_nope)
q[:, : ql_nope.shape[1], : self.kv_lora_rank].copy_(ql_nope)
if self.q_head_size > self.kv_lora_rank:
q[..., self.kv_lora_rank :].zero_()
else:
ops.concat_mla_q(ql_nope, q_pe, q)
else:
actual_num_heads = q.shape[1]
if q.shape[-1] < self.q_head_size:
# DCP gathers the unpadded NoPE query.
q = F.pad(q, (0, self.q_head_size - q.shape[-1]))
num_actual_toks = q.shape[0]