@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 80eb01e93d

[Bugfix] DSV4 TP16 garbage output (#47493)

Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Jeff (Junze) Ma committed Jul 7, 2026 at 21:04 UTC 80eb01e93dcd21695210a6a9c43a82b25536d74d

5 files changed +61 -6

vllm/models/deepseek_v4/attention.py

+4 -4

@@ -618,7 +618,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

618 dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,

619 compress_ratio=self.compress_ratio,

620 cache_dtype_str=self.kv_cache_dtype,

621 - alignment=576 if uses_fp8_ds_mla_layout else None,

621 + alignment=576 if uses_fp8_ds_mla_layout else 512,

622 model_version="deepseek_v4",

623 kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),

624 )

@@ -648,15 +648,15 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):

648 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

649 # head_dim already carries the fp8 scale padding

650 # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout.

651 + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"

652 return MLAAttentionSpec(

653 block_size=self.cache_config.block_size,

654 num_kv_heads=1,

655 head_size=self.head_dim,

656 dtype=self.dtype,

657 compress_ratio=self.compress_ratio,

657 - # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with

658 - # the indexer's compressor state cache. V3.2 keeps the legacy layout.

659 - alignment=576,

658 + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).

659 + alignment=576 if uses_fp8_ds_mla_layout else 512,

660 )

661

662 def forward(self): ...

vllm/models/deepseek_v4/common/ops/cache_utils.py

+35

@@ -654,6 +654,8 @@ def build_flashinfer_mixed_sparse_indices(

654 topk: int,

655 decode_compressed_indices_are_local: bool = False,

656 decode_is_valid_token: torch.Tensor | None = None,

657 + swa_block_span: int | None = None,

658 + compressed_block_span: int | None = None,

659 ) -> tuple[torch.Tensor, torch.Tensor]:

660 """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.

661

@@ -730,6 +732,13 @@ def build_flashinfer_mixed_sparse_indices(

732 max_block_size = max(window_block_size, topk_block_size)

733 num_warps = 4 if max_block_size >= 256 else 1

734

735 + # block_span = page_stride / token_stride; == block_size (no-op) for unpacked KV.

736 + swa_span = swa_block_size if swa_block_span is None else swa_block_span

737 + compressed_span = (

738 + compressed_block_size

739 + if compressed_block_span is None

740 + else compressed_block_span

741 + )

742 _build_flashinfer_mixed_sparse_indices_kernel[(num_tokens,)](

743 sparse_indices,

744 sparse_indices.stride(0),

@@ -748,9 +757,11 @@ def build_flashinfer_mixed_sparse_indices(

757 swa_block_table,

758 swa_block_table.stride(0),

759 swa_block_size,

760 + swa_span,

761 compressed_block_table,

762 compressed_block_table.stride(0),

763 compressed_block_size,

764 + compressed_span,

765 NUM_DECODE_TOKENS=num_decode_tokens,

766 WINDOW_SIZE=window_size,

767 COMPRESS_RATIO=compress_ratio,

@@ -767,6 +778,18 @@ def build_flashinfer_mixed_sparse_indices(

778 return sparse_indices, sparse_topk_lens

779

780

781 +@triton.jit

782 +def _remap_flashinfer_index(values, block_size, block_span):

783 + # FlashInfer's DSv4 kernel indexes sparse KV by physical token stride, so

784 + # packed pages (#44577) need block*block_size+off -> block*block_span+off.

785 + # TODO: remove once flashinfer-ai/flashinfer#3856 is fixed.

786 + is_valid = values >= 0

787 + safe_values = tl.where(is_valid, values, 0)

788 + values = (safe_values // block_size) * block_span

789 + values += safe_values % block_size

790 + return tl.where(is_valid, values, -1)

791 +

792 +

793 @triton.jit(

794 do_not_specialize=[

795 "sparse_indices_stride",

@@ -775,8 +798,10 @@ def build_flashinfer_mixed_sparse_indices(

798 "prefill_topk_stride",

799 "swa_block_table_stride",

800 "swa_block_size",

801 + "swa_block_span",

802 "compressed_block_table_stride",

803 "compressed_block_size",

804 + "compressed_block_span",

805 "NUM_DECODE_TOKENS",

806 "PREFILL_TOPK_STRIDE",

807 ]

@@ -799,9 +824,11 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

824 swa_block_table_ptr,

825 swa_block_table_stride,

826 swa_block_size,

827 + swa_block_span,

828 compressed_block_table_ptr,

829 compressed_block_table_stride,

830 compressed_block_size,

831 + compressed_block_span,

832 NUM_DECODE_TOKENS,

833 WINDOW_SIZE: tl.constexpr,

834 COMPRESS_RATIO: tl.constexpr,

@@ -825,6 +852,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

852 mask=mask,

853 other=-1,

854 )

855 + values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)

856 tl.store(

857 sparse_indices_ptr + token_idx * sparse_indices_stride + offset,

858 values,

@@ -858,6 +886,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

886 values = block_numbers * compressed_block_size + block_offsets

887 values = tl.where(is_valid, values, -1)

888 compressed_len += tl.sum((is_valid & token_valid).to(tl.int32), axis=0)

889 + values = _remap_flashinfer_index(

890 + values, compressed_block_size, compressed_block_span

891 + )

892 tl.store(

893 sparse_indices_ptr

894 + token_idx * sparse_indices_stride

@@ -904,6 +935,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

935 block_offsets = pos_offset % swa_block_size

936 slot_ids = block_numbers * swa_block_size + block_offsets

937 slot_ids = tl.where(offset < swa_len, slot_ids, -1)

938 + slot_ids = _remap_flashinfer_index(slot_ids, swa_block_size, swa_block_span)

939 tl.store(

940 sparse_indices_ptr + token_idx * sparse_indices_stride + offset,

941 slot_ids,

@@ -930,6 +962,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

962 block_offsets = local_idx % compressed_block_size

963 slot_ids = block_numbers * compressed_block_size + block_offsets

964 slot_ids = tl.where((offset < topk_len) & is_valid, slot_ids, -1)

965 + slot_ids = _remap_flashinfer_index(

966 + slot_ids, compressed_block_size, compressed_block_span

967 + )

968 tl.store(

969 sparse_indices_ptr

970 + token_idx * sparse_indices_stride

vllm/models/deepseek_v4/compressor.py

+1 -1

@@ -162,7 +162,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):

162 head_size=self.state_dim,

163 dtype=self.dtype,

164 sliding_window=self.sliding_window,

165 - alignment=576 if uses_fp8_ds_mla_layout else None,

165 + alignment=576 if uses_fp8_ds_mla_layout else 512,

166 )

167

168 def forward(self): ...

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

+19

@@ -45,6 +45,21 @@ def _get_flashinfer_dsv4_workspace(device: torch.device) -> torch.Tensor:

45 return workspace

48 +def _packed_block_span(pool: torch.Tensor) -> int:

49 + """Per-block stride of ``pool`` in tokens (``stride(0)//stride(-2)``): ==

50 + block_size for unpacked KV, larger when packed (#44577). Raises if not

51 + token-aligned."""

52 + block_stride = pool.stride(0)

53 + token_stride = pool.stride(-2)

54 + if block_stride % token_stride != 0:

55 + raise NotImplementedError(

56 + "FLASHINFER_MLA_SPARSE_DSV4 packed KV requires the per-block stride "

57 + f"({block_stride}) to be a multiple of the per-token stride "

58 + f"({token_stride}); this layout is not supported yet."

59 + )

60 + return block_stride // token_stride

61 +

62 +

63 class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):

64 """FlashInfer backend using the DSv4 sparse metadata/cache layout.

@@ -368,6 +383,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

383 )

384 cached_sparse = swa_metadata.flashinfer_sparse_index_cache.get(cache_key, None)

385 if cached_sparse is None:

386 + swa_block_span = _packed_block_span(swa_k_cache)

387 + compressed_block_span = _packed_block_span(compressed_kv_cache)

388 sparse_indices, sparse_topk_lens = build_flashinfer_mixed_sparse_indices(

389 decode_swa_indices,

390 decode_compressed_indices,

@@ -385,6 +402,8 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

402 top_k,

403 decode_compressed_indices_are_local=decode_compressed_indices_are_local,

404 decode_is_valid_token=decode_is_valid_token,

405 + swa_block_span=swa_block_span,

406 + compressed_block_span=compressed_block_span,

407 )

408 if cache_key != "c4a":

409 swa_metadata.flashinfer_sparse_index_cache[cache_key] = (

vllm/v1/attention/backends/mla/sparse_swa.py

+2 -1

@@ -90,7 +90,8 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):

90 dtype=self.dtype,

91 sliding_window=self.window_size,

92 cache_dtype_str=self.cache_config.cache_dtype,

93 - alignment=576 if uses_fp8_ds_mla_layout else None,

93 + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).

94 + alignment=576 if uses_fp8_ds_mla_layout else 512,

95 model_version="deepseek_v4",

96 kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),

97 )