diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 906d594ac5..5628a6d0d7 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -618,7 +618,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, compress_ratio=self.compress_ratio, cache_dtype_str=self.kv_cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else None, + alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), ) @@ -648,15 +648,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # head_dim already carries the fp8 scale padding # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, compress_ratio=self.compress_ratio, - # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with - # the indexer's compressor state cache. V3.2 keeps the legacy layout. - alignment=576, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, ) def forward(self): ... diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index ffaec528aa..55106d06af 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -654,6 +654,8 @@ def build_flashinfer_mixed_sparse_indices( topk: int, decode_compressed_indices_are_local: bool = False, decode_is_valid_token: torch.Tensor | None = None, + swa_block_span: int | None = None, + compressed_block_span: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. @@ -730,6 +732,13 @@ def build_flashinfer_mixed_sparse_indices( max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 + # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV. + swa_span = swa_block_size if swa_block_span is None else swa_block_span + compressed_span = ( + compressed_block_size + if compressed_block_span is None + else compressed_block_span + ) _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)]( sparse_indices, sparse_indices.stride(0), @@ -748,9 +757,11 @@ def build_flashinfer_mixed_sparse_indices( swa_block_table, swa_block_table.stride(0), swa_block_size, + swa_span, compressed_block_table, compressed_block_table.stride(0), compressed_block_size, + compressed_span, NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, COMPRESS_RATIO=compress_ratio, @@ -767,6 +778,18 @@ def build_flashinfer_mixed_sparse_indices( return sparse_indices, sparse_topk_lens +@triton.jit +def _remap_flashinfer_index(values, block_size, block_span): + # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so + # packed pages (#44577) need block*block_size+off -> block*block_span+off. + # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed. + is_valid = values >= 0 + safe_values = tl.where(is_valid, values, 0) + values = (safe_values // block_size) * block_span + values += safe_values % block_size + return tl.where(is_valid, values, -1) + + @triton.jit( do_not_specialize=[ "sparse_indices_stride", @@ -775,8 +798,10 @@ def build_flashinfer_mixed_sparse_indices( "prefill_topk_stride", "swa_block_table_stride", "swa_block_size", + "swa_block_span", "compressed_block_table_stride", "compressed_block_size", + "compressed_block_span", "NUM_DECODE_TOKENS", "PREFILL_TOPK_STRIDE", ] @@ -799,9 +824,11 @@ def _build_flashinfer_mixed_sparse_indices_kernel( swa_block_table_ptr, swa_block_table_stride, swa_block_size, + swa_block_span, compressed_block_table_ptr, compressed_block_table_stride, compressed_block_size, + compressed_block_span, NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, COMPRESS_RATIO: tl.constexpr, @@ -825,6 +852,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( mask=mask, other=-1, ) + values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride + offset, values, @@ -858,6 +886,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( values = block_numbers * compressed_block_size + block_offsets values = tl.where(is_valid, values, -1) compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0) + values = _remap_flashinfer_index( + values, compressed_block_size, compressed_block_span + ) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride @@ -904,6 +935,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( block_offsets = pos_offset % swa_block_size slot_ids = block_numbers * swa_block_size + block_offsets slot_ids = tl.where(offset < swa_len, slot_ids, -1) + slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride + offset, slot_ids, @@ -930,6 +962,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( block_offsets = local_idx % compressed_block_size slot_ids = block_numbers * compressed_block_size + block_offsets slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1) + slot_ids = _remap_flashinfer_index( + slot_ids, compressed_block_size, compressed_block_span + ) tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 48cd234038..24838c237c 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -162,7 +162,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase): head_size=self.state_dim, dtype=self.dtype, sliding_window=self.sliding_window, - alignment=576 if uses_fp8_ds_mla_layout else None, + alignment=576 if uses_fp8_ds_mla_layout else 512, ) def forward(self): ... diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index f35fa03252..1848c1930d 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -45,6 +45,21 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor: 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 + + class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend): """FlashInfer backend using the DSv4 sparse metadata/cache layout. @@ -368,6 +383,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): ) 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, @@ -385,6 +402,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): 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, ) if cache_key != "c4a": swa_metadata.flashinfer_sparse_index_cache[cache_key] = ( diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index ca5be3a302..9e54b62a00 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -90,7 +90,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, - alignment=576 if uses_fp8_ds_mla_layout else None, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), )