From 66a2ae784a38687077fad383a6f6b0467205fadf Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 21 Aug 2026 10:05:31 +0000 Subject: [PATCH 1/6] refactor: separate sparse FlashMLA attention --- .../backends/cuda/attention/__init__.py | 8 +- .../pytorch/backends/cuda/attention/mla.py | 290 ++++-------------- .../backends/cuda/attention/sparse_mla.py | 240 +++++++++++++++ lmdeploy/pytorch/models/deepseek_v32.py | 3 +- tests/pytorch/kernel/test_mla_attention.py | 76 +++-- 5 files changed, 351 insertions(+), 266 deletions(-) create mode 100644 lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py diff --git a/lmdeploy/pytorch/backends/cuda/attention/__init__.py b/lmdeploy/pytorch/backends/cuda/attention/__init__.py index c703c065e9..76780ea3d1 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/__init__.py +++ b/lmdeploy/pytorch/backends/cuda/attention/__init__.py @@ -72,7 +72,7 @@ class TritonAttentionBuilder(AttentionBuilder[TritonAttentionMetadata]): """Triton attention builder. This builder selects the appropriate attention implementation based on: - 1. use_flash_mla: Use FlashMLAImpl for MLA models + 1. use_flash_mla: Use dense or sparse FlashMLA for MLA models 2. enable_fa3: Use FA3Impl if FA3 is available and supported 3. Default: Use TritonAttentionImpl as fallback """ @@ -89,6 +89,7 @@ def build( logit_softcapping: float = 0.0, causal: bool = True, use_flash_mla: bool = False, + mla_index_topk: int | None = None, learnable_sink: bool = False, block_sparse_size: int = 1, **kwargs, @@ -106,6 +107,7 @@ def build( logit_softcapping: Logit softcapping value (for Gemma 2). causal: Whether to use causal attention. use_flash_mla: Whether to use Flash MLA implementation. + mla_index_topk: Sparse MLA top-k width, or ``None`` for dense MLA. learnable_sink: Whether to use learnable sink tokens. block_sparse_size: Block sparse attention size. **kwargs: Additional arguments. @@ -132,6 +134,10 @@ def build( enable_fa3 = _enable_fa3(alibi, learnable_sink, block_sparse_size, head_size) if use_flash_mla is True: + if mla_index_topk is not None: + logger.debug('Build FlashMLASparseImpl Attention') + from .sparse_mla import FlashMLASparseImpl + return FlashMLASparseImpl(mla_index_topk=mla_index_topk, **common_args) logger.debug('Build FlashMLAImpl Attention') from .mla import FlashMLAImpl return FlashMLAImpl(use_fa3=use_fa3, **common_args) diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index ed26803f91..806b066072 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -1,6 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. -import functools from collections.abc import Hashable from dataclasses import dataclass from typing import Any @@ -106,22 +105,22 @@ class FlashMLAAttentionMetaBuilder( """Build metadata requested by one selected FlashMLA configuration.""" num_attention_heads: int + index_topk: int | None = None @property def key(self) -> Hashable: - return (type(self), self.num_attention_heads) + return (type(self), self.num_attention_heads, self.index_topk) def build(self, step_context, sequence_metadata) -> FlashMLAAttentionMetadata: if not step_context.is_decoding: return FlashMLAAttentionMetadata() batch_size = sequence_metadata.q_seqlens.size(0) - model_config = step_context.model_config return build_flash_mla_metadata( sequence_metadata, num_attention_heads=self.num_attention_heads, decoding_query_len=step_context.input_ids.size(1) // batch_size, - is_fp8_kvcache=model_config.use_mla_fp8_cache, - index_topk=model_config.mla_index_topk, + is_fp8_kvcache=step_context.model_config.use_mla_fp8_cache, + index_topk=self.index_topk, ) def apply_legacy_metadata(self, attn_metadata, metadata: FlashMLAAttentionMetadata) -> None: @@ -130,13 +129,12 @@ def apply_legacy_metadata(self, attn_metadata, metadata: FlashMLAAttentionMetada def make_cudagraph_buffer(self, graph_meta, input_buffers, step_context) -> FlashMLAAttentionMetadata: - model_config = step_context.model_config return _build_flash_mla_metadata( torch.ones(graph_meta.max_batchs, dtype=torch.int32, device=graph_meta.device), num_attention_heads=self.num_attention_heads, decoding_query_len=graph_meta.decode_query_len, - is_fp8_kvcache=model_config.use_mla_fp8_cache, - index_topk=model_config.mla_index_topk, + is_fp8_kvcache=step_context.model_config.use_mla_fp8_cache, + index_topk=self.index_topk, ) def fill_cudagraph_buffer(self, graph_meta, input_buffers, step_context, @@ -149,13 +147,12 @@ def fill_cudagraph_buffer(self, graph_meta, input_buffers, step_context, assert buffer.num_splits is None return buffer - model_config = step_context.model_config metadata = _build_flash_mla_metadata( input_buffers['kv_seqlens'], num_attention_heads=self.num_attention_heads, decoding_query_len=graph_meta.decode_query_len, - is_fp8_kvcache=model_config.use_mla_fp8_cache, - index_topk=model_config.mla_index_topk, + is_fp8_kvcache=step_context.model_config.use_mla_fp8_cache, + index_topk=self.index_topk, ) tile_scheduler_metadata.copy_(metadata.tile_scheduler_metadata) assert buffer.num_splits is not None and metadata.num_splits is not None @@ -183,109 +180,19 @@ def _cdiv(a, b): return (a + b - 1) // b -def _try_dynamic_compile(func, *args, **kwargs): - """Try compile.""" - try: - compiled_func = torch.compile(func, dynamic=True) - compiled_func(*args, **kwargs) - return compiled_func - except Exception: - return func - - -class NSAIndicesUpdater: - """NSA indices updater. - - Flash MLA sparse attention requires different indice format for prefill and decoding. This module is used to update - the indices to meet the requirements. - """ - - def __init__(self): - self._update_decode_funcs = dict() - self._update_prefill_func = None - - def _update_decode_impl(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, - block_size: int, block_stride: int, token_stride: int, - index_stride: int, expand_block_offsets: bool) -> torch.Tensor: - """Map logical decode indices to the selected cache layout.""" - batch_size = block_offsets.size(0) - query_len = nsa_indices.size(0) // batch_size - if expand_block_offsets: - block_offsets = block_offsets[:, None, :].expand(-1, query_len, -1).flatten(0, 1) - block_ids = nsa_indices // block_size - block_ids = block_ids.clamp_min(0) - block_ids = block_offsets.gather(1, block_ids) - block_remain = nsa_indices % block_size - ret = (block_ids * block_stride + block_remain * token_stride) // index_stride - ret[nsa_indices < 0] = -1 - return ret.unflatten(0, (batch_size, query_len)) - - def _update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, - max_q_seqlen: int, block_size: int, block_stride: int, - token_stride: int, index_stride: int) -> torch.Tensor: - """Dispatch a cached specialization for one decode layout.""" - expand_block_offsets = max_q_seqlen != 1 - key = (expand_block_offsets, block_size, block_stride, token_stride, index_stride) - args = (nsa_indices, block_offsets, block_size, block_stride, token_stride, - index_stride, expand_block_offsets) - func = self._update_decode_funcs.get(key) - if func is None: - func = _try_dynamic_compile(self._update_decode_impl, *args) - self._update_decode_funcs[key] = func - return func(*args) - - def update_decode(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int, - block_size: int) -> torch.Tensor: - """Update for decode.""" - return self._update_decode(nsa_indices, block_offsets, max_q_seqlen, block_size, - block_stride=block_size, token_stride=1, index_stride=1) - - def update_decode_strided(self, nsa_indices: torch.Tensor, block_offsets: torch.Tensor, max_q_seqlen: int, - block_size: int, block_stride: int, token_stride: int, - index_stride: int) -> torch.Tensor: - """Update decode indices for strided cache storage.""" - return self._update_decode(nsa_indices, block_offsets, max_q_seqlen, block_size, - block_stride, token_stride, index_stride) - - def _update_prefill_impl(self, nsa_indices: torch.Tensor, q_seqlens: torch.Tensor, cu_seqlens_k: torch.Tensor): - """Update for prefill impl.""" - num_tokens = nsa_indices.size(0) - repeat_cu_seqlens_k = torch.repeat_interleave(cu_seqlens_k[:-1], q_seqlens, output_size=num_tokens) - neg_mask = nsa_indices < 0 - nsa_indices = nsa_indices + repeat_cu_seqlens_k[:, None] - nsa_indices[neg_mask] = -1 - return nsa_indices[:, None] - - def update_prefill(self, nsa_indices: torch.Tensor, q_seqlens: torch.Tensor, cu_seqlens_k: torch.Tensor): - """Update for prefill.""" - if self._update_prefill_func is None: - self._update_prefill_func = _try_dynamic_compile(self._update_prefill_impl, nsa_indices, q_seqlens, - cu_seqlens_k) - - return self._update_prefill_func(nsa_indices, q_seqlens, cu_seqlens_k) - - @staticmethod - @functools.cache - def build(): - return NSAIndicesUpdater() +class FlashMLAImpl(TritonAttentionImpl): + """Dense MLA attention. + Prefill: FA3 when available; otherwise the Triton MLA kernel. + Decode: paged FlashMLA. -class FlashMLAImpl(TritonAttentionImpl): - """Flash MLA (Multi-head Latent Attention) implementation. - - This implementation supports multiple execution paths: - - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache - - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over a zero-copy BF16 cache view - - Sparse FlashMLA prefill: Uses flash_mla_sparse_fwd for NSA attention - - Prefill with FA3: Uses flash_attn_varlen_func with split q_rope/q_nope - - Prefill fallback: Uses custom Triton kernel + DSA-specific index mapping and sparse execution live in + :class:`FlashMLASparseImpl`. """ # MLA-specific constants - _MLA_HEAD_ALIGNMENT = 64 # Query heads must be multiple of 64 for flash_mla _MLA_NOPE_SIZE = 512 # Size of non-positional embeddings _MLA_SCALE_SIZE = 16 # Size of FP8 quantization scales - _BF16_CACHE_INDEX_STRIDE = 64 # Address BF16 cache rows in aligned 128-byte units def __init__( self, @@ -325,14 +232,11 @@ def __init__( from lmdeploy.pytorch.kernels.cuda.fill_kv_cache import fill_kv_cache_blocked_fp8 from lmdeploy.pytorch.kernels.cuda.flatten_kv_cache import flatten_kv_cache_mla_fp8 self.flash_mla_with_kvcache = flash_mla.flash_mla_with_kvcache - self.flash_mla_sparse_fwd = None self.fill_kv_cache_blocked_fp8 = fill_kv_cache_blocked_fp8 self.flatten_kv_cache_mla_fp8 = flatten_kv_cache_mla_fp8 assert num_kv_heads == 1, 'MLA requires num kv heads equal to 1' self.use_fa3 = use_fa3 - self.nsa_updater = NSAIndicesUpdater.build() - def get_step_metadata_provider(self): """Describe metadata required by this selected implementation.""" return FlashMLAAttentionMetaBuilder(num_attention_heads=self.num_heads) @@ -344,26 +248,17 @@ def _get_scheduler_metadata(self, attn_metadata: TritonAttentionMetadata): assert isinstance(kernel_metadata, FlashMLAAttentionMetadata) return kernel_metadata.tile_scheduler_metadata, kernel_metadata.num_splits - def _get_flash_mla_sparse_fwd(self): - if self.flash_mla_sparse_fwd is not None: - return self.flash_mla_sparse_fwd - - try: - import flash_mla - self.flash_mla_sparse_fwd = flash_mla.flash_mla_sparse_fwd - return self.flash_mla_sparse_fwd - except Exception: - logger.exception('Can not import flash_mla_sparse_fwd from flash_mla.') - - def _decode_paged_flash_mla( + def _decoding_paged( self, query: torch.Tensor, k_cache: torch.Tensor, - nsa_indices: torch.Tensor, attn_metadata: TritonAttentionMetadata, + indices: torch.Tensor = None, + causal: bool = None, ): - """Flash mla decoding.""" - causal = self.causal + """Run paged FlashMLA decode with optional provider-ready indices.""" + if causal is None: + causal = self.causal kv_seqlens = attn_metadata.kv_seqlens block_offsets = attn_metadata.block_offsets is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn @@ -379,19 +274,6 @@ def _decode_paged_flash_mla( tile_scheduler_metadata, num_splits = self._get_scheduler_metadata(attn_metadata) - # update nsa indice according to flash-mla requirement - if nsa_indices is not None: - block_size = k_cache.size(1) - nsa_indices = self.nsa_updater.update_decode(nsa_indices, block_offsets, max_q_seqlen, block_size) - causal = False - # The new FlashMLASchedMeta API uses a sparse decoder that only - # accepts 64 or 128 query heads. The old API stores metadata in a - # tensor and supports the unpadded TP head count. - if not isinstance(tile_scheduler_metadata, torch.Tensor): - pad_heads = -num_q_heads % self._MLA_HEAD_ALIGNMENT - if pad_heads: - query = torch.nn.functional.pad(query, (0, 0, 0, pad_heads)) - attn_output, _ = self.flash_mla_with_kvcache(query, k_cache=k_cache, block_table=block_offsets, @@ -402,60 +284,12 @@ def _decode_paged_flash_mla( num_splits=num_splits, causal=causal, is_fp8_kvcache=is_fp8_kvcache, - indices=nsa_indices) + indices=indices) attn_output = attn_output[:, :, :num_q_heads] attn_output = attn_output.flatten(0, 1) return attn_output - def _flash_mla_sparse(self, query: torch.Tensor, indexed_k: torch.Tensor, - nsa_indices: torch.Tensor) -> torch.Tensor: - """Run sparse FlashMLA over index-addressable BF16 KV storage.""" - flash_mla_sparse_fwd = self._get_flash_mla_sparse_fwd() - num_q_heads = query.size(1) - # flash_mla_sparse_fwd requires query heads to be multiple of alignment - if num_q_heads % self._MLA_HEAD_ALIGNMENT != 0: - padding = self._MLA_HEAD_ALIGNMENT - num_q_heads % self._MLA_HEAD_ALIGNMENT - query = torch.nn.functional.pad(query, (0, 0, 0, padding)) - - output = flash_mla_sparse_fwd( - query, - indexed_k, - nsa_indices, - sm_scale=self.scale, - ) - attn_output = output[0] - attn_output = attn_output[:, :num_q_heads] - return attn_output - - def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, nsa_indices: torch.Tensor, - attn_metadata: TritonAttentionMetadata) -> torch.Tensor: - """Run sparse prefill over the flattened BF16 KV cache.""" - nsa_indices = self.nsa_updater.update_prefill(nsa_indices, attn_metadata.q_seqlens, - attn_metadata.cu_seqlens_k) - return self._flash_mla_sparse(query, flatten_k, nsa_indices) - - def _decode_bf16_sparse_flash_mla(self, query: torch.Tensor, k_cache: torch.Tensor, nsa_indices: torch.Tensor, - attn_metadata: TritonAttentionMetadata) -> torch.Tensor: - """Run sparse decode over a zero-copy BF16 paged-cache view.""" - assert query.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 query' - assert k_cache.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 KV cache' - block_size = k_cache.size(1) - max_q_seqlen = self._get_max_q_seqlen(query, attn_metadata) - # The cache shares a block record with indexer state, so flattening it - # copies the full cache capacity. Expose the same storage in aligned - # address units and let sparse indices point directly at each KV row. - index_stride = self._BF16_CACHE_INDEX_STRIDE - block_stride, token_stride = k_cache.stride()[:2] - last_token_offset = ((k_cache.size(0) - 1) * block_stride + (block_size - 1) * token_stride) - storage_rows = last_token_offset // index_stride + 1 - storage_k = k_cache.as_strided((storage_rows, *k_cache.shape[2:]), - (index_stride, *k_cache.stride()[2:])) - nsa_indices = self.nsa_updater.update_decode_strided(nsa_indices, attn_metadata.block_offsets, max_q_seqlen, - block_size, block_stride, token_stride, index_stride) - nsa_indices = nsa_indices.flatten(0, 1)[:, None] - return self._flash_mla_sparse(query, storage_k, nsa_indices) - def _prefill_triton( self, query: torch.Tensor, @@ -539,15 +373,15 @@ def _prefill_fa3( ) return attn_output - def run_flatten_kv_cache(self, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - attn_metadata: TritonAttentionMetadata, - out_dtype: torch.dtype, - is_nsa: bool, - k_scales_zeros: torch.Tensor = None, - v_scales_zeros: torch.Tensor = None): - """Flatten kv cache for prefill.""" + def _flatten_prefill_kv_cache(self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + out_dtype: torch.dtype, + kv_layout: str, + k_scales_zeros: torch.Tensor = None, + v_scales_zeros: torch.Tensor = None): + """Flatten paged KV into the layout required by prefill.""" kv_start_loc = attn_metadata.kv_start_loc kv_seqlens = attn_metadata.kv_seqlens @@ -555,15 +389,13 @@ def run_flatten_kv_cache(self, kv_flatten_size = attn_metadata.kv_flatten_size quant_policy = attn_metadata.quant_policy is_fp8_kvcache = k_cache.dtype == torch.float8_e4m3fn - BLOCK_BS = k_cache.size(1) + block_size = k_cache.size(1) # pad one more block to avoid invalid kv visit - if self.use_fa3 or is_nsa: + if kv_layout == 'shd': out_size = kv_flatten_size - flatten_kv_layout = 'shd' else: - out_size = (_cdiv(kv_flatten_size, BLOCK_BS) * BLOCK_BS + BLOCK_BS) - flatten_kv_layout = 'hsd' + out_size = _cdiv(kv_flatten_size, block_size) * block_size + block_size if is_fp8_kvcache: flatten_k = self.flatten_kv_cache_mla_fp8( @@ -573,7 +405,7 @@ def run_flatten_kv_cache(self, start_loc=kv_start_loc, out_size=out_size, out_dtype=out_dtype, - flatten_kv_layout=flatten_kv_layout, + flatten_kv_layout=kv_layout, ) flatten_v = flatten_k[..., :self._MLA_NOPE_SIZE] else: @@ -588,7 +420,7 @@ def run_flatten_kv_cache(self, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, quant_policy=quant_policy, - flatten_kv_layout=flatten_kv_layout, + flatten_kv_layout=kv_layout, ) return flatten_k, flatten_v @@ -681,23 +513,20 @@ def _forward_decoding( attn_metadata: TritonAttentionMetadata, nsa_indices: torch.Tensor = None, ) -> torch.Tensor: - """Forward pass for decoding stage. - - Uses paged FlashMLA decode by default. Sparse NSA with BF16 cache uses - sparse FlashMLA decode. + """Forward pass for dense MLA decoding. Args: query: Query tensor. k_cache: Key cache tensor. attn_metadata: Attention metadata. - nsa_indices: Optional sparse attention indices. + nsa_indices: Must be ``None`` for dense MLA. Returns: Attention output tensor. """ - if nsa_indices is not None and k_cache.dtype != torch.float8_e4m3fn: - return self._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, attn_metadata) - return self._decode_paged_flash_mla(query, k_cache, nsa_indices, attn_metadata) + if nsa_indices is not None: + raise RuntimeError('Sparse MLA indices require FlashMLASparseImpl.') + return self._decoding_paged(query, k_cache, attn_metadata) def _forward_prefill( self, @@ -709,43 +538,37 @@ def _forward_prefill( k_scales_zeros: torch.Tensor = None, v_scales_zeros: torch.Tensor = None, ) -> torch.Tensor: - """Forward pass for prefill stage. - - Supports three execution paths: - 1. Sparse FlashMLA: flash_mla_sparse_fwd over BF16 flattened KV - 2. FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope - 3. Triton fallback: Custom Triton kernel implementation + """Forward pass for dense MLA prefill. Args: query: Query tensor. k_cache: Key cache tensor. v_cache: Value cache tensor. attn_metadata: Attention metadata. - nsa_indices: Optional sparse attention indices. + nsa_indices: Must be ``None`` for dense MLA. k_scales_zeros: Key quantization scales/zeros. v_scales_zeros: Value quantization scales/zeros. Returns: Attention output tensor. """ - # Flatten KV cache once for all prefill paths - flatten_k, flatten_v = self.run_flatten_kv_cache( + if nsa_indices is not None: + raise RuntimeError('Sparse MLA indices require FlashMLASparseImpl.') + + kv_layout = 'shd' if self.use_fa3 else 'hsd' + flatten_k, flatten_v = self._flatten_prefill_kv_cache( k_cache, v_cache, attn_metadata, out_dtype=query.dtype, - is_nsa=nsa_indices is not None, + kv_layout=kv_layout, k_scales_zeros=k_scales_zeros, v_scales_zeros=v_scales_zeros, ) - # Dispatch to appropriate prefill implementation - if nsa_indices is not None: - return self._prefill_sparse(query, flatten_k, nsa_indices, attn_metadata) - elif self.use_fa3: + if self.use_fa3: return self._prefill_fa3(query, flatten_k, attn_metadata) - else: - return self._prefill_triton(query, flatten_k, flatten_v, attn_metadata) + return self._prefill_triton(query, flatten_k, flatten_v, attn_metadata) def forward( self, @@ -760,20 +583,15 @@ def forward( nsa_indices: torch.Tensor = None, **kwargs, ) -> torch.Tensor: - """Forward pass for MLA attention computation. + """Forward pass for dense MLA attention. This method handles both prefill and decoding stages by: 1. Computing max query sequence length 2. Filling KV cache if new key/value are provided 3. Dispatching to the cache-specific stage implementation - Architecture: - - Paged FlashMLA decode: Uses flash_mla_with_kvcache with FP8 paged KV cache - - Sparse FlashMLA decode: Uses flash_mla_sparse_fwd over a zero-copy BF16 cache view - - Prefill: Three paths based on availability and requirements - * Sparse FlashMLA: flash_mla_sparse_fwd - * FA3 optimized: flash_attn_varlen_func with split q_rope/q_nope - * Triton fallback: Custom triton kernel + Prefill: FA3 when available; otherwise the Triton MLA kernel. + Decode: paged FlashMLA. Args: query: Query tensor. @@ -784,7 +602,7 @@ def forward( attn_metadata: Attention metadata containing stage info and indices. k_scales_zeros: Key quantization scales/zeros. v_scales_zeros: Value quantization scales/zeros. - nsa_indices: Optional sparse attention indices. + nsa_indices: Must be ``None`` for dense MLA. Returns: Attention output tensor. diff --git a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py new file mode 100644 index 0000000000..59fc739bf0 --- /dev/null +++ b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py @@ -0,0 +1,240 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import functools + +import torch + +from lmdeploy.utils import get_logger + +from .default import TritonAttentionMetadata +from .mla import FlashMLAAttentionMetaBuilder, FlashMLAImpl + +logger = get_logger('lmdeploy') + + +def _try_dynamic_compile(func, *args, **kwargs): + """Try compile.""" + try: + compiled_func = torch.compile(func, dynamic=True) + compiled_func(*args, **kwargs) + return compiled_func + except Exception: + return func + + +class FlashMLAIndexMapper: + """Map logical DSA top-k indices to FlashMLA storage indices.""" + + def __init__(self): + self._map_decode_funcs = dict() + self._map_prefill_func = None + + def _map_decode_impl(self, indices: torch.Tensor, block_offsets: torch.Tensor, + block_size: int, block_stride: int, token_stride: int, + index_stride: int, expand_block_offsets: bool) -> torch.Tensor: + """Map logical decode indices to the selected cache layout.""" + batch_size = block_offsets.size(0) + query_len = indices.size(0) // batch_size + if expand_block_offsets: + block_offsets = block_offsets[:, None, :].expand(-1, query_len, -1).flatten(0, 1) + block_ids = indices // block_size + block_ids = block_ids.clamp_min(0) + block_ids = block_offsets.gather(1, block_ids) + block_remain = indices % block_size + mapped_indices = (block_ids * block_stride + block_remain * token_stride) // index_stride + mapped_indices[indices < 0] = -1 + return mapped_indices.unflatten(0, (batch_size, query_len)) + + def _map_decode(self, indices: torch.Tensor, block_offsets: torch.Tensor, + max_q_seqlen: int, block_size: int, block_stride: int, + token_stride: int, index_stride: int) -> torch.Tensor: + """Dispatch a cached specialization for one decode layout.""" + expand_block_offsets = max_q_seqlen != 1 + key = (expand_block_offsets, block_size, block_stride, token_stride, index_stride) + args = (indices, block_offsets, block_size, block_stride, token_stride, + index_stride, expand_block_offsets) + func = self._map_decode_funcs.get(key) + if func is None: + func = _try_dynamic_compile(self._map_decode_impl, *args) + self._map_decode_funcs[key] = func + return func(*args) + + def map_paged_decode(self, indices: torch.Tensor, block_offsets: torch.Tensor, + max_q_seqlen: int, block_size: int) -> torch.Tensor: + """Map logical indices to paged-cache token offsets.""" + return self._map_decode(indices, block_offsets, max_q_seqlen, block_size, + block_stride=block_size, token_stride=1, index_stride=1) + + def map_strided_decode(self, indices: torch.Tensor, block_offsets: torch.Tensor, + max_q_seqlen: int, block_size: int, block_stride: int, + token_stride: int, index_stride: int) -> torch.Tensor: + """Map logical indices to aligned offsets in a strided cache view.""" + return self._map_decode(indices, block_offsets, max_q_seqlen, block_size, + block_stride, token_stride, index_stride) + + def _map_flat_prefill_impl(self, indices: torch.Tensor, q_seqlens: torch.Tensor, + cu_seqlens_k: torch.Tensor): + """Map request-local prefill indices into the flattened KV buffer.""" + num_tokens = indices.size(0) + kv_offsets = torch.repeat_interleave(cu_seqlens_k[:-1], q_seqlens, output_size=num_tokens) + invalid = indices < 0 + indices = indices + kv_offsets[:, None] + indices[invalid] = -1 + return indices[:, None] + + def map_flat_prefill(self, indices: torch.Tensor, q_seqlens: torch.Tensor, + cu_seqlens_k: torch.Tensor): + """Map request-local prefill indices into the flattened KV buffer.""" + if self._map_prefill_func is None: + self._map_prefill_func = _try_dynamic_compile(self._map_flat_prefill_impl, + indices, q_seqlens, cu_seqlens_k) + return self._map_prefill_func(indices, q_seqlens, cu_seqlens_k) + + @staticmethod + @functools.cache + def build(): + """Return the process-local mapper shared by all DSA layers.""" + return FlashMLAIndexMapper() + + +class FlashMLASparseImpl(FlashMLAImpl): + """Sparse DSA attention using FlashMLA kernels. + + Prefill: ``flash_mla_sparse_fwd`` over flattened BF16 KV. + Decode: ``flash_mla_sparse_fwd`` over a zero-copy BF16 cache view, or + ``flash_mla_with_kvcache`` over the packed FP8 cache. + """ + + _MLA_HEAD_ALIGNMENT = 64 + _BF16_CACHE_INDEX_STRIDE = 64 + + def __init__(self, mla_index_topk: int, **kwargs): + super().__init__(**kwargs) + self.mla_index_topk = mla_index_topk + self.flash_mla_sparse_fwd = None + self.index_mapper = FlashMLAIndexMapper.build() + + def get_step_metadata_provider(self): + """Describe metadata required by sparse FlashMLA.""" + return FlashMLAAttentionMetaBuilder(num_attention_heads=self.num_heads, + index_topk=self.mla_index_topk) + + def _get_flash_mla_sparse_fwd(self): + if self.flash_mla_sparse_fwd is not None: + return self.flash_mla_sparse_fwd + + try: + import flash_mla + self.flash_mla_sparse_fwd = flash_mla.flash_mla_sparse_fwd + return self.flash_mla_sparse_fwd + except Exception: + logger.exception('Can not import flash_mla_sparse_fwd from flash_mla.') + + def _flash_mla_sparse_forward(self, query: torch.Tensor, indexed_kv: torch.Tensor, + indices: torch.Tensor) -> torch.Tensor: + """Run sparse FlashMLA over index-addressable BF16 KV storage.""" + flash_mla_sparse_fwd = self._get_flash_mla_sparse_fwd() + num_q_heads = query.size(1) + pad_heads = -num_q_heads % self._MLA_HEAD_ALIGNMENT + if pad_heads: + query = torch.nn.functional.pad(query, (0, 0, 0, pad_heads)) + + attn_output = flash_mla_sparse_fwd(query, indexed_kv, indices, sm_scale=self.scale)[0] + return attn_output[:, :num_q_heads] + + def _prefill_sparse(self, query: torch.Tensor, flatten_k: torch.Tensor, + nsa_indices: torch.Tensor, + attn_metadata: TritonAttentionMetadata) -> torch.Tensor: + """Run sparse prefill over flattened BF16 KV.""" + indices = self.index_mapper.map_flat_prefill(nsa_indices, + attn_metadata.q_seqlens, + attn_metadata.cu_seqlens_k) + return self._flash_mla_sparse_forward(query, flatten_k, indices) + + def _decoding_sparse_bf16(self, query: torch.Tensor, k_cache: torch.Tensor, + nsa_indices: torch.Tensor, + attn_metadata: TritonAttentionMetadata) -> torch.Tensor: + """Run sparse decode over a zero-copy BF16 paged-cache view.""" + assert query.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 query' + assert k_cache.dtype == torch.bfloat16, 'BF16 sparse MLA requires a bfloat16 KV cache' + block_size = k_cache.size(1) + max_q_seqlen = self._get_max_q_seqlen(query, attn_metadata) + + # Expose the paged cache in FlashMLA's aligned addressing units without + # copying its full capacity. + index_stride = self._BF16_CACHE_INDEX_STRIDE + block_stride, token_stride = k_cache.stride()[:2] + last_token_offset = ((k_cache.size(0) - 1) * block_stride + + (block_size - 1) * token_stride) + storage_rows = last_token_offset // index_stride + 1 + storage_k = k_cache.as_strided((storage_rows, *k_cache.shape[2:]), + (index_stride, *k_cache.stride()[2:])) + indices = self.index_mapper.map_strided_decode( + nsa_indices, + attn_metadata.block_offsets, + max_q_seqlen, + block_size, + block_stride, + token_stride, + index_stride, + ) + indices = indices.flatten(0, 1)[:, None] + return self._flash_mla_sparse_forward(query, storage_k, indices) + + def _decoding_sparse_fp8(self, query: torch.Tensor, k_cache: torch.Tensor, + nsa_indices: torch.Tensor, + attn_metadata: TritonAttentionMetadata) -> torch.Tensor: + """Run sparse decode directly over the packed FP8 paged cache.""" + max_q_seqlen = query.size(0) // attn_metadata.q_seqlens.size(0) + indices = self.index_mapper.map_paged_decode( + nsa_indices, + attn_metadata.block_offsets, + max_q_seqlen, + k_cache.size(1), + ) + + num_q_heads = query.size(1) + scheduler_metadata, _ = self._get_scheduler_metadata(attn_metadata) + if not isinstance(scheduler_metadata, torch.Tensor): + pad_heads = -num_q_heads % self._MLA_HEAD_ALIGNMENT + if pad_heads: + query = torch.nn.functional.pad(query, (0, 0, 0, pad_heads)) + + output = self._decoding_paged(query, + k_cache, + attn_metadata, + indices=indices, + causal=False) + return output[:, :num_q_heads] + + def _forward_decoding(self, query: torch.Tensor, k_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + nsa_indices: torch.Tensor = None) -> torch.Tensor: + """Forward pass for sparse MLA decoding.""" + if nsa_indices is None: + raise RuntimeError('Sparse MLA requires DSA top-k indices.') + if k_cache.dtype == torch.float8_e4m3fn: + return self._decoding_sparse_fp8(query, k_cache, nsa_indices, attn_metadata) + return self._decoding_sparse_bf16(query, k_cache, nsa_indices, attn_metadata) + + def _forward_prefill(self, + query: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + attn_metadata: TritonAttentionMetadata, + nsa_indices: torch.Tensor = None, + k_scales_zeros: torch.Tensor = None, + v_scales_zeros: torch.Tensor = None) -> torch.Tensor: + """Forward pass for sparse MLA prefill.""" + if nsa_indices is None: + raise RuntimeError('Sparse MLA requires DSA top-k indices.') + flatten_k, _ = self._flatten_prefill_kv_cache( + k_cache, + v_cache, + attn_metadata, + out_dtype=query.dtype, + kv_layout='shd', + k_scales_zeros=k_scales_zeros, + v_scales_zeros=v_scales_zeros, + ) + return self._prefill_sparse(query, flatten_k, nsa_indices, attn_metadata) diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index 0ab08cdd14..ad3f72f06a 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -342,7 +342,8 @@ def __init__(self, num_kv_heads=num_key_value_heads, v_head_size=config.kv_lora_rank, num_replicate_kv_heads=num_replicate_kv_heads, - use_flash_mla=use_flash_mla) + use_flash_mla=use_flash_mla, + mla_index_topk=config.index_topk) self.vc = DeepseekV2BMM(self.num_heads, config.kv_lora_rank, self.v_head_dim, dtype=dtype, device=device) self.o_proj = build_o_proj( diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 0ada98ae57..0946fafe25 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -5,53 +5,73 @@ import pytest import torch +from lmdeploy.pytorch.backends.cuda import attention as attention_module +from lmdeploy.pytorch.backends.cuda.attention import TritonAttentionBuilder from lmdeploy.pytorch.backends.cuda.attention import mla as mla_module -from lmdeploy.pytorch.backends.cuda.attention.mla import FlashMLAImpl, NSAIndicesUpdater +from lmdeploy.pytorch.backends.cuda.attention import sparse_mla as sparse_mla_module +from lmdeploy.pytorch.backends.cuda.attention.sparse_mla import ( + FlashMLAIndexMapper, + FlashMLASparseImpl, +) from lmdeploy.pytorch.backends.cuda.op_backend import CudaOpsBackend def _disable_dynamic_compile(monkeypatch): - monkeypatch.setattr(mla_module, '_try_dynamic_compile', lambda func, *args, **kwargs: func) + monkeypatch.setattr(sparse_mla_module, '_try_dynamic_compile', lambda func, *args, **kwargs: func) -def test_nsa_decode_indices_update(monkeypatch): +def test_flash_mla_builder_selects_sparse_impl(monkeypatch): + dense_output = object() + sparse_output = object() + dense_impl = Mock(return_value=dense_output) + sparse_impl = Mock(return_value=sparse_output) + monkeypatch.setattr(attention_module, '_enable_fa3', lambda *args: False) + monkeypatch.setattr(mla_module, 'FlashMLAImpl', dense_impl) + monkeypatch.setattr(sparse_mla_module, 'FlashMLASparseImpl', sparse_impl) + kwargs = dict(num_heads=64, head_size=576, num_kv_heads=1, use_flash_mla=True) + + assert TritonAttentionBuilder.build(**kwargs) is dense_output + assert TritonAttentionBuilder.build(**kwargs, mla_index_topk=2048) is sparse_output + + +def test_flash_mla_decode_index_mapping(monkeypatch): _disable_dynamic_compile(monkeypatch) - updater = NSAIndicesUpdater() + mapper = FlashMLAIndexMapper() block_offsets = torch.tensor([[100, 101, 102], [200, 201, 202]]) nsa_indices = torch.tensor([[0, 17, -1], [32, 1, 16], [0, 33, 47], [32, 1, 16]]) - output = updater.update_decode(nsa_indices, block_offsets, max_q_seqlen=2, block_size=16) + output = mapper.map_paged_decode(nsa_indices, block_offsets, max_q_seqlen=2, block_size=16) expected = torch.tensor([[[1600, 1617, -1], [1632, 1601, 1616]], [[3200, 3233, 3247], [3232, 3201, 3216]]]) assert torch.equal(output, expected) nsa_indices = torch.tensor([[0, 17], [32, -1]]) - output = updater.update_decode(nsa_indices, block_offsets, max_q_seqlen=1, block_size=16) + output = mapper.map_paged_decode(nsa_indices, block_offsets, max_q_seqlen=1, block_size=16) assert torch.equal(output, torch.tensor([[[1600, 1617]], [[3232, -1]]])) -def test_nsa_decode_indices_update_caches_query_modes(monkeypatch): +def test_flash_mla_decode_index_mapping_caches_query_modes(monkeypatch): compile_func = Mock(side_effect=lambda func, *args, **kwargs: func) - monkeypatch.setattr(mla_module, '_try_dynamic_compile', compile_func) - updater = NSAIndicesUpdater() + monkeypatch.setattr(sparse_mla_module, '_try_dynamic_compile', compile_func) + mapper = FlashMLAIndexMapper() block_offsets = torch.tensor([[100, 101, 102], [200, 201, 202]]) single_indices = torch.tensor([[0, 17], [32, -1]]) multi_indices = torch.tensor([[0, 17], [32, 1], [0, 33], [32, 1]]) - updater.update_decode(single_indices, block_offsets, max_q_seqlen=1, block_size=16) - updater.update_decode(single_indices, block_offsets, max_q_seqlen=1, block_size=16) - updater.update_decode(multi_indices, block_offsets, max_q_seqlen=2, block_size=16) - updater.update_decode(multi_indices.repeat_interleave(2, dim=0), block_offsets, - max_q_seqlen=4, block_size=16) + mapper.map_paged_decode(single_indices, block_offsets, max_q_seqlen=1, block_size=16) + mapper.map_paged_decode(single_indices, block_offsets, max_q_seqlen=1, block_size=16) + mapper.map_paged_decode(multi_indices, block_offsets, max_q_seqlen=2, block_size=16) + mapper.map_paged_decode(multi_indices.repeat_interleave(2, dim=0), block_offsets, + max_q_seqlen=4, block_size=16) assert compile_func.call_count == 2 def test_bf16_sparse_decode_uses_strided_cache_view(monkeypatch): _disable_dynamic_compile(monkeypatch) - impl = object.__new__(FlashMLAImpl) - impl.nsa_updater = NSAIndicesUpdater() - impl._flash_mla_sparse = Mock(return_value=torch.empty(4, 64, 512, dtype=torch.bfloat16)) + impl = object.__new__(FlashMLASparseImpl) + impl.index_mapper = FlashMLAIndexMapper() + impl._flash_mla_sparse_forward = Mock(return_value=torch.empty(4, 64, 512, dtype=torch.bfloat16)) query = torch.empty(4, 64, 576, dtype=torch.bfloat16) block_size = 16 @@ -63,9 +83,9 @@ def test_bf16_sparse_decode_uses_strided_cache_view(monkeypatch): q_seqlens=torch.tensor([2, 2]), block_offsets=torch.tensor([[1, 2, 0], [2, 0, 1]])) - impl._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, metadata) + impl._decoding_sparse_bf16(query, k_cache, nsa_indices, metadata) - sparse_query, storage_k, global_indices = impl._flash_mla_sparse.call_args.args + sparse_query, storage_k, global_indices = impl._flash_mla_sparse_forward.call_args.args assert sparse_query is query assert storage_k.untyped_storage().data_ptr() == k_cache.untyped_storage().data_ptr() assert storage_k.stride() == (64, 576, 1) @@ -79,10 +99,10 @@ def test_bf16_sparse_decode_strided_cache_matches_contiguous_cache(monkeypatch): pytest.skip('FlashMLA BF16 sparse attention requires an SM90 GPU') _disable_dynamic_compile(monkeypatch) - impl = object.__new__(FlashMLAImpl) + impl = object.__new__(FlashMLASparseImpl) impl.scale = 576**-0.5 impl.flash_mla_sparse_fwd = None - impl.nsa_updater = NSAIndicesUpdater() + impl.index_mapper = FlashMLAIndexMapper() batch_size = 2 query_len = 2 @@ -99,22 +119,22 @@ def test_bf16_sparse_decode_strided_cache_matches_contiguous_cache(monkeypatch): q_seqlens=torch.full((batch_size, ), query_len, dtype=torch.int32, device='cuda'), block_offsets=block_offsets) - output = impl._decode_bf16_sparse_flash_mla(query, k_cache, nsa_indices, metadata) + output = impl._decoding_sparse_bf16(query, k_cache, nsa_indices, metadata) contiguous_k = k_cache.flatten(0, 1) - contiguous_indices = impl.nsa_updater.update_decode(nsa_indices, block_offsets, query_len, block_size) + contiguous_indices = impl.index_mapper.map_paged_decode(nsa_indices, block_offsets, query_len, block_size) contiguous_indices = contiguous_indices.flatten(0, 1)[:, None] - expected = impl._flash_mla_sparse(query, contiguous_k, contiguous_indices) + expected = impl._flash_mla_sparse_forward(query, contiguous_k, contiguous_indices) torch.testing.assert_close(output, expected) def test_fp8_sparse_decode_pads_tp_query_heads_for_aligned_kernel(): - impl = object.__new__(FlashMLAImpl) + impl = object.__new__(FlashMLASparseImpl) impl.causal = True impl.scale = 1.0 impl.v_head_size = 512 - impl.nsa_updater = Mock() - impl.nsa_updater.update_decode.return_value = torch.zeros(2, 3, 4, dtype=torch.int32) + impl.index_mapper = Mock() + impl.index_mapper.map_paged_decode.return_value = torch.zeros(2, 3, 4, dtype=torch.int32) impl.flash_mla_with_kvcache = Mock( return_value=(torch.empty(2, 3, 64, 512, dtype=torch.bfloat16), None)) @@ -128,7 +148,7 @@ def test_fp8_sparse_decode_pads_tp_query_heads_for_aligned_kernel(): num_splits=None, ) - output = impl._decode_paged_flash_mla(query, k_cache, torch.zeros(6, 4, dtype=torch.int32), metadata) + output = impl._decoding_sparse_fp8(query, k_cache, torch.zeros(6, 4, dtype=torch.int32), metadata) padded_query = impl.flash_mla_with_kvcache.call_args.args[0] assert padded_query.shape == (2, 3, 64, 576) From 6f93267b4c34f54b9b8b757658aac6000625022e Mon Sep 17 00:00:00 2001 From: zxy Date: Fri, 21 Aug 2026 12:57:56 +0000 Subject: [PATCH 2/6] fix: bound DSA prefill score memory --- lmdeploy/pytorch/backends/cuda/nsa.py | 115 ++++++++++++++++----- lmdeploy/pytorch/engine/executor/base.py | 13 +++ lmdeploy/pytorch/envs.py | 3 + tests/pytorch/engine/test_executor_base.py | 32 ++++++ tests/pytorch/nn/test_nsa.py | 68 ++++++++++++ 5 files changed, 208 insertions(+), 23 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 96201f750c..5658318847 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -6,6 +6,7 @@ import torch from torch import Tensor +from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.backends.cuda.step_metadata import ( CudaAttentionMetaBuilder, CudaSequenceMetadata, @@ -33,6 +34,15 @@ logger = get_logger('lmdeploy') +_FP32_BYTES = 4 + + +def _get_max_score_rows(max_kv_seqlen: int, max_logits_bytes: int) -> int: + """Return the query rows fitting in a bounded FP32 score tensor.""" + if max_kv_seqlen <= 0: + return 1 + return max(1, max_logits_bytes // (max_kv_seqlen * _FP32_BYTES)) + def _get_dsa_indexer_k_cache_views(indexer_k_cache: Tensor, head_dim: int) -> tuple[Tensor, Tensor]: @@ -273,6 +283,7 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int, fill: int) self.fill = fill # TODO: configable scale fmt self.scale_fmt = 'ue8m0' + self.max_logits_bytes = _envs.dsa_indexer_max_logits_mb * (1 << 20) self._sparse_index_topk = _get_sparse_index_topk(topk) self._step_meta_group: int | None = None register_step_metadata_impl(self) @@ -297,26 +308,10 @@ def _compute_scores(self, q: Tensor, q_s: Tensor, """Compute dense index scores with DeepGEMM or the Triton fallback.""" score_meta = meta.score_meta if isinstance(score_meta, _DeepGemmContiguousScoreMeta): - k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( - indexer_k_cache, q.size(-1)) - flat_k, flat_k_s = flatten_dsa_indexer_k_cache( - k_cache, - k_s_cache[..., 0], - meta.cu_seqlen_k, - meta.k_seqlens, - meta.block_offset, - out_size=meta.kv_flatten_size, - ) - return _get_deep_gemm().fp8_fp4_mqa_logits( - q=(q, None), - kv=(flat_k, flat_k_s), - weights=q_s, - cu_seq_len_k_start=score_meta.k_starts, - cu_seq_len_k_end=score_meta.k_ends, - clean_logits=False, - max_seqlen_k=score_meta.max_kv_seqlen, - logits_dtype=torch.float32, - ) + flat_k, flat_k_s = self._flatten_prefill_k( + indexer_k_cache, q.size(-1), meta) + return self._compute_contiguous_scores( + q, q_s, flat_k, flat_k_s, score_meta) if isinstance(score_meta, _DeepGemmPagedScoreMeta): # Paged MQA reads the packed cache directly and requires its compact # ``entries * (D + 4)`` byte block stride. @@ -346,25 +341,99 @@ def _compute_scores(self, q: Tensor, q_s: Tensor, max_k_seqlen=meta.max_kv_seqlen, causal=True) - def _select_topk(self, scores: Tensor, meta: NSAIndexMeta) -> Tensor: + def _flatten_prefill_k(self, indexer_k_cache: Tensor, head_dim: int, + meta: NSAIndexMeta) -> tuple[Tensor, Tensor]: + """Flatten the paged indexer K cache once for prefill scoring.""" + k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( + indexer_k_cache, head_dim) + return flatten_dsa_indexer_k_cache( + k_cache, + k_s_cache[..., 0], + meta.cu_seqlen_k, + meta.k_seqlens, + meta.block_offset, + out_size=meta.kv_flatten_size, + ) + + def _compute_contiguous_scores( + self, + q: Tensor, + q_s: Tensor, + flat_k: Tensor, + flat_k_s: Tensor, + score_meta: _DeepGemmContiguousScoreMeta, + row_slice: slice = slice(None)) -> Tensor: + """Compute one contiguous DeepGEMM score-row slice.""" + return _get_deep_gemm().fp8_fp4_mqa_logits( + q=(q, None), + kv=(flat_k, flat_k_s), + weights=q_s, + cu_seq_len_k_start=score_meta.k_starts[row_slice], + cu_seq_len_k_end=score_meta.k_ends[row_slice], + clean_logits=False, + max_seqlen_k=score_meta.max_kv_seqlen, + logits_dtype=torch.float32, + ) + + def _select_topk(self, scores: Tensor, meta: NSAIndexMeta, + row_slice: slice = slice(None)) -> Tensor: """Select sparse-attention positions from dense index scores.""" + kv_seqlens = meta.indexer_kv_seqlens[row_slice] + # Both selectors consume q_seqlens only when kv_seqlens still has one + # entry per request. DSA metadata already expands it to one entry per + # score row, including for a query-row chunk. if self._sparse_index_topk is not None: return self._sparse_index_topk(scores, meta.q_seqlens, - meta.indexer_kv_seqlens, + kv_seqlens, self.topk, fill=self.fill, descending=True, sorted=False) return bitonic_topk(scores, meta.q_seqlens, - meta.indexer_kv_seqlens, + kv_seqlens, self.topk, fill=self.fill, descending=True) + def _score_and_select_contiguous( + self, q: Tensor, q_s: Tensor, indexer_k_cache: Tensor, + meta: NSAIndexMeta, + score_meta: _DeepGemmContiguousScoreMeta) -> Tensor: + """Bound prefill score memory by chunking only the query rows.""" + flat_k, flat_k_s = self._flatten_prefill_k( + indexer_k_cache, q.size(-1), meta) + max_rows = _get_max_score_rows(score_meta.max_kv_seqlen, + self.max_logits_bytes) + num_rows = q.size(0) + if num_rows <= max_rows: + scores = self._compute_contiguous_scores( + q, q_s, flat_k, flat_k_s, score_meta) + return self._select_topk(scores, meta) + + logger.debug('Split DSA prefill scores into %d chunks with at most %d query rows.', + (num_rows + max_rows - 1) // max_rows, max_rows) + out = torch.empty((num_rows, self.topk), + dtype=torch.int32, + device=q.device) + for start in range(0, num_rows, max_rows): + end = min(start + max_rows, num_rows) + row_slice = slice(start, end) + scores = self._compute_contiguous_scores( + q[row_slice], q_s[row_slice], flat_k, flat_k_s, + score_meta, row_slice) + selected = self._select_topk(scores, meta, row_slice) + out[row_slice].copy_(selected) + del scores, selected + return out + def _score_and_select(self, q: Tensor, q_s: Tensor, indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor: + score_meta = meta.score_meta + if isinstance(score_meta, _DeepGemmContiguousScoreMeta): + return self._score_and_select_contiguous( + q, q_s, indexer_k_cache, meta, score_meta) scores = self._compute_scores(q, q_s, indexer_k_cache, meta) return self._select_topk(scores, meta) diff --git a/lmdeploy/pytorch/engine/executor/base.py b/lmdeploy/pytorch/engine/executor/base.py index aed35e2ba1..5766a4c091 100644 --- a/lmdeploy/pytorch/engine/executor/base.py +++ b/lmdeploy/pytorch/engine/executor/base.py @@ -4,6 +4,7 @@ import contextlib from typing import Any, NamedTuple +from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.config import BackendConfig, CacheConfig, DistConfig, MiscConfig, ModelConfig, SpecDecodeConfig from lmdeploy.pytorch.disagg.config import EngineRole from lmdeploy.pytorch.disagg.conn.protocol import DistServeInitRequest, DistServeKVTransferEndpointInfo @@ -221,6 +222,12 @@ def _get_rank_cache_block_sizes(self, num_ranks: int, cache_block_size: _CacheBl for rank in range(num_ranks) ] + def _get_dsa_score_workspace_size(self) -> int: + """Return the bounded sparse-indexer score workspace in bytes.""" + if getattr(self.model_config, 'mla_index_topk', None) is None: + return 0 + return _envs.dsa_indexer_max_logits_mb * (1 << 20) + def _get_runtime_size(self, free_mems: list[int], cache_block_size: _CacheBlockSize, vocab_size: int) -> tuple[int, int]: """Find best prefill num.""" @@ -228,12 +235,14 @@ def _get_runtime_size(self, free_mems: list[int], cache_block_size: _CacheBlockS max_prefill_token_num = self.cache_config.max_prefill_token_num max_batches = self.cache_config.max_batches rank_cache_block_sizes = self._get_rank_cache_block_sizes(len(free_mems), cache_block_size) + dsa_score_workspace = self._get_dsa_score_workspace_size() runtime_cache_size = 0 while max_prefill_token_num > 0: # Runtime buffers scale mostly with the prefill token budget and # logits/vocab size. They are not pageable KV cache, so reserve # them before applying the KV cache memory ratio. runtime_cache_size = int((max_prefill_token_num + max_batches * 2) * vocab_size * 2) + runtime_cache_size += dsa_score_workspace available_mems = [int((free_mem - runtime_cache_size) * cache_max_entry_count) for free_mem in free_mems] # Keep at least a small number of KV blocks after runtime reserve. # If not possible, reduce the prefill token budget and try again. @@ -393,6 +402,10 @@ def _get_cache_block_sizes(self, spec_cache_config: CacheConfig | None, def _reserve_runtime_mem(self, free_mems: list[int], cache_block_size: _CacheBlockSize, spec_cache_config: CacheConfig | None) -> list[int]: """Reserve runtime memory and update prefill token limit if needed.""" + dsa_score_workspace = self._get_dsa_score_workspace_size() + if dsa_score_workspace > 0: + logger.info('Reserve %d MiB for DSA prefill score workspace.', + dsa_score_workspace >> 20) runtime_mem, max_prefill_token_num = self._get_runtime_size(free_mems, cache_block_size, self.model_config.vocab_size) if self.cache_config.max_prefill_token_num != max_prefill_token_num: diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py index b5b8b0c175..206117bb02 100644 --- a/lmdeploy/pytorch/envs.py +++ b/lmdeploy/pytorch/envs.py @@ -227,6 +227,9 @@ def _patched_get_env( # deepseekv3.2 indexer fusion disable_dsa_indexer_fusion = env_to_bool('LMDEPLOY_DISABLE_DSA_INDEXER_FUSION', False) + # DSA indexer score memory + dsa_indexer_max_logits_mb = max(1, env_to_int('LMDEPLOY_DSA_INDEXER_MAX_LOGITS_MB', 512)) + # cudagraph # fake capture flag for debug cudagraph padding behavior fake_capture = env_to_bool('LMDEPLOY_FAKE_CUDA_GRAPH_CAPTURE', False) diff --git a/tests/pytorch/engine/test_executor_base.py b/tests/pytorch/engine/test_executor_base.py index d371dff3b4..13a18750f4 100644 --- a/tests/pytorch/engine/test_executor_base.py +++ b/tests/pytorch/engine/test_executor_base.py @@ -10,6 +10,7 @@ from lmdeploy.pytorch.disagg.config import EngineRole from lmdeploy.pytorch.engine.cache_engine import CacheEngine, StateCacheEngine from lmdeploy.pytorch.engine.config_builder import ConfigBuilder +from lmdeploy.pytorch.engine.executor import base as executor_base from lmdeploy.pytorch.engine.executor.base import ExecutorBase, _CacheBlockSize @@ -93,6 +94,37 @@ def test_get_num_gpu_blocks_rejects_empty_cache_block(): ExecutorBase._get_num_gpu_blocks(available_mem=4096, cache_block_size=0) +def test_dsa_score_workspace_uses_configured_logits_budget(monkeypatch): + monkeypatch.setattr(executor_base._envs, 'dsa_indexer_max_logits_mb', 7) + executor = object.__new__(ExecutorBase) + executor.model_config = SimpleNamespace(mla_index_topk=2048) + + assert executor._get_dsa_score_workspace_size() == 7 << 20 + + executor.model_config.mla_index_topk = None + assert executor._get_dsa_score_workspace_size() == 0 + + +def test_runtime_size_reserves_dsa_score_workspace(monkeypatch): + monkeypatch.setattr(executor_base._envs, 'dsa_indexer_max_logits_mb', 1) + executor = object.__new__(ExecutorBase) + executor.model_config = SimpleNamespace(mla_index_topk=2048) + executor.cache_config = SimpleNamespace( + cache_max_entry_count=1.0, + max_prefill_token_num=16, + max_batches=2, + ) + executor.dist_config = SimpleNamespace(attn_tp=1) + executor.specdecode_config = None + + runtime_size, max_prefill_token_num = executor._get_runtime_size( + [4 << 20], _CacheBlockSize(target=1024), vocab_size=100) + + generic_runtime_size = (16 + 2 * 2) * 100 * 2 + assert runtime_size == (1 << 20) + generic_runtime_size + assert max_prefill_token_num == 16 + + def test_sync_spec_cache_block_size_updates_kernel_block_size(): executor = object.__new__(ExecutorBase) executor.cache_config = CacheConfig(max_batches=1, diff --git a/tests/pytorch/nn/test_nsa.py b/tests/pytorch/nn/test_nsa.py index cc1fcc417d..86e2ee6c8f 100644 --- a/tests/pytorch/nn/test_nsa.py +++ b/tests/pytorch/nn/test_nsa.py @@ -55,6 +55,68 @@ def test_indexer_meta_builds_causal_rows(): assert meta.max_kv_seqlen == 8 +@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA backend') +def test_deepgemm_prefill_scores_are_chunked_by_logits_budget(monkeypatch): + from lmdeploy.pytorch.backends.cuda import nsa as cuda_nsa + + class FakeDeepGemm: + + def __init__(self): + self.row_counts = [] + + def fp8_fp4_mqa_logits(self, q, kv, weights, cu_seq_len_k_start, + cu_seq_len_k_end, **kwargs): + query = q[0] + self.row_counts.append(query.size(0)) + assert kv == ('flat_k', 'flat_k_s') + assert weights.size(0) == query.size(0) + assert cu_seq_len_k_start.size(0) == query.size(0) + assert cu_seq_len_k_end.size(0) == query.size(0) + return query[:, :1] * 10 + torch.arange(4) + + deep_gemm = FakeDeepGemm() + monkeypatch.setattr(cuda_nsa, '_get_deep_gemm', lambda: deep_gemm) + + impl = object.__new__(cuda_nsa.TritonNSAIndexFP8) + impl.topk = 2 + impl.fill = -1 + impl.max_logits_bytes = 2 * 4 * 4 + flatten_calls = [] + + def fake_flatten(indexer_k_cache, head_dim, meta): + flatten_calls.append((indexer_k_cache, head_dim, meta)) + return 'flat_k', 'flat_k_s' + + def fake_topk(scores, q_seqlens, kv_seqlens, k, **kwargs): + assert q_seqlens.tolist() == [5] + assert kv_seqlens.numel() == scores.size(0) + return torch.topk(scores, k).indices.to(torch.int32) + + impl._flatten_prefill_k = fake_flatten + impl._sparse_index_topk = fake_topk + q = torch.arange(5, dtype=torch.float32).unsqueeze(-1) + q_s = torch.ones(5) + indexer_k_cache = object() + score_meta = cuda_nsa._DeepGemmContiguousScoreMeta( + k_starts=torch.arange(5, dtype=torch.int32), + k_ends=torch.arange(5, dtype=torch.int32) + 4, + max_kv_seqlen=4, + ) + meta = SimpleNamespace( + score_meta=score_meta, + indexer_kv_seqlens=torch.full((5, ), 4, dtype=torch.int32), + q_seqlens=torch.tensor([5], dtype=torch.int32), + ) + + selected = impl._score_and_select(q, q_s, indexer_k_cache, meta) + + assert cuda_nsa._get_max_score_rows(4, impl.max_logits_bytes) == 2 + assert deep_gemm.row_counts == [2, 2, 1] + assert len(flatten_calls) == 1 + assert flatten_calls[0][:2] == (indexer_k_cache, 1) + assert torch.equal(selected, torch.tensor([[3, 2]] * 5, dtype=torch.int32)) + + @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9, reason='requires CUDA device with cc>=9.0') @@ -86,6 +148,12 @@ def test_deepgemm_prefill_scores_match_triton(): topk=2, softmax_scale=1.0, block_size=128, fill=-1) deepgemm_scores = impl._compute_scores(q, q_s, packed_cache, meta) + expected_topk = impl._select_topk(deepgemm_scores, meta) + impl.max_logits_bytes = 2 * meta.max_kv_seqlen * 4 + chunked_topk = impl._score_and_select(q, q_s, packed_cache, meta) + torch.testing.assert_close(chunked_topk.sort().values, + expected_topk.sort().values) + meta.score_meta = None triton_scores = impl._compute_scores(q, q_s, packed_cache, meta) assert deepgemm_scores.shape == triton_scores.shape == (5, 8) From 69dd860ba758e4aa00f3238ed55b566bb43fad21 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 24 Aug 2026 03:17:37 +0000 Subject: [PATCH 3/6] refactor: align DSA score selection paths --- lmdeploy/pytorch/backends/cuda/nsa.py | 85 ++++++++++++--------------- tests/pytorch/nn/test_nsa.py | 19 +++++- 2 files changed, 53 insertions(+), 51 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index 5658318847..aa1414b7e2 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -34,14 +34,13 @@ logger = get_logger('lmdeploy') -_FP32_BYTES = 4 - def _get_max_score_rows(max_kv_seqlen: int, max_logits_bytes: int) -> int: """Return the query rows fitting in a bounded FP32 score tensor.""" if max_kv_seqlen <= 0: return 1 - return max(1, max_logits_bytes // (max_kv_seqlen * _FP32_BYTES)) + _fp32_bytes = 4 + return max(1, max_logits_bytes // (max_kv_seqlen * _fp32_bytes)) def _get_dsa_indexer_k_cache_views(indexer_k_cache: Tensor, @@ -303,44 +302,6 @@ def get_step_metadata(self, attn_metadata) -> NSAIndexMeta: assert isinstance(meta, NSAIndexMeta) return meta - def _compute_scores(self, q: Tensor, q_s: Tensor, - indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor: - """Compute dense index scores with DeepGEMM or the Triton fallback.""" - score_meta = meta.score_meta - if isinstance(score_meta, _DeepGemmContiguousScoreMeta): - flat_k, flat_k_s = self._flatten_prefill_k( - indexer_k_cache, q.size(-1), meta) - return self._compute_contiguous_scores( - q, q_s, flat_k, flat_k_s, score_meta) - if isinstance(score_meta, _DeepGemmPagedScoreMeta): - # Paged MQA reads the packed cache directly and requires its compact - # ``entries * (D + 4)`` byte block stride. - return _get_deep_gemm().fp8_fp4_paged_mqa_logits( - q=(q[:, None], None), - kv_cache=indexer_k_cache, - weights=q_s, - context_lens=score_meta.context_lens, - block_table=score_meta.block_offsets, - schedule_meta=score_meta.schedule, - max_context_len=score_meta.max_kv_seqlen, - clean_logits=False, - logits_dtype=torch.float32, - ) - - _warn_triton_index_scoring() - k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( - indexer_k_cache, q.size(-1)) - return fp8_index(q, - q_s, - k_cache, - k_s_cache[..., 0], - meta.cu_seqlen_q, - meta.k_seqlens, - meta.block_offset, - max_q_seqlen=meta.max_q_seqlen, - max_k_seqlen=meta.max_kv_seqlen, - causal=True) - def _flatten_prefill_k(self, indexer_k_cache: Tensor, head_dim: int, meta: NSAIndexMeta) -> tuple[Tensor, Tensor]: """Flatten the paged indexer K cache once for prefill scoring.""" @@ -355,7 +316,7 @@ def _flatten_prefill_k(self, indexer_k_cache: Tensor, head_dim: int, out_size=meta.kv_flatten_size, ) - def _compute_contiguous_scores( + def _compute_prefill_scores( self, q: Tensor, q_s: Tensor, @@ -363,7 +324,7 @@ def _compute_contiguous_scores( flat_k_s: Tensor, score_meta: _DeepGemmContiguousScoreMeta, row_slice: slice = slice(None)) -> Tensor: - """Compute one contiguous DeepGEMM score-row slice.""" + """Compute one DeepGEMM prefill score-row slice.""" return _get_deep_gemm().fp8_fp4_mqa_logits( q=(q, None), kv=(flat_k, flat_k_s), @@ -397,7 +358,7 @@ def _select_topk(self, scores: Tensor, meta: NSAIndexMeta, fill=self.fill, descending=True) - def _score_and_select_contiguous( + def _score_and_select_prefill( self, q: Tensor, q_s: Tensor, indexer_k_cache: Tensor, meta: NSAIndexMeta, score_meta: _DeepGemmContiguousScoreMeta) -> Tensor: @@ -408,7 +369,7 @@ def _score_and_select_contiguous( self.max_logits_bytes) num_rows = q.size(0) if num_rows <= max_rows: - scores = self._compute_contiguous_scores( + scores = self._compute_prefill_scores( q, q_s, flat_k, flat_k_s, score_meta) return self._select_topk(scores, meta) @@ -420,7 +381,7 @@ def _score_and_select_contiguous( for start in range(0, num_rows, max_rows): end = min(start + max_rows, num_rows) row_slice = slice(start, end) - scores = self._compute_contiguous_scores( + scores = self._compute_prefill_scores( q[row_slice], q_s[row_slice], flat_k, flat_k_s, score_meta, row_slice) selected = self._select_topk(scores, meta, row_slice) @@ -432,9 +393,37 @@ def _score_and_select(self, q: Tensor, q_s: Tensor, indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor: score_meta = meta.score_meta if isinstance(score_meta, _DeepGemmContiguousScoreMeta): - return self._score_and_select_contiguous( + return self._score_and_select_prefill( q, q_s, indexer_k_cache, meta, score_meta) - scores = self._compute_scores(q, q_s, indexer_k_cache, meta) + + if isinstance(score_meta, _DeepGemmPagedScoreMeta): + # Paged MQA reads the packed cache directly and requires its compact + # ``entries * (D + 4)`` byte block stride. + scores = _get_deep_gemm().fp8_fp4_paged_mqa_logits( + q=(q[:, None], None), + kv_cache=indexer_k_cache, + weights=q_s, + context_lens=score_meta.context_lens, + block_table=score_meta.block_offsets, + schedule_meta=score_meta.schedule, + max_context_len=score_meta.max_kv_seqlen, + clean_logits=False, + logits_dtype=torch.float32, + ) + else: + _warn_triton_index_scoring() + k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( + indexer_k_cache, q.size(-1)) + scores = fp8_index(q, + q_s, + k_cache, + k_s_cache[..., 0], + meta.cu_seqlen_q, + meta.k_seqlens, + meta.block_offset, + max_q_seqlen=meta.max_q_seqlen, + max_k_seqlen=meta.max_kv_seqlen, + causal=True) return self._select_topk(scores, meta) def forward(self, q: Tensor, k: Tensor, weights: Tensor, diff --git a/tests/pytorch/nn/test_nsa.py b/tests/pytorch/nn/test_nsa.py index 86e2ee6c8f..a1de30b911 100644 --- a/tests/pytorch/nn/test_nsa.py +++ b/tests/pytorch/nn/test_nsa.py @@ -145,9 +145,11 @@ def test_deepgemm_prefill_scores_match_triton(): k_cache.copy_(torch.randn_like(k_cache.float()).to(k_cache.dtype)) k_s_cache.copy_(torch.rand_like(k_s_cache) * 0.01) impl = cuda_nsa.TritonNSAIndexFP8( - topk=2, softmax_scale=1.0, block_size=128, fill=-1) + topk=2048, softmax_scale=1.0, block_size=128, fill=-1) - deepgemm_scores = impl._compute_scores(q, q_s, packed_cache, meta) + flat_k, flat_k_s = impl._flatten_prefill_k(packed_cache, q.size(-1), meta) + deepgemm_scores = impl._compute_prefill_scores( + q, q_s, flat_k, flat_k_s, meta.score_meta) expected_topk = impl._select_topk(deepgemm_scores, meta) impl.max_logits_bytes = 2 * meta.max_kv_seqlen * 4 chunked_topk = impl._score_and_select(q, q_s, packed_cache, meta) @@ -155,7 +157,18 @@ def test_deepgemm_prefill_scores_match_triton(): expected_topk.sort().values) meta.score_meta = None - triton_scores = impl._compute_scores(q, q_s, packed_cache, meta) + triton_scores = cuda_nsa.fp8_index( + q, + q_s, + k_cache, + k_s_cache[..., 0], + meta.cu_seqlen_q, + meta.k_seqlens, + meta.block_offset, + max_q_seqlen=meta.max_q_seqlen, + max_k_seqlen=meta.max_kv_seqlen, + causal=True, + ) assert deepgemm_scores.shape == triton_scores.shape == (5, 8) for row, row_len in enumerate(meta.indexer_kv_seqlens.tolist()): torch.testing.assert_close( From da39a4950cdcbbf0e557c66b86134a0c0d0c3261 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 24 Aug 2026 08:09:39 +0000 Subject: [PATCH 4/6] perf: optimize short DSA prefill Skip DSA score selection when top-k covers the full prefill context and route the equivalent dense attention through FA3 when available. Preserve MTP index reuse, default sparse MLA caches to BF16 unless FP8 is requested, and cover dispatch and cache-policy behavior. --- .../backends/cuda/attention/__init__.py | 4 +- .../backends/cuda/attention/sparse_mla.py | 13 +++- lmdeploy/pytorch/backends/cuda/nsa.py | 38 +++++++++-- lmdeploy/pytorch/backends/nsa.py | 7 +- .../pytorch/configurations/deepseek_v32.py | 2 +- .../pytorch/configurations/glm_moe_dsa.py | 1 - lmdeploy/pytorch/envs.py | 2 +- lmdeploy/pytorch/models/deepseek_v32.py | 9 ++- lmdeploy/pytorch/models/glm_moe_dsa.py | 24 +++++-- lmdeploy/pytorch/nn/nsa.py | 12 +++- tests/pytorch/config/test_model_config.py | 31 +++++++++ tests/pytorch/kernel/test_mla_attention.py | 29 ++++++++ tests/pytorch/nn/test_nsa.py | 66 +++++++++++++++++++ .../spec_decode/test_glm_moe_dsa_mtp.py | 10 +++ 14 files changed, 225 insertions(+), 23 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/attention/__init__.py b/lmdeploy/pytorch/backends/cuda/attention/__init__.py index 76780ea3d1..21f8cafb35 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/__init__.py +++ b/lmdeploy/pytorch/backends/cuda/attention/__init__.py @@ -137,7 +137,9 @@ def build( if mla_index_topk is not None: logger.debug('Build FlashMLASparseImpl Attention') from .sparse_mla import FlashMLASparseImpl - return FlashMLASparseImpl(mla_index_topk=mla_index_topk, **common_args) + return FlashMLASparseImpl(mla_index_topk=mla_index_topk, + use_fa3=use_fa3, + **common_args) logger.debug('Build FlashMLAImpl Attention') from .mla import FlashMLAImpl return FlashMLAImpl(use_fa3=use_fa3, **common_args) diff --git a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py index 59fc739bf0..cad1969ce8 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/sparse_mla.py @@ -100,7 +100,8 @@ def build(): class FlashMLASparseImpl(FlashMLAImpl): """Sparse DSA attention using FlashMLA kernels. - Prefill: ``flash_mla_sparse_fwd`` over flattened BF16 KV. + Prefill: dense MLA when top-k covers the sequence; otherwise + ``flash_mla_sparse_fwd`` over flattened BF16 KV. Decode: ``flash_mla_sparse_fwd`` over a zero-copy BF16 cache view, or ``flash_mla_with_kvcache`` over the packed FP8 cache. """ @@ -226,6 +227,16 @@ def _forward_prefill(self, k_scales_zeros: torch.Tensor = None, v_scales_zeros: torch.Tensor = None) -> torch.Tensor: """Forward pass for sparse MLA prefill.""" + if attn_metadata.max_kv_seqlen <= self.mla_index_topk: + # Top-k contains every valid key, so dense attention is equivalent + # and avoids sparse index mapping and kernel overhead. + return super()._forward_prefill(query, + k_cache, + v_cache, + attn_metadata, + nsa_indices=None, + k_scales_zeros=k_scales_zeros, + v_scales_zeros=v_scales_zeros) if nsa_indices is None: raise RuntimeError('Sparse MLA requires DSA top-k indices.') flatten_k, _ = self._flatten_prefill_kv_cache( diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index aa1414b7e2..fc8f16a2c5 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -274,12 +274,15 @@ def fill_cudagraph_buffer(self, graph_meta, input_buffers, step_context, class TritonNSAIndexFP8(BaseNSAIndexFP8): - def __init__(self, topk: int, softmax_scale: float, block_size: int, fill: int) -> None: + def __init__(self, topk: int, softmax_scale: float, block_size: int, + fill: int, + allow_short_prefill_scoring_skip: bool = False) -> None: super().__init__() self.topk = topk self.softmax_scale = softmax_scale self.block_size = block_size self.fill = fill + self._allow_short_prefill_scoring_skip = allow_short_prefill_scoring_skip # TODO: configable scale fmt self.scale_fmt = 'ue8m0' self.max_logits_bytes = _envs.dsa_indexer_max_logits_mb * (1 << 20) @@ -287,6 +290,19 @@ def __init__(self, topk: int, softmax_scale: float, block_size: int, fill: int) self._step_meta_group: int | None = None register_step_metadata_impl(self) + def _should_skip_scoring(self, meta: NSAIndexMeta) -> bool: + """Whether dense prefill makes index scoring unnecessary.""" + return (self._allow_short_prefill_scoring_skip and not meta.is_decoding + and meta.max_kv_seqlen <= self.topk) + + def _maybe_score_and_select(self, q: Tensor, q_s: Tensor, + indexer_k_cache: Tensor, + meta: NSAIndexMeta) -> Tensor | None: + """Score after the caller has preserved K for future decode.""" + if self._should_skip_scoring(meta): + return None + return self._score_and_select(q, q_s, indexer_k_cache, meta) + def get_step_metadata_provider(self): """Describe metadata required by the selected DSA indexer.""" return DSAIndexerMetaBuilder() @@ -427,7 +443,7 @@ def _score_and_select(self, q: Tensor, q_s: Tensor, return self._select_topk(scores, meta) def forward(self, q: Tensor, k: Tensor, weights: Tensor, - indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor: + indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor | None: assert q.dim() == 3 assert k.dim() == 2 k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( @@ -451,11 +467,11 @@ def forward(self, q: Tensor, k: Tensor, weights: Tensor, block_offsets=meta.block_offset, group_size=self.block_size, scale_fmt=self.scale_fmt) - return self._score_and_select(q, q_s, indexer_k_cache, meta) + return self._maybe_score_and_select(q, q_s, indexer_k_cache, meta) def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos: Tensor, sin: Tensor, indexer_k_cache: Tensor, norm_eps: float, head_gate_scale: float, - rope_interleaved: bool, meta: NSAIndexMeta) -> Tensor: + rope_interleaved: bool, meta: NSAIndexMeta) -> Tensor | None: """Prepare FP8 Q and write K cache without allocating rotated BF16 Q/K.""" k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( @@ -480,11 +496,19 @@ def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tens max_q_seqlen=meta.max_q_seqlen, eps=norm_eps, rope_interleaved=rope_interleaved) - return self._score_and_select(q, q_s, indexer_k_cache, meta) + return self._maybe_score_and_select(q, q_s, indexer_k_cache, meta) class TritonNSAIndexFP8Builder(BaseNSAIndexFP8Builder): @staticmethod - def build(topk: int, softmax_scale: float, block_size: int = 128, fill: int = -1) -> BaseNSAIndexFP8: - return TritonNSAIndexFP8(topk, softmax_scale=softmax_scale, block_size=block_size, fill=fill) + def build(topk: int, softmax_scale: float, block_size: int = 128, + fill: int = -1, + allow_short_prefill_scoring_skip: bool = False) -> BaseNSAIndexFP8: + return TritonNSAIndexFP8( + topk, + softmax_scale=softmax_scale, + block_size=block_size, + fill=fill, + allow_short_prefill_scoring_skip=allow_short_prefill_scoring_skip, + ) diff --git a/lmdeploy/pytorch/backends/nsa.py b/lmdeploy/pytorch/backends/nsa.py index 3d68239073..6338ba9115 100644 --- a/lmdeploy/pytorch/backends/nsa.py +++ b/lmdeploy/pytorch/backends/nsa.py @@ -85,14 +85,14 @@ def get_step_metadata(self, attn_metadata) -> NSAIndexMeta: @abstractmethod def forward(self, q: Tensor, k: Tensor, weights: Tensor, - indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor: + indexer_k_cache: Tensor, meta: NSAIndexMeta) -> Tensor | None: """forward.""" raise NotImplementedError('Not implemented.') @abstractmethod def forward_fused(self, q: Tensor, k: Tensor, weights: Tensor, norm_weight: Tensor, norm_bias: Tensor, cos: Tensor, sin: Tensor, indexer_k_cache: Tensor, norm_eps: float, head_gate_scale: float, - rope_interleaved: bool, meta: NSAIndexMeta) -> Tensor: + rope_interleaved: bool, meta: NSAIndexMeta) -> Tensor | None: """Forward with fused DSA indexer preparation.""" raise NotImplementedError('Not implemented.') @@ -100,6 +100,7 @@ class BaseNSAIndexFP8Builder: @staticmethod @abstractmethod - def build(topk: int, softmax_scale: float, block_size: int = 128, fill: int = -1) -> BaseNSAIndexFP8: + def build(topk: int, softmax_scale: float, block_size: int = 128, fill: int = -1, + allow_short_prefill_scoring_skip: bool = False) -> BaseNSAIndexFP8: """Build layer implementation.""" raise NotImplementedError('Not implemented.') diff --git a/lmdeploy/pytorch/configurations/deepseek_v32.py b/lmdeploy/pytorch/configurations/deepseek_v32.py index 1360b22e4f..4999ae7978 100644 --- a/lmdeploy/pytorch/configurations/deepseek_v32.py +++ b/lmdeploy/pytorch/configurations/deepseek_v32.py @@ -68,7 +68,7 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): index_k_shape = ([hf_config.index_head_dim], torch.float8_e4m3fn) index_k_scale_shape = ([1], torch.float32) config.cache_shapes = [index_k_shape, index_k_scale_shape] - config.mla_kv_cache_dtype = 'fp8_ds_mla' + config.mla_kv_cache_dtype = 'bfloat16' config.mla_index_topk = hf_config.index_topk config.check_env_func = _check_env_v32 config.post_build_func = _finalize_v32_cache_specs diff --git a/lmdeploy/pytorch/configurations/glm_moe_dsa.py b/lmdeploy/pytorch/configurations/glm_moe_dsa.py index 968ecdc4fe..c726e7d6b0 100644 --- a/lmdeploy/pytorch/configurations/glm_moe_dsa.py +++ b/lmdeploy/pytorch/configurations/glm_moe_dsa.py @@ -33,7 +33,6 @@ def build(cls, hf_config, model_path: str | None = None, **kwargs): hf_config.head_dim = hf_config.qk_rope_head_dim config = super().build(hf_config, model_path=model_path, **kwargs) - config.mla_kv_cache_dtype = 'bfloat16' if is_draft_model: hf_config.architectures[0] = 'GlmMoeDsaMTPModel' return config diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py index 206117bb02..7212a9df3f 100644 --- a/lmdeploy/pytorch/envs.py +++ b/lmdeploy/pytorch/envs.py @@ -224,7 +224,7 @@ def _patched_get_env( # qwen3.5 recurrent_state dtype fp32_mamba_ssm_dtype = env_to_bool('LMDEPLOY_FP32_MAMBA_SSM_DTYPE', False) - # deepseekv3.2 indexer fusion + # DSA indexer fusion disable_dsa_indexer_fusion = env_to_bool('LMDEPLOY_DISABLE_DSA_INDEXER_FUSION', False) # DSA indexer score memory diff --git a/lmdeploy/pytorch/models/deepseek_v32.py b/lmdeploy/pytorch/models/deepseek_v32.py index ad3f72f06a..32d1db683b 100644 --- a/lmdeploy/pytorch/models/deepseek_v32.py +++ b/lmdeploy/pytorch/models/deepseek_v32.py @@ -189,7 +189,14 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic self.k_norm = LayerNorm(self.head_dim, device=device) self.softmax_scale = self.head_dim**-0.5 self.apply_rotary_pos_emb = ApplyRotaryEmb() - self.indexer_topk = IndexerTopKFP8(self.index_topk, self.softmax_scale, block_size=128, fill=-1) + self.indexer_topk = IndexerTopKFP8( + self.index_topk, + self.softmax_scale, + block_size=128, + fill=-1, + # MTP may reuse its first iteration's indices in later drafts. + allow_short_prefill_scoring_skip=layer_idx < config.num_hidden_layers, + ) def forward(self, x: torch.Tensor, diff --git a/lmdeploy/pytorch/models/glm_moe_dsa.py b/lmdeploy/pytorch/models/glm_moe_dsa.py index a54c226dce..320204bc3a 100644 --- a/lmdeploy/pytorch/models/glm_moe_dsa.py +++ b/lmdeploy/pytorch/models/glm_moe_dsa.py @@ -88,7 +88,14 @@ def __init__(self, config: Any, layer_idx: int, dtype: torch.dtype = None, devic self.k_norm = LayerNorm(self.head_dim, device=device) self.softmax_scale = self.head_dim**-0.5 self.apply_rotary_pos_emb = ApplyRotaryEmb() - self.indexer_topk = IndexerTopKFP8(self.index_topk, self.softmax_scale, block_size=128, fill=-1) + self.indexer_topk = IndexerTopKFP8( + self.index_topk, + self.softmax_scale, + block_size=128, + fill=-1, + # MTP may reuse its first iteration's indices in later drafts. + allow_short_prefill_scoring_skip=layer_idx < config.num_hidden_layers, + ) def _apply_rotary_pos_emb(self, q_pe: torch.Tensor, k_pe: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor]): @@ -150,6 +157,8 @@ class DSATopKIndicesBuffer(nn.Module): def __init__(self, topk: int): super().__init__() self.topk = topk + # None means no full indexer has written the current buffer yet. + self._has_indices: bool | None = None self.register_buffer('indices', None, persistent=False) def _target_capacity(self, num_tokens: int) -> int: @@ -172,14 +181,19 @@ def ensure(self, num_tokens: int, device: torch.device) -> torch.Tensor: self.indices = torch.empty(capacity, self.topk, dtype=torch.int32, device=device) return self.indices[:num_tokens] - def write(self, topk_indices: torch.Tensor) -> torch.Tensor: - """Copy freshly computed top-k indices into the shared buffer.""" + def write(self, topk_indices: torch.Tensor | None) -> torch.Tensor | None: + """Store indices, or mark a dense prefill as index-free.""" + self._has_indices = topk_indices is not None + if topk_indices is None: + return None buffer = self.ensure(topk_indices.size(0), topk_indices.device) buffer.copy_(topk_indices) return buffer - def read(self, num_tokens: int, device: torch.device) -> torch.Tensor: - """Read indices previously written by a full indexer layer.""" + def read(self, num_tokens: int, device: torch.device) -> torch.Tensor | None: + """Read the current indices for a shared indexer layer.""" + if self._has_indices is False: + return None if self.indices is None or self.indices.size(0) < num_tokens or self.indices.device != device: raise RuntimeError('DSA top-k indices are reused before the shared buffer is populated.') return self.indices[:num_tokens] diff --git a/lmdeploy/pytorch/nn/nsa.py b/lmdeploy/pytorch/nn/nsa.py index f92e126ab3..fc38d5ff67 100644 --- a/lmdeploy/pytorch/nn/nsa.py +++ b/lmdeploy/pytorch/nn/nsa.py @@ -15,11 +15,19 @@ def get_dsa_indexer_k_cache(layer_idx: int) -> Tensor: class IndexerTopKFP8(nn.Module): - def __init__(self, topk: int, softmax_scale: float, block_size: int = 128, fill: int = -1): + def __init__(self, topk: int, softmax_scale: float, block_size: int = 128, + fill: int = -1, + allow_short_prefill_scoring_skip: bool = False): super().__init__() backend = get_backend() index_builder = backend.get_layer_impl_builder(OpType.NSAIndexFP8) - self.index_impl = index_builder.build(topk, softmax_scale, block_size, fill) + self.index_impl = index_builder.build( + topk, + softmax_scale, + block_size, + fill, + allow_short_prefill_scoring_skip=allow_short_prefill_scoring_skip, + ) def forward( self, diff --git a/tests/pytorch/config/test_model_config.py b/tests/pytorch/config/test_model_config.py index c4c94edb42..2b77274044 100644 --- a/tests/pytorch/config/test_model_config.py +++ b/tests/pytorch/config/test_model_config.py @@ -87,6 +87,37 @@ def test_get_num_qkv_head_by_tp_requires_divisible_heads(): model_config.get_num_qkv_head_by_tp() +def test_dsa_models_default_to_bf16_mla_cache(monkeypatch): + from lmdeploy.pytorch.configurations.deepseek_v2 import DeepseekV2ModelConfigBuilder + from lmdeploy.pytorch.configurations.deepseek_v32 import DeepseekV32ModelConfigBuilder + from lmdeploy.pytorch.configurations.glm_moe_dsa import GlmMoeDsaModelConfigBuilder + + def fake_build(cls, hf_config, model_path=None, **kwargs): + return _make_model_config(num_attention_heads=64, + num_key_value_heads=1) + + monkeypatch.setattr(DeepseekV2ModelConfigBuilder, 'build', classmethod(fake_build)) + deepseek_config = DeepseekV32ModelConfigBuilder.build( + SimpleNamespace( + use_flash_mla=True, + index_head_dim=128, + index_topk=2048, + )) + glm_config = GlmMoeDsaModelConfigBuilder.build( + SimpleNamespace( + use_flash_mla=True, + index_head_dim=128, + index_topk=2048, + qk_head_dim=192, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + quantization_config=None, + )) + + assert deepseek_config.mla_kv_cache_dtype == 'bfloat16' + assert glm_config.mla_kv_cache_dtype == 'bfloat16' + + @pytest.mark.parametrize(('block_size', 'kernel_block_size', 'expected_block_size'), [ (192, 64, 256), (256, 128, 256), diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 0946fafe25..4083d751bc 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -26,12 +26,14 @@ def test_flash_mla_builder_selects_sparse_impl(monkeypatch): dense_impl = Mock(return_value=dense_output) sparse_impl = Mock(return_value=sparse_output) monkeypatch.setattr(attention_module, '_enable_fa3', lambda *args: False) + monkeypatch.setattr(attention_module, 'use_fa3', True) monkeypatch.setattr(mla_module, 'FlashMLAImpl', dense_impl) monkeypatch.setattr(sparse_mla_module, 'FlashMLASparseImpl', sparse_impl) kwargs = dict(num_heads=64, head_size=576, num_kv_heads=1, use_flash_mla=True) assert TritonAttentionBuilder.build(**kwargs) is dense_output assert TritonAttentionBuilder.build(**kwargs, mla_index_topk=2048) is sparse_output + assert sparse_impl.call_args.kwargs['use_fa3'] is True def test_flash_mla_decode_index_mapping(monkeypatch): @@ -163,3 +165,30 @@ def test_bf16_sparse_decode_skips_fp8_flashmla_metadata(): assert metadata.block_offsets.dtype == torch.int32 assert not hasattr(metadata, 'tile_scheduler_metadata') + + +def test_sparse_mla_prefill_routes_by_kv_length(monkeypatch): + dense_output = object() + sparse_output = object() + dense_prefill = Mock(return_value=dense_output) + monkeypatch.setattr(mla_module.FlashMLAImpl, '_forward_prefill', dense_prefill) + impl = object.__new__(FlashMLASparseImpl) + impl.mla_index_topk = 2048 + impl._flatten_prefill_kv_cache = Mock(return_value=(Mock(), Mock())) + impl._prefill_sparse = Mock(return_value=sparse_output) + query, k_cache, v_cache, nsa_indices = (Mock() for _ in range(4)) + + dense = impl._forward_prefill(query, + k_cache, + v_cache, + SimpleNamespace(max_kv_seqlen=2048), + nsa_indices=None) + sparse = impl._forward_prefill(query, + k_cache, + v_cache, + SimpleNamespace(max_kv_seqlen=2049), + nsa_indices=nsa_indices) + + assert dense is dense_output + assert sparse is sparse_output + assert dense_prefill.call_args.kwargs['nsa_indices'] is None diff --git a/tests/pytorch/nn/test_nsa.py b/tests/pytorch/nn/test_nsa.py index a1de30b911..5dc417a41f 100644 --- a/tests/pytorch/nn/test_nsa.py +++ b/tests/pytorch/nn/test_nsa.py @@ -184,6 +184,72 @@ def test_indexer_meta_skips_reused_mtp_topk(): assert not nsa.should_skip_nsa_indexer([dict(skip_topk=False)]) +@pytest.mark.parametrize( + ('allow_skip', 'is_decoding', 'max_kv_seqlen', 'should_skip'), + [(True, False, 2048, True), (True, False, 2049, False), + (True, True, 2048, False), (False, False, 2048, False)]) +def test_short_prefill_caches_indexer_k_before_optional_scoring( + monkeypatch, allow_skip, is_decoding, max_kv_seqlen, should_skip): + from lmdeploy.pytorch.backends.cuda import nsa as cuda_nsa + + prepare_q = torch.tensor(0) + monkeypatch.setattr(cuda_nsa, '_get_dsa_indexer_k_cache_views', + lambda *_: (torch.empty(0), torch.empty(1, 1))) + monkeypatch.setattr(cuda_nsa, 'prepare_dsa_indexer_q', + lambda *args, **kwargs: (prepare_q, prepare_q)) + + k_cache_calls = [] + + def prepare_k_cache(*args, **kwargs): + k_cache_calls.append((args, kwargs)) + + monkeypatch.setattr(cuda_nsa, 'prepare_dsa_indexer_k_cache', + prepare_k_cache) + + impl = object.__new__(cuda_nsa.TritonNSAIndexFP8) + impl.topk = 2048 + impl.softmax_scale = 1.0 + impl._allow_short_prefill_scoring_skip = allow_skip + selected = torch.tensor(1) + score_calls = [] + + def score_and_select(*args): + score_calls.append(args) + return selected + + impl._score_and_select = score_and_select + meta = SimpleNamespace( + is_decoding=is_decoding, + max_kv_seqlen=max_kv_seqlen, + cu_seqlen_q=torch.tensor([0, 1]), + k_seqlens=torch.tensor([1]), + block_offset=torch.tensor([[0]]), + max_q_seqlen=1, + ) + + output = impl.forward_fused( + torch.empty(1), + torch.empty(1), + torch.empty(1), + torch.empty(1), + torch.empty(1), + torch.empty(1), + torch.empty(1), + torch.empty(1), + norm_eps=1e-6, + head_gate_scale=1.0, + rope_interleaved=False, + meta=meta, + ) + + assert len(k_cache_calls) == 1 + assert len(score_calls) == (0 if should_skip else 1) + if should_skip: + assert output is None + else: + assert output is selected + + @pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA backend') def test_sparse_index_topk_is_resolved_at_init(monkeypatch): from lmdeploy.pytorch.backends.cuda import nsa as cuda_nsa diff --git a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py index 1cc7fe4f91..f2bd74e817 100644 --- a/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py +++ b/tests/pytorch/spec_decode/test_glm_moe_dsa_mtp.py @@ -188,3 +188,13 @@ def test_topk_buffer_compaction_preserves_selected_order(): compacted = buffer.compact(torch.tensor([2, 0])) assert torch.equal(compacted, rows[[2, 0]]) + + +def test_topk_buffer_does_not_reuse_stale_indices_for_dense_prefill(): + buffer = DSATopKIndicesBuffer(topk=2) + rows = torch.tensor([[0, 1]], dtype=torch.int32) + buffer.write(rows) + + buffer.write(None) + + assert buffer.read(num_tokens=1, device=rows.device) is None From 99878358830460b55387c4d9d7ab8e62153ce2b3 Mon Sep 17 00:00:00 2001 From: zxy Date: Mon, 24 Aug 2026 08:22:54 +0000 Subject: [PATCH 5/6] chore: clarify DSA logits memory bound Document DeepGEMM's dense logits allocation and the TMA reason for query-only chunking. Remove the redundant builder-level BF16 cache test. --- lmdeploy/pytorch/backends/cuda/nsa.py | 7 +++++ tests/pytorch/config/test_model_config.py | 31 ----------------------- 2 files changed, 7 insertions(+), 31 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index fc8f16a2c5..d98ff259f1 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -39,6 +39,11 @@ def _get_max_score_rows(max_kv_seqlen: int, max_logits_bytes: int) -> int: """Return the query rows fitting in a bounded FP32 score tensor.""" if max_kv_seqlen <= 0: return 1 + # DeepGEMM materializes an aligned [query_rows, max_kv_seqlen] output + # before top-k selection: + # https://github.com/deepseek-ai/DeepGEMM/blob/88965b078186ee7510ab9fc4f1d5ebc19adfa8d1/csrc/apis/attention.hpp#L155-L171 + # Bounding flattened KV alone therefore does not bound the M * N logits + # allocation; limit M so its FP32 payload stays within the runtime budget. _fp32_bytes = 4 return max(1, max_logits_bytes // (max_kv_seqlen * _fp32_bytes)) @@ -379,6 +384,8 @@ def _score_and_select_prefill( meta: NSAIndexMeta, score_meta: _DeepGemmContiguousScoreMeta) -> Tensor: """Bound prefill score memory by chunking only the query rows.""" + # Keep KV contiguous for DeepGEMM's TMA descriptors; slicing only Q + # avoids the alignment failures caused by per-request KV views. flat_k, flat_k_s = self._flatten_prefill_k( indexer_k_cache, q.size(-1), meta) max_rows = _get_max_score_rows(score_meta.max_kv_seqlen, diff --git a/tests/pytorch/config/test_model_config.py b/tests/pytorch/config/test_model_config.py index 2b77274044..c4c94edb42 100644 --- a/tests/pytorch/config/test_model_config.py +++ b/tests/pytorch/config/test_model_config.py @@ -87,37 +87,6 @@ def test_get_num_qkv_head_by_tp_requires_divisible_heads(): model_config.get_num_qkv_head_by_tp() -def test_dsa_models_default_to_bf16_mla_cache(monkeypatch): - from lmdeploy.pytorch.configurations.deepseek_v2 import DeepseekV2ModelConfigBuilder - from lmdeploy.pytorch.configurations.deepseek_v32 import DeepseekV32ModelConfigBuilder - from lmdeploy.pytorch.configurations.glm_moe_dsa import GlmMoeDsaModelConfigBuilder - - def fake_build(cls, hf_config, model_path=None, **kwargs): - return _make_model_config(num_attention_heads=64, - num_key_value_heads=1) - - monkeypatch.setattr(DeepseekV2ModelConfigBuilder, 'build', classmethod(fake_build)) - deepseek_config = DeepseekV32ModelConfigBuilder.build( - SimpleNamespace( - use_flash_mla=True, - index_head_dim=128, - index_topk=2048, - )) - glm_config = GlmMoeDsaModelConfigBuilder.build( - SimpleNamespace( - use_flash_mla=True, - index_head_dim=128, - index_topk=2048, - qk_head_dim=192, - qk_nope_head_dim=128, - qk_rope_head_dim=64, - quantization_config=None, - )) - - assert deepseek_config.mla_kv_cache_dtype == 'bfloat16' - assert glm_config.mla_kv_cache_dtype == 'bfloat16' - - @pytest.mark.parametrize(('block_size', 'kernel_block_size', 'expected_block_size'), [ (192, 64, 256), (256, 128, 256), From bd068e63f2b5b2006645bbe3d22904e5fd2cd71d Mon Sep 17 00:00:00 2001 From: zxy Date: Tue, 25 Aug 2026 04:42:51 +0000 Subject: [PATCH 6/6] fix: correct DSA prefill fallbacks --- .../pytorch/backends/cuda/attention/mla.py | 4 +++ lmdeploy/pytorch/backends/cuda/nsa.py | 5 ++++ tests/pytorch/kernel/test_mla_attention.py | 26 +++++++++++++++++++ tests/pytorch/nn/test_nsa.py | 4 +++ 4 files changed, 39 insertions(+) diff --git a/lmdeploy/pytorch/backends/cuda/attention/mla.py b/lmdeploy/pytorch/backends/cuda/attention/mla.py index 806b066072..9178209c68 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/mla.py +++ b/lmdeploy/pytorch/backends/cuda/attention/mla.py @@ -422,6 +422,10 @@ def _flatten_prefill_kv_cache(self, quant_policy=quant_policy, flatten_kv_layout=kv_layout, ) + if flatten_v.size(-1) == 0: + # BF16 MLA stores the latent value in the leading K payload; + # its standalone V cache is intentionally empty. + flatten_v = flatten_k[..., :self.v_head_size] return flatten_k, flatten_v diff --git a/lmdeploy/pytorch/backends/cuda/nsa.py b/lmdeploy/pytorch/backends/cuda/nsa.py index d98ff259f1..2ee985904d 100644 --- a/lmdeploy/pytorch/backends/cuda/nsa.py +++ b/lmdeploy/pytorch/backends/cuda/nsa.py @@ -435,6 +435,11 @@ def _score_and_select(self, q: Tensor, q_s: Tensor, ) else: _warn_triton_index_scoring() + score_bytes = q.size(0) * meta.max_kv_seqlen * 4 + if score_bytes > self.max_logits_bytes: + raise RuntimeError( + 'DSA index scoring exceeds the configured logits memory budget; ' + 'a compatible DeepGEMM installation is required.') k_cache, k_s_cache = _get_dsa_indexer_k_cache_views( indexer_k_cache, q.size(-1)) scores = fp8_index(q, diff --git a/tests/pytorch/kernel/test_mla_attention.py b/tests/pytorch/kernel/test_mla_attention.py index 4083d751bc..0dfb4bb182 100644 --- a/tests/pytorch/kernel/test_mla_attention.py +++ b/tests/pytorch/kernel/test_mla_attention.py @@ -167,6 +167,32 @@ def test_bf16_sparse_decode_skips_fp8_flashmla_metadata(): assert not hasattr(metadata, 'tile_scheduler_metadata') +def test_bf16_mla_flatten_uses_shared_k_latent_as_value(): + impl = object.__new__(mla_module.FlashMLAImpl) + impl.v_head_size = 512 + flatten_k = torch.empty(3, 1, 576, dtype=torch.bfloat16) + impl.flatten_kv_cache = Mock( + return_value=(flatten_k, torch.empty(3, 1, 0, dtype=torch.bfloat16))) + metadata = SimpleNamespace( + kv_start_loc=torch.tensor([0]), + kv_seqlens=torch.tensor([3]), + block_offsets=torch.tensor([[0]]), + kv_flatten_size=3, + quant_policy=0, + ) + + _, flatten_v = impl._flatten_prefill_kv_cache( + torch.empty(1, 4, 1, 576, dtype=torch.bfloat16), + torch.empty(1, 4, 1, 0, dtype=torch.bfloat16), + metadata, + out_dtype=torch.bfloat16, + kv_layout='hsd', + ) + + assert flatten_v.shape == (3, 1, 512) + assert flatten_v.untyped_storage().data_ptr() == flatten_k.untyped_storage().data_ptr() + + def test_sparse_mla_prefill_routes_by_kv_length(monkeypatch): dense_output = object() sparse_output = object() diff --git a/tests/pytorch/nn/test_nsa.py b/tests/pytorch/nn/test_nsa.py index 5dc417a41f..7acb4298ff 100644 --- a/tests/pytorch/nn/test_nsa.py +++ b/tests/pytorch/nn/test_nsa.py @@ -178,6 +178,10 @@ def test_deepgemm_prefill_scores_match_triton(): atol=1e-3, ) + impl.max_logits_bytes = triton_scores.numel() * triton_scores.element_size() - 1 + with pytest.raises(RuntimeError, match='DeepGEMM installation is required'): + impl._score_and_select(q, q_s, packed_cache, meta) + def test_indexer_meta_skips_reused_mtp_topk(): assert nsa.should_skip_nsa_indexer([dict(skip_topk=True)])