929 lines
37 KiB
Python
929 lines
37 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""DeepSeek V4 FlashInfer sparse MLA backend."""
|
|
|
|
from typing import TYPE_CHECKING, ClassVar, cast
|
|
|
|
import torch
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.config.cache import CacheDType
|
|
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,
|
|
deep_gemm_fp8_o_proj,
|
|
)
|
|
from vllm.models.deepseek_v4.sparse_mla import (
|
|
DeepseekV4FlashMLAMetadata,
|
|
DeepseekV4SparseMLABackend,
|
|
DeepseekV4SparseMLAMetadataBuilder,
|
|
)
|
|
from vllm.platforms.interface import DeviceCapability
|
|
from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
|
|
from vllm.v1.attention.backend import AttentionCGSupport, MultipleOf
|
|
from vllm.v1.attention.backends.mla.compressor_utils import (
|
|
get_dspark_swa_index_width,
|
|
)
|
|
from vllm.v1.attention.backends.mla.sparse_swa import (
|
|
DeepseekSparseSWABackend,
|
|
DeepseekSparseSWAMetadataBuilder,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
|
|
|
|
_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
|
|
_flashinfer_dsv4_workspace_by_device: dict[torch.device, torch.Tensor] = {}
|
|
|
|
|
|
def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:
|
|
workspace = _flashinfer_dsv4_workspace_by_device.get(device)
|
|
if workspace is None:
|
|
workspace = torch.zeros(
|
|
_FLASHINFER_DSV4_WORKSPACE_BUFFER_SIZE,
|
|
dtype=torch.uint8,
|
|
device=device,
|
|
)
|
|
_flashinfer_dsv4_workspace_by_device[device] = workspace
|
|
return workspace
|
|
|
|
|
|
def _packed_block_span(pool: torch.Tensor) -> int:
|
|
"""Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): ==
|
|
block_size for unpacked KV, larger when packed (#44577). Raises if not
|
|
token-aligned."""
|
|
block_stride = pool.stride(0)
|
|
token_stride = pool.stride(-2)
|
|
if block_stride % token_stride != 0:
|
|
raise NotImplementedError(
|
|
"FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride "
|
|
f"({block_stride}) to be a multiple of the per-token stride "
|
|
f"({token_stride}); this layout is not supported yet."
|
|
)
|
|
return block_stride // token_stride
|
|
|
|
|
|
# Sparse MLA h_q counts accepted natively (flashinfer>=0.6.14, #3545).
|
|
_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128)
|
|
|
|
|
|
def _pad_to_supported_q_heads(num_heads: int) -> int:
|
|
for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS:
|
|
if num_heads <= supported:
|
|
return supported
|
|
raise ValueError(
|
|
f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
|
|
"(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})."
|
|
)
|
|
|
|
|
|
def _required_sm120_sparse_topk(vllm_config: VllmConfig, window_size: int) -> int:
|
|
"""Return the SM120 DSV4 SWA specialization needed by this model."""
|
|
if not vllm_config.attention_config.use_non_causal:
|
|
return window_size
|
|
speculative_config = vllm_config.speculative_config
|
|
if speculative_config is None:
|
|
return window_size
|
|
return get_dspark_swa_index_width(
|
|
window_size,
|
|
speculative_config.num_speculative_tokens,
|
|
)
|
|
|
|
|
|
class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
|
|
"""FlashInfer backend using the DSv4 sparse metadata/cache layout.
|
|
|
|
Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``
|
|
"""
|
|
|
|
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
|
|
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_builder_cls() -> type["DeepseekV4FlashInferSparseMLAMetadataBuilder"]:
|
|
return DeepseekV4FlashInferSparseMLAMetadataBuilder
|
|
|
|
|
|
class DeepseekV4FlashInferSparseMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
|
|
"""Varlen-capable metadata builder for the FlashInfer sparse MLA backend."""
|
|
|
|
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
|
|
|
|
|
|
class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
|
|
"""SWA metadata for the FlashInfer sparse decode path (varlen decode)."""
|
|
|
|
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
|
|
|
|
|
|
class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
|
|
@staticmethod
|
|
def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
|
|
return DeepseekSparseSWAFlashInferMetadataBuilder
|
|
|
|
|
|
class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
|
|
"""FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""
|
|
|
|
backend_cls = DeepseekV4FlashInferMLASparseBackend
|
|
swa_backend_cls = DeepseekSparseSWAFlashInferBackend
|
|
use_fp8_ds_mla_layout: ClassVar[bool] = False
|
|
|
|
@classmethod
|
|
def get_padded_num_q_heads(cls, num_heads: int) -> int:
|
|
return _pad_to_supported_q_heads(num_heads)
|
|
|
|
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)
|
|
self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe()
|
|
# Per-tensor FP8 scale buffers + precomputed scalar BMM scales. Only the
|
|
# per-tensor FP8 cache path consumes these; bf16 reads ``self.scale``.
|
|
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,
|
|
)
|
|
# 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
|
|
|
|
def forward_mqa(
|
|
self,
|
|
q: torch.Tensor,
|
|
kv: torch.Tensor,
|
|
positions: torch.Tensor,
|
|
output: torch.Tensor,
|
|
) -> None:
|
|
# The TRTLLM-gen kernel requires h_q in {64, 128}, so the output buffer
|
|
# is allocated at the padded head count while q arrives at the local
|
|
# head count; _forward pads q to match before the launcher.
|
|
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:
|
|
# Warmup dummy run: FlashInfer reads the cache directly and lazily
|
|
# allocates its workspace, so nothing to reserve here.
|
|
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(
|
|
q=q,
|
|
kv_cache=self_kv_cache,
|
|
swa_k_cache=swa_kv_cache,
|
|
swa_metadata=swa_metadata,
|
|
attn_metadata=flashmla_metadata,
|
|
swa_only=swa_only,
|
|
output=output,
|
|
)
|
|
|
|
def _build_sparse_index_metadata(
|
|
self,
|
|
kv_cache: torch.Tensor | None,
|
|
swa_k_cache: torch.Tensor,
|
|
swa_metadata: "DeepseekSparseSWAMetadata",
|
|
attn_metadata: DeepseekV4FlashMLAMetadata | None,
|
|
swa_only: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Build the combined sparse-index tensors for the mixed batch.
|
|
|
|
Returns ``(compressed_kv_cache, seq_lens, sparse_indices,
|
|
sparse_topk_lens)``.
|
|
"""
|
|
num_decodes = swa_metadata.num_decodes
|
|
num_prefills = swa_metadata.num_prefills
|
|
num_decode_tokens = swa_metadata.num_decode_tokens
|
|
num_prefill_tokens = swa_metadata.num_prefill_tokens
|
|
num_reqs = num_decodes + num_prefills
|
|
num_tokens = num_decode_tokens + num_prefill_tokens
|
|
|
|
assert swa_metadata.seq_lens is not None
|
|
assert swa_metadata.query_start_loc is not None
|
|
assert swa_metadata.token_to_req_indices is not None
|
|
assert swa_metadata.decode_swa_indices is not None
|
|
assert swa_metadata.block_table is not None
|
|
|
|
decode_swa_indices = swa_metadata.decode_swa_indices.reshape(
|
|
num_decode_tokens, swa_metadata.decode_swa_width
|
|
)
|
|
decode_compressed_topk_lens = None
|
|
decode_compressed_indices_are_local = False
|
|
decode_is_valid_token = None
|
|
|
|
if swa_only:
|
|
assert self.topk_indices_buffer is not None
|
|
compressed_kv_cache = swa_k_cache
|
|
decode_compressed_indices = None
|
|
prefill_topk_indices = self.topk_indices_buffer[
|
|
num_decode_tokens:num_tokens, :0
|
|
]
|
|
compressed_block_table = None
|
|
compressed_block_size = swa_metadata.block_size
|
|
top_k = 0
|
|
else:
|
|
assert kv_cache is not None
|
|
assert attn_metadata is not None
|
|
compressed_kv_cache = kv_cache
|
|
compressed_block_table = attn_metadata.block_table[:num_reqs]
|
|
compressed_block_size = attn_metadata.block_size // self.compress_ratio
|
|
|
|
if self.compress_ratio == 4:
|
|
assert self.topk_indices_buffer is not None
|
|
if num_prefill_tokens > 0:
|
|
prefill_topk_indices = self.topk_indices_buffer[
|
|
num_decode_tokens:num_tokens
|
|
]
|
|
top_k = prefill_topk_indices.shape[-1]
|
|
else:
|
|
prefill_topk_indices = self.topk_indices_buffer[:0, :0]
|
|
top_k = 0
|
|
|
|
decode_compressed_indices_are_local = True
|
|
assert swa_metadata.is_valid_token is not None
|
|
decode_is_valid_token = swa_metadata.is_valid_token[:num_decode_tokens]
|
|
if num_decode_tokens > 0:
|
|
decode_compressed_indices = self.topk_indices_buffer[
|
|
:num_decode_tokens
|
|
]
|
|
else:
|
|
# Keep the logical width aligned with the mixed-batch case so
|
|
# pure-prefill steps reuse the same Triton specialization.
|
|
decode_compressed_indices = prefill_topk_indices[:0]
|
|
else:
|
|
if num_prefill_tokens > 0:
|
|
assert attn_metadata.c128a_prefill_topk_indices is not None
|
|
prefill_topk_indices = attn_metadata.c128a_prefill_topk_indices
|
|
top_k = prefill_topk_indices.shape[-1]
|
|
else:
|
|
prefill_topk_indices = decode_swa_indices[:0, :0]
|
|
top_k = 0
|
|
|
|
if num_decode_tokens > 0:
|
|
assert attn_metadata.c128a_global_decode_topk_indices is not None
|
|
assert attn_metadata.c128a_decode_topk_lens is not None
|
|
decode_compressed_indices = (
|
|
attn_metadata.c128a_global_decode_topk_indices.view(
|
|
num_decode_tokens, -1
|
|
)
|
|
)
|
|
decode_compressed_topk_lens = attn_metadata.c128a_decode_topk_lens
|
|
if num_prefill_tokens == 0:
|
|
prefill_topk_indices = decode_compressed_indices[:0, :0]
|
|
else:
|
|
decode_compressed_indices = prefill_topk_indices[:0]
|
|
decode_compressed_topk_lens = swa_metadata.seq_lens[:0]
|
|
|
|
query_start_loc = swa_metadata.query_start_loc[: num_reqs + 1]
|
|
seq_lens = swa_metadata.seq_lens[:num_reqs]
|
|
assert seq_lens.dtype == torch.int32
|
|
# cache for SWA-only and C128A that build the same mixed sparse indices
|
|
# C4A stays uncached.
|
|
cache_key = (
|
|
"swa_only"
|
|
if swa_only
|
|
else ("c128a" if self.compress_ratio == 128 else "c4a")
|
|
)
|
|
cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None)
|
|
if cached_sparse is None:
|
|
swa_block_span = _packed_block_span(swa_k_cache)
|
|
compressed_block_span = _packed_block_span(compressed_kv_cache)
|
|
sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(
|
|
decode_swa_indices,
|
|
decode_compressed_indices,
|
|
decode_compressed_topk_lens,
|
|
prefill_topk_indices[:num_prefill_tokens],
|
|
query_start_loc,
|
|
seq_lens,
|
|
swa_metadata.token_to_req_indices[:num_tokens],
|
|
swa_metadata.block_table[:num_reqs],
|
|
swa_metadata.block_size,
|
|
compressed_block_table,
|
|
compressed_block_size,
|
|
self.window_size,
|
|
self.compress_ratio,
|
|
top_k,
|
|
decode_compressed_indices_are_local=decode_compressed_indices_are_local,
|
|
decode_is_valid_token=decode_is_valid_token,
|
|
swa_block_span=swa_block_span,
|
|
compressed_block_span=compressed_block_span,
|
|
prefill_left_visible=swa_metadata.prefill_left_visible,
|
|
prefill_right_visible=swa_metadata.prefill_right_visible,
|
|
# getattr for tests that bypass __init__ via object.__new__.
|
|
max_image_tokens=getattr(self, "max_image_tokens", 0),
|
|
)
|
|
if cache_key != "c4a":
|
|
swa_metadata.flashinfer_sparse_index_cache[cache_key] = (
|
|
sparse_indices,
|
|
sparse_topk_lens,
|
|
)
|
|
else:
|
|
sparse_indices, sparse_topk_lens = cached_sparse
|
|
return compressed_kv_cache, seq_lens, sparse_indices, sparse_topk_lens
|
|
|
|
def _forward(
|
|
self,
|
|
q: torch.Tensor,
|
|
kv_cache: torch.Tensor | None,
|
|
swa_k_cache: torch.Tensor,
|
|
swa_metadata: "DeepseekSparseSWAMetadata",
|
|
attn_metadata: DeepseekV4FlashMLAMetadata | None,
|
|
swa_only: bool,
|
|
output: torch.Tensor,
|
|
) -> None:
|
|
assert self.kv_cache_torch_dtype in (torch.bfloat16, torch.float8_e4m3fn)
|
|
num_decodes = swa_metadata.num_decodes
|
|
num_prefills = swa_metadata.num_prefills
|
|
num_decode_tokens = swa_metadata.num_decode_tokens
|
|
num_prefill_tokens = swa_metadata.num_prefill_tokens
|
|
num_reqs = num_decodes + num_prefills
|
|
num_tokens = num_decode_tokens + num_prefill_tokens
|
|
if num_tokens == 0:
|
|
return
|
|
|
|
(
|
|
compressed_kv_cache,
|
|
seq_lens,
|
|
sparse_indices,
|
|
sparse_topk_lens,
|
|
) = self._build_sparse_index_metadata(
|
|
kv_cache=kv_cache,
|
|
swa_k_cache=swa_k_cache,
|
|
swa_metadata=swa_metadata,
|
|
attn_metadata=attn_metadata,
|
|
swa_only=swa_only,
|
|
)
|
|
|
|
# CUDA graph execution can pad q/output past the scheduled token count;
|
|
# restrict to the real tokens (the launcher validates sparse indices).
|
|
query = q[:num_tokens]
|
|
output = output[:num_tokens]
|
|
bmm1_scale: float | torch.Tensor = self.scale
|
|
bmm2_scale: float | torch.Tensor = 1.0
|
|
if self.kv_cache_torch_dtype == torch.float8_e4m3fn:
|
|
assert query.dtype == torch.float8_e4m3fn
|
|
bmm1_scale = self._flashinfer_fp8_bmm1_scale
|
|
bmm2_scale = self._flashinfer_fp8_bmm2_scale
|
|
else:
|
|
assert query.dtype == torch.bfloat16
|
|
query = query.contiguous()
|
|
|
|
# The TRTLLM-gen sparse-MLA kernel requires h_q in {64, 128}; zero-pad
|
|
# the query heads to the allocated output head count. Padded heads attend
|
|
# to the shared KV and are sliced off downstream (output is padded too).
|
|
padded_heads = output.shape[1]
|
|
if query.shape[1] < padded_heads:
|
|
padded_query = query.new_zeros(
|
|
(query.shape[0], padded_heads, query.shape[2])
|
|
)
|
|
padded_query[:, : query.shape[1], :] = query
|
|
query = padded_query
|
|
|
|
workspace = _get_flashinfer_dsv4_workspace(q.device)
|
|
query_start_loc = swa_metadata.query_start_loc
|
|
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 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]
|
|
flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
|
|
query=query[:num_decode_tokens],
|
|
swa_kv_cache=swa_k_cache,
|
|
workspace_buffer=workspace,
|
|
sparse_indices=sparse_indices[:num_decode_tokens],
|
|
compressed_kv_cache=compressed_kv_cache,
|
|
sparse_topk_lens=sparse_topk_lens[:num_decode_tokens],
|
|
seq_lens=seq_lens[:num_decodes],
|
|
out=output[:num_decode_tokens],
|
|
bmm1_scale=bmm1_scale,
|
|
bmm2_scale=bmm2_scale,
|
|
sinks=self.attn_sink,
|
|
cum_seq_lens_q=decode_cu,
|
|
max_q_len=swa_metadata.max_decode_query_len,
|
|
)
|
|
|
|
if num_prefill_tokens > 0:
|
|
# The prefill query view re-anchors at offset 0, so rebase the
|
|
# cumulative query offsets to start at 0.
|
|
prefill_cu = (
|
|
query_start_loc[num_decodes : num_reqs + 1]
|
|
- query_start_loc[num_decodes]
|
|
)
|
|
prefill_cu_cpu = query_start_loc_cpu[num_decodes : num_reqs + 1]
|
|
prefill_lens_cpu = prefill_cu_cpu[1:] - prefill_cu_cpu[:-1]
|
|
flashinfer_trtllm_batch_decode_sparse_mla_dsv4(
|
|
query=query[num_decode_tokens:num_tokens],
|
|
swa_kv_cache=swa_k_cache,
|
|
workspace_buffer=workspace,
|
|
sparse_indices=sparse_indices[num_decode_tokens:num_tokens],
|
|
compressed_kv_cache=compressed_kv_cache,
|
|
sparse_topk_lens=sparse_topk_lens[num_decode_tokens:num_tokens],
|
|
seq_lens=seq_lens[num_decodes:num_reqs],
|
|
out=output[num_decode_tokens:num_tokens],
|
|
bmm1_scale=bmm1_scale,
|
|
bmm2_scale=bmm2_scale,
|
|
sinks=self.attn_sink,
|
|
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
|
|
swa_backend_cls = DeepseekSparseSWAFlashInferBackend
|
|
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:
|
|
return _pad_to_supported_q_heads(num_heads)
|
|
|
|
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, vllm_config: VllmConfig, *args, **kwargs) -> None:
|
|
super().__init__(vllm_config, *args, **kwargs)
|
|
from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120_config
|
|
|
|
required_topk = _required_sm120_sparse_topk(vllm_config, self.window_size)
|
|
if not has_flashinfer_sparse_mla_sm120_config(self.padded_heads, required_topk):
|
|
raise RuntimeError(
|
|
"FLASHINFER_MLA_SPARSE_DSV4 on SM120 requires a FlashInfer "
|
|
"DSV4 sparse MLA decode specialization for "
|
|
f"(num_q_heads={self.padded_heads}, top_k={required_topk}). "
|
|
"Install a FlashInfer build containing "
|
|
"flashinfer-ai/flashinfer#4380."
|
|
)
|
|
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,
|
|
)
|