From 9f9e1dac26ff0379651dd7e8ca409b573ccfce54 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Wed, 16 Sep 2026 11:09:00 -0700 Subject: [PATCH] [Bugfix][Model] Restore causal image SWA for DeepSeek V4.1 (#57152) Signed-off-by: Juntian Liu Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../attention/test_deepseek_v4_swa_visible.py | 87 ++++++++++- vllm/models/deepseek_v41/attention.py | 8 - .../deepseek_v41/common/ops/cache_utils.py | 139 ++++-------------- .../nvidia/flash_mla_mega_attn.py | 19 +-- .../deepseek_v41/nvidia/flashinfer_sparse.py | 4 - vllm/models/deepseek_v41/nvidia/flashmla.py | 21 +-- .../configs/deepseek_v41.py | 8 - vllm/v1/attention/backends/mla/sparse_swa.py | 10 +- 8 files changed, 122 insertions(+), 174 deletions(-) diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py index d18b7c671a..1b5097de8e 100644 --- a/tests/v1/attention/test_deepseek_v4_swa_visible.py +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -462,13 +462,31 @@ def test_flashinfer_mixed_sparse_indices_with_image_spans(): assert actual == exp_rows[1:] -def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder: - overrides: dict = {"sliding_window": WINDOW} +def make_builder( + vision: bool, + window: int = WINDOW, + max_image_tokens: int = MAX_IMG, + max_num_batched_tokens: int = 64, + model_name: str = "meta-llama/Meta-Llama-3-8B", + builder_cls: type[ + DeepseekSparseSWAMetadataBuilder + ] = DeepseekSparseSWAMetadataBuilder, + mm_prefix_clamp_sliding_window: bool = True, +) -> DeepseekSparseSWAMetadataBuilder: + overrides: dict = {"sliding_window": window} if vision: - overrides.update(vision_n_layers=2, vision_max_n_token=MAX_IMG) + # Emulate the V4 vision config, whose mm_prefix_clamp_sliding_window + # gates the in-kernel SWA widening. V4.1 (causal image tokens) leaves + # it off. + overrides.update( + vision_n_layers=2, + vision_max_n_token=max_image_tokens, + mm_prefix_clamp_sliding_window=mm_prefix_clamp_sliding_window, + ) vllm_config = create_vllm_config( + model_name=model_name, max_model_len=4096, - max_num_batched_tokens=64, + max_num_batched_tokens=max_num_batched_tokens, max_num_seqs=8, hf_config_override=overrides, ) @@ -477,11 +495,11 @@ def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder: num_kv_heads=1, head_size=512, dtype=torch.bfloat16, - sliding_window=WINDOW, + sliding_window=window, cache_dtype_str="auto", model_version="deepseek_v4", ) - return DeepseekSparseSWAMetadataBuilder( + return builder_cls( kv_cache_spec=spec, layer_names=["layer0"], vllm_config=vllm_config, @@ -578,3 +596,60 @@ def test_builder_text_model_unchanged(): ) assert md.prefill_swa_lens.cpu().tolist() == lens assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("compress_ratio", [0, 1, 2]) +@pytest.mark.parametrize("query_len", [100, 4000]) +def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_len): + """V4.1 image tokens use the same causal SWA as text, including in chunks.""" + from vllm.models.deepseek_v41.common.ops.cache_utils import ( + combine_topk_swa_indices as combine_v41, + ) + from vllm.models.deepseek_v41.sparse_mla import DeepseekV41SparseSWAMetadataBuilder + + seq_len, window, max_image_tokens = 4000, 128, 2048 + spans = [(1900, 3947)] + builder = make_builder( + True, + window, + max_image_tokens, + query_len, + builder_cls=DeepseekV41SparseSWAMetadataBuilder, + mm_prefix_clamp_sliding_window=False, + ) + assert builder.max_image_tokens == 0 + md = build_metadata(builder, [seq_len], [query_len], {0: spans}) + assert md.prefill_left_visible is None + assert md.prefill_right_visible is None + assert md.prefill_swa_indices.shape[-1] == window + plan = md.get_prefill_chunk_plan( + compress_ratio, prefill_chunk_size=4, has_compressed=compress_ratio > 0 + ) + [(start, end, n, m)] = plan + gather_len = int(md.prefill_gather_lens[0]) + assert (start, end, m - n) == (0, 1, gather_len) + assert gather_len == query_len + min(seq_len - query_len, window - 1) + indices, lens = combine_v41( + torch.empty(query_len, 0, dtype=torch.int32, device="cuda"), + md.query_start_loc, + md.prefill_seq_lens, + md.prefill_gather_lens, + window, + compress_ratio, + 0, + m, + n, + ) + indices, lens = indices.cpu(), lens.cpu().tolist() + paged_indices = md.prefill_swa_indices.cpu() + paged_lens = md.prefill_swa_lens.cpu().tolist() + for token, pos in enumerate(range(seq_len - query_len, seq_len)): + lo, hi = max(0, pos - window + 1), pos + 1 + row = indices[token, : lens[token]] + assert lens[token] == hi - lo + assert torch.all((row >= n) & (row < m)) + absolute_keys = row - n + seq_len - gather_len + assert absolute_keys.tolist() == list(range(lo, hi)) + assert paged_lens[token] == hi - lo + assert paged_indices[token, 0, : hi - lo].tolist() == list(range(lo, hi)) diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 4275d25a76..5b26a98c1c 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -288,13 +288,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.n_groups = config.o_groups self.n_local_groups = self.n_groups // tp_size self.window_size = config.sliding_window - # Vision variant: image spans are visible bidirectionally, widening - # prefill SWA index rows by up to max_image_tokens columns. - self.max_image_tokens = ( - getattr(config, "vision_max_n_token", 0) - if getattr(config, "vision_n_layers", 0) > 0 - else 0 - ) # ---- v4.1 sparse-attention topology ---- # compress_ratios has one entry per layer (MTP layers included): # 0 = pure sliding window, 1 = full-length compressed cache, @@ -576,7 +569,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup( window_size=self.window_size, block_size=self.swa_cache_layer.block_size, - max_image_tokens=self.max_image_tokens, ) if self.compress_ratio > 1: diff --git a/vllm/models/deepseek_v41/common/ops/cache_utils.py b/vllm/models/deepseek_v41/common/ops/cache_utils.py index 9a251e0748..e7c3650768 100644 --- a/vllm/models/deepseek_v41/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v41/common/ops/cache_utils.py @@ -815,16 +815,12 @@ def combine_topk_swa_indices( M: int, N: int, out: tuple[torch.Tensor, torch.Tensor] | None = None, - left_visible: torch.Tensor | None = None, - right_visible: torch.Tensor | None = None, - max_image_tokens: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = topk_indices.shape[0] - # max_image_tokens widens the SWA column region for in-image - # bidirectional visibility; the width is fixed per model so the - # caller-provided workspace stays valid on image-free batches. + # The width is fixed per model so the caller-provided workspace stays + # valid on every batch. combined_topk = ( - (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) @@ -853,9 +849,6 @@ def combine_topk_swa_indices( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, - left_visible=left_visible, - right_visible=right_visible, - max_image_tokens=max_image_tokens, ) return combined_indices, combined_lens @@ -924,7 +917,6 @@ class CombineTopkSwaIndicesKernel( TOP_K: int COMPRESS_RATIO: int WINDOW_SIZE: int - IMAGE_WIDTH: int PADDED_TOP_K: int input_variant: TritonPointerInputVariant @@ -946,14 +938,11 @@ class CombineTopkSwaIndicesKernel( query_start_loc_ptr, seq_lens_ptr, gather_lens_ptr, - left_visible_ptr, - right_visible_ptr, M, N, TOP_K: tl.constexpr, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, - IMAGE_WIDTH: tl.constexpr, PADDED_TOP_K: tl.constexpr, ): batch_idx = tl.program_id(0) @@ -988,19 +977,8 @@ class CombineTopkSwaIndicesKernel( # (integer div by constexpr 0 is UB and yields garbage, which # turns the store index below negative). topk_len = 0 - if IMAGE_WIDTH > 0: - # In-image bidirectional visibility: the window starts up to - # max(left - (window - 1), 0) positions earlier and extends - # `right` positions past the query token (both 0 outside - # image spans, reducing to the plain causal window). - left = tl.load(left_visible_ptr + token_idx) - right = tl.load(right_visible_ptr + token_idx) - else: - left = 0 - right = 0 - left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) - swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) - swa_len = pos + right - swa_start + 1 + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0) + swa_len = pos - swa_start + 1 offset = tl.arange(0, PADDED_TOP_K) mask = offset < topk_len @@ -1014,18 +992,17 @@ class CombineTopkSwaIndicesKernel( mask=mask, ) # Index into gathered buffer: N + (position - gather_start) - # For positions [swa_start, pos + right], the buffer indices are: - # [N + swa_start - gather_start, N + pos + right - gather_start] - for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE): - swa_offset = i + tl.arange(0, WINDOW_SIZE) - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + swa_offset, - M * batch_idx + N + swa_start + swa_offset - gather_start, - mask=swa_offset < swa_len, - ) + # For positions [swa_start, pos], the buffer indices are: + # [N + swa_start - gather_start, N + pos - gather_start] + swa_offset = tl.arange(0, WINDOW_SIZE) + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + swa_offset, + M * batch_idx + N + swa_start + swa_offset - gather_start, + mask=swa_offset < swa_len, + ) combined_len = topk_len + swa_len tl.store(combined_lens_ptr + token_idx, combined_len) @@ -1041,7 +1018,6 @@ class CombineTopkSwaIndicesKernel( topk: int, compress_ratio: int, WINDOW_SIZE: int, - image_width: int = 0, ) -> CompileKey: padded_topk = next_power_of_2(topk_width) input_variant = TritonPointerInputVariant.from_alignment( @@ -1054,7 +1030,6 @@ class CombineTopkSwaIndicesKernel( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=WINDOW_SIZE, - IMAGE_WIDTH=image_width, PADDED_TOP_K=padded_topk, input_variant=input_variant, ) @@ -1064,19 +1039,10 @@ class CombineTopkSwaIndicesKernel( return [] window_size = _hf_config_int(vllm_config, "sliding_window", 128) - image_width = ( - _hf_config_int(vllm_config, "vision_max_n_token", 0) - if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0 - else 0 - ) - # Warm both the plain-window variant (batches without image spans) - # and the in-image bidirectional variant. - image_widths = [0, image_width] if image_width > 0 else [0] return self._trace_dispatch(self.dispatch)( _combine_topk_swa_warmup_inputs(vllm_config), _COMBINE_TOPK_SWA_POINTER_INPUTS, WINDOW_SIZE=window_size, - image_width=image_widths, ) def compile(self, compile_key: CompileKey) -> None: @@ -1093,14 +1059,11 @@ class CombineTopkSwaIndicesKernel( input_variant.pointer("query_start_loc", torch.int32), input_variant.pointer("seq_lens", torch.int32), input_variant.pointer("gather_lens", torch.int32), - int32_ptr, - int32_ptr, 1, # do not specialize M 1, # do not specialize N TOP_K=compile_key.TOP_K, COMPRESS_RATIO=compile_key.COMPRESS_RATIO, WINDOW_SIZE=compile_key.WINDOW_SIZE, - IMAGE_WIDTH=compile_key.IMAGE_WIDTH, PADDED_TOP_K=compile_key.PADDED_TOP_K, grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS), ) @@ -1119,13 +1082,8 @@ class CombineTopkSwaIndicesKernel( TOP_K: int, COMPRESS_RATIO: int, WINDOW_SIZE: int, - left_visible: torch.Tensor | None = None, - right_visible: torch.Tensor | None = None, - max_image_tokens: int = 0, ) -> None: num_reqs = seq_lens.shape[0] - has_image = left_visible is not None - image_width = max_image_tokens if has_image else 0 self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)]( combined_indices, combined_indices.stride(0), @@ -1135,14 +1093,11 @@ class CombineTopkSwaIndicesKernel( query_start_loc, seq_lens, gather_lens, - left_visible if has_image else topk_indices, - right_visible if has_image else topk_indices, M, N, TOP_K=TOP_K, COMPRESS_RATIO=COMPRESS_RATIO, WINDOW_SIZE=WINDOW_SIZE, - IMAGE_WIDTH=image_width, PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), ) @@ -1169,9 +1124,6 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token: torch.Tensor | None = None, swa_block_span: int | None = None, compressed_block_span: int | None = None, - prefill_left_visible: torch.Tensor | None = None, - prefill_right_visible: torch.Tensor | None = None, - max_image_tokens: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. @@ -1181,21 +1133,11 @@ def build_flashinfer_mixed_sparse_indices( per token). Decode tokens read precomputed SWA/compressed indices; prefill tokens derive their SWA window from the position and translate local compressed indices to global slots via the block tables. - - When ``prefill_left_visible``/``prefill_right_visible`` are given (vision - variant), the SWA column region widens by ``max_image_tokens`` and prefill - tokens inside an image span get a bidirectionally widened window; decode - rows are padded with -1 across the extra columns. """ assert decode_swa_indices.dtype == torch.int32 assert decode_swa_indices.dim() == 2 - swa_index_width = decode_swa_indices.shape[-1] - assert swa_index_width >= window_size - has_image = prefill_left_visible is not None - image_width = max_image_tokens if has_image else 0 - swa_total_width = swa_index_width + image_width - if has_image: - assert prefill_right_visible is not None + swa_total_width = decode_swa_indices.shape[-1] + assert swa_total_width >= window_size if decode_compressed_topk_lens is not None: assert decode_compressed_topk_lens.dtype == torch.int32 assert prefill_topk_indices.dtype == torch.int32 @@ -1254,7 +1196,7 @@ def build_flashinfer_mixed_sparse_indices( if num_tokens == 0: return sparse_indices, sparse_topk_lens - window_block_size = triton.next_power_of_2(max(swa_index_width, 1)) + window_block_size = triton.next_power_of_2(max(swa_total_width, 1)) topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 @@ -1278,8 +1220,6 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token, prefill_topk_indices, prefill_topk_indices.stride(0), - prefill_left_visible if has_image else token_to_req_indices, - prefill_right_visible if has_image else token_to_req_indices, query_start_loc, seq_lens, token_to_req_indices, @@ -1293,8 +1233,7 @@ def build_flashinfer_mixed_sparse_indices( compressed_span, NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, - SWA_INDEX_WIDTH=swa_index_width, - IMAGE_WIDTH=image_width, + SWA_INDEX_WIDTH=swa_total_width, COMPRESS_RATIO=compress_ratio, TOP_K=topk, PADDED_TOP_K=padded_topk, @@ -1349,8 +1288,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel( decode_is_valid_token_ptr, prefill_topk_indices_ptr, prefill_topk_stride, - left_visible_ptr, - right_visible_ptr, query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, @@ -1365,7 +1302,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel( NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, SWA_INDEX_WIDTH: tl.constexpr, - IMAGE_WIDTH: tl.constexpr, COMPRESS_RATIO: tl.constexpr, TOP_K: tl.constexpr, PADDED_TOP_K: tl.constexpr, @@ -1377,18 +1313,14 @@ def _build_flashinfer_mixed_sparse_indices_kernel( TOPK_BLOCK_SIZE: tl.constexpr, ): token_idx = tl.program_id(0) - # Total SWA column count; > SWA_INDEX_WIDTH only for in-image - # bidirectional visibility (vision variant), where the extra columns are - # -1-padded for decode rows. - SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH if token_idx < NUM_DECODE_TOKENS: - for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < SWA_TOTAL_WIDTH + mask = offset < SWA_INDEX_WIDTH values = tl.load( decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, - mask=offset < SWA_INDEX_WIDTH, + mask=mask, other=-1, ) values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) @@ -1431,7 +1363,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_TOTAL_WIDTH + + SWA_INDEX_WIDTH + offset, values, mask=mask, @@ -1445,7 +1377,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( else: compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len) return prefill_idx = token_idx - NUM_DECODE_TOKENS @@ -1457,26 +1389,17 @@ def _build_flashinfer_mixed_sparse_indices_kernel( start_pos = seq_len - query_len token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query - if IMAGE_WIDTH > 0: - # In-image bidirectional visibility: window starts up to - # max(left - (window - 1), 0) earlier and extends `right` past pos. - left = tl.load(left_visible_ptr + token_idx) - right = tl.load(right_visible_ptr + token_idx) - else: - left = 0 - right = 0 - left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) - swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) - swa_len = pos + right - swa_start_pos + 1 + swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0) + swa_len = pos - swa_start_pos + 1 if COMPRESS_RATIO > 0: topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) else: # SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB. topk_len = 0 - for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < SWA_TOTAL_WIDTH + mask = offset < SWA_INDEX_WIDTH pos_offset = swa_start_pos + offset block_indices = pos_offset // swa_block_size block_numbers = tl.load( @@ -1520,10 +1443,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_TOTAL_WIDTH + + SWA_INDEX_WIDTH + offset, slot_ids, mask=mask, ) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len) diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 233d2f80f9..9722673dd1 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -311,7 +311,7 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention): else: assert self.topk_indices_buffer is not None top_k = self.topk_indices_buffer.shape[-1] - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + combined_topk = round_up(top_k + self.window_size, 128) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16), ((self.max_num_batched_tokens, combined_topk), torch.int32), @@ -411,7 +411,7 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention): ) assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + combined_topk = round_up(top_k + self.window_size, 128) for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan: chunk_size = chunk_end - chunk_start kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous( @@ -457,21 +457,6 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention): chunk_m, chunk_n, out=(idx_ws[: qe - qs], lens_ws[: qe - qs]), - left_visible=( - swa_metadata.prefill_left_visible[ - num_decode_tokens + qs : num_decode_tokens + qe - ] - if swa_metadata.prefill_left_visible is not None - else None - ), - right_visible=( - swa_metadata.prefill_right_visible[ - num_decode_tokens + qs : num_decode_tokens + qe - ] - if swa_metadata.prefill_right_visible is not None - else None - ), - max_image_tokens=self.max_image_tokens, ) chunk_out = _token_slice(out, token_base + qs, token_base + qe) # Mega attention over the gathered non-paged bf16 KV (RoPE already diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 0b43cfad19..409725434a 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -417,10 +417,6 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): decode_is_valid_token=decode_is_valid_token, swa_block_span=swa_block_span, compressed_block_span=compressed_block_span, - prefill_left_visible=swa_metadata.prefill_left_visible, - prefill_right_visible=swa_metadata.prefill_right_visible, - # getattr for tests that bypass __init__ via object.__new__. - max_image_tokens=getattr(self, "max_image_tokens", 0), ) if swa_only: swa_metadata.flashinfer_sparse_index_cache["swa_only"] = ( diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index 47a7b35c76..870098975b 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -120,9 +120,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): else: assert self.topk_indices_buffer is not None top_k = self.topk_indices_buffer.shape[-1] - combined_topk = round_up( - top_k + self.window_size + self.max_image_tokens, 128 - ) + combined_topk = round_up(top_k + self.window_size, 128) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), ((self.max_num_batched_tokens, combined_topk), torch.int32), @@ -303,7 +301,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): ) assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) + combined_topk = round_up(top_k + self.window_size, 128) for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan: chunk_size = chunk_end - chunk_start workspace = workspace_manager.get_simultaneous( @@ -361,21 +359,6 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): chunk_M, chunk_N, out=(combined_indices_out, combined_lens_out), - left_visible=( - swa_metadata.prefill_left_visible[ - num_decode_tokens + query_start : num_decode_tokens + query_end - ] - if swa_metadata.prefill_left_visible is not None - else None - ), - right_visible=( - swa_metadata.prefill_right_visible[ - num_decode_tokens + query_start : num_decode_tokens + query_end - ] - if swa_metadata.prefill_right_visible is not None - else None - ), - max_image_tokens=self.max_image_tokens, ) flash_mla_sparse_fwd( q=q[query_start:query_end], diff --git a/vllm/transformers_utils/configs/deepseek_v41.py b/vllm/transformers_utils/configs/deepseek_v41.py index 8f33245cd6..d8b4335438 100644 --- a/vllm/transformers_utils/configs/deepseek_v41.py +++ b/vllm/transformers_utils/configs/deepseek_v41.py @@ -60,11 +60,3 @@ class DeepseekV41Config(PretrainedConfig): self.vision_max_n_token = vision_config.get("max_image_tokens", 1024) self.vision_min_pixels = vision_config.get("min_pixels", 295936) self.vision_max_wh_ratio = vision_config.get("max_wh_ratio") - # The vision variant needs the mm-prefix plumbing (atomic image-span - # prefill + sparse-SWA window widening); the base arch-config - # convertor reads this config-level attribute. - self.is_mm_prefix_lm = vision_n_layers > 0 - # The sparse-SWA index kernels widen the window within image spans - # in-kernel, so mm-prefix ranges longer than sliding_window must be - # kept (they are only consumed by those kernels). - self.mm_prefix_clamp_sliding_window = vision_n_layers > 0 diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 700f328a37..e88631c32b 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -452,13 +452,15 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): assert hasattr(hf_config, "sliding_window") self.window_size = hf_config.sliding_window - # Vision variant: image spans (up to vision_max_n_token tokens) are + # V4 vision variant: image spans (up to vision_max_n_token tokens) are # visible bidirectionally, so prefill index rows widen from - # window_size to window_size + max_image_tokens. Text-only models keep - # max_image_tokens == 0 and take the original code paths everywhere. + # window_size to window_size + max_image_tokens. The V4 config sets + # mm_prefix_clamp_sliding_window exactly for these in-kernel-widened + # ranges; V4.1 image tokens use the plain causal window, and text-only + # models keep max_image_tokens == 0 everywhere. self.max_image_tokens = ( getattr(hf_config, "vision_max_n_token", 0) - if getattr(hf_config, "vision_n_layers", 0) > 0 + if getattr(hf_config, "mm_prefix_clamp_sliding_window", False) else 0 ) self.prefill_index_width = self.window_size + self.max_image_tokens