@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9f9e1dac26

[Bugfix][Model] Restore causal image SWA for DeepSeek V4.1 (#57152)

Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>

Juntian Liu committed Sep 16, 2026 at 11:09 UTC 9f9e1dac26ff0379651dd7e8ca409b573ccfce54

8 files changed +122 -174

tests/v1/attention/test_deepseek_v4_swa_visible.py

+81 -6

@@ -462,13 +462,31 @@ def test_flashinfer_mixed_sparse_indices_with_image_spans():

462 assert actual == exp_rows[1:]

463

464

465 -def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder:

466 - overrides: dict = {"sliding_window": WINDOW}

465 +def make_builder(

466 + vision: bool,

467 + window: int = WINDOW,

468 + max_image_tokens: int = MAX_IMG,

469 + max_num_batched_tokens: int = 64,

470 + model_name: str = "meta-llama/Meta-Llama-3-8B",

471 + builder_cls: type[

472 + DeepseekSparseSWAMetadataBuilder

473 + ] = DeepseekSparseSWAMetadataBuilder,

474 + mm_prefix_clamp_sliding_window: bool = True,

475 +) -> DeepseekSparseSWAMetadataBuilder:

476 + overrides: dict = {"sliding_window": window}

477 if vision:

468 - overrides.update(vision_n_layers=2, vision_max_n_token=MAX_IMG)

478 + # Emulate the V4 vision config, whose mm_prefix_clamp_sliding_window

479 + # gates the in-kernel SWA widening. V4.1 (causal image tokens) leaves

480 + # it off.

481 + overrides.update(

482 + vision_n_layers=2,

483 + vision_max_n_token=max_image_tokens,

484 + mm_prefix_clamp_sliding_window=mm_prefix_clamp_sliding_window,

485 + )

486 vllm_config = create_vllm_config(

487 + model_name=model_name,

488 max_model_len=4096,

471 - max_num_batched_tokens=64,

489 + max_num_batched_tokens=max_num_batched_tokens,

490 max_num_seqs=8,

491 hf_config_override=overrides,

492 )

@@ -477,11 +495,11 @@ def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder:

495 num_kv_heads=1,

496 head_size=512,

497 dtype=torch.bfloat16,

480 - sliding_window=WINDOW,

498 + sliding_window=window,

499 cache_dtype_str="auto",

500 model_version="deepseek_v4",

501 )

484 - return DeepseekSparseSWAMetadataBuilder(

502 + return builder_cls(

503 kv_cache_spec=spec,

504 layer_names=["layer0"],

505 vllm_config=vllm_config,

@@ -578,3 +596,60 @@ def test_builder_text_model_unchanged():

596 )

597 assert md.prefill_swa_lens.cpu().tolist() == lens

598 assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows

599 +

600 +

601 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

602 +@pytest.mark.parametrize("compress_ratio", [0, 1, 2])

603 +@pytest.mark.parametrize("query_len", [100, 4000])

604 +def test_v41_image_prefill_uses_causal_swa(compress_ratio, query_len):

605 + """V4.1 image tokens use the same causal SWA as text, including in chunks."""

606 + from vllm.models.deepseek_v41.common.ops.cache_utils import (

607 + combine_topk_swa_indices as combine_v41,

608 + )

609 + from vllm.models.deepseek_v41.sparse_mla import DeepseekV41SparseSWAMetadataBuilder

610 +

611 + seq_len, window, max_image_tokens = 4000, 128, 2048

612 + spans = [(1900, 3947)]

613 + builder = make_builder(

614 + True,

615 + window,

616 + max_image_tokens,

617 + query_len,

618 + builder_cls=DeepseekV41SparseSWAMetadataBuilder,

619 + mm_prefix_clamp_sliding_window=False,

620 + )

621 + assert builder.max_image_tokens == 0

622 + md = build_metadata(builder, [seq_len], [query_len], {0: spans})

623 + assert md.prefill_left_visible is None

624 + assert md.prefill_right_visible is None

625 + assert md.prefill_swa_indices.shape[-1] == window

626 + plan = md.get_prefill_chunk_plan(

627 + compress_ratio, prefill_chunk_size=4, has_compressed=compress_ratio > 0

628 + )

629 + [(start, end, n, m)] = plan

630 + gather_len = int(md.prefill_gather_lens[0])

631 + assert (start, end, m - n) == (0, 1, gather_len)

632 + assert gather_len == query_len + min(seq_len - query_len, window - 1)

633 + indices, lens = combine_v41(

634 + torch.empty(query_len, 0, dtype=torch.int32, device="cuda"),

635 + md.query_start_loc,

636 + md.prefill_seq_lens,

637 + md.prefill_gather_lens,

638 + window,

639 + compress_ratio,

640 + 0,

641 + m,

642 + n,

643 + )

644 + indices, lens = indices.cpu(), lens.cpu().tolist()

645 + paged_indices = md.prefill_swa_indices.cpu()

646 + paged_lens = md.prefill_swa_lens.cpu().tolist()

647 + for token, pos in enumerate(range(seq_len - query_len, seq_len)):

648 + lo, hi = max(0, pos - window + 1), pos + 1

649 + row = indices[token, : lens[token]]

650 + assert lens[token] == hi - lo

651 + assert torch.all((row >= n) & (row < m))

652 + absolute_keys = row - n + seq_len - gather_len

653 + assert absolute_keys.tolist() == list(range(lo, hi))

654 + assert paged_lens[token] == hi - lo

655 + assert paged_indices[token, 0, : hi - lo].tolist() == list(range(lo, hi))

vllm/models/deepseek_v41/attention.py

-8

@@ -288,13 +288,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

288 self.n_groups = config.o_groups

289 self.n_local_groups = self.n_groups // tp_size

290 self.window_size = config.sliding_window

291 - # Vision variant: image spans are visible bidirectionally, widening

292 - # prefill SWA index rows by up to max_image_tokens columns.

293 - self.max_image_tokens = (

294 - getattr(config, "vision_max_n_token", 0)

295 - if getattr(config, "vision_n_layers", 0) > 0

296 - else 0

297 - )

291 # ---- v4.1 sparse-attention topology ----

292 # compress_ratios has one entry per layer (MTP layers included):

293 # 0 = pure sliding window, 1 = full-length compressed cache,

@@ -576,7 +569,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):

569 _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(

570 window_size=self.window_size,

571 block_size=self.swa_cache_layer.block_size,

579 - max_image_tokens=self.max_image_tokens,

572 )

573

574 if self.compress_ratio > 1:

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

+31 -108

@@ -815,16 +815,12 @@ def combine_topk_swa_indices(

815 M: int,

816 N: int,

817 out: tuple[torch.Tensor, torch.Tensor] | None = None,

818 - left_visible: torch.Tensor | None = None,

819 - right_visible: torch.Tensor | None = None,

820 - max_image_tokens: int = 0,

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

819 num_tokens = topk_indices.shape[0]

823 - # max_image_tokens widens the SWA column region for in-image

824 - # bidirectional visibility; the width is fixed per model so the

825 - # caller-provided workspace stays valid on image-free batches.

820 + # The width is fixed per model so the caller-provided workspace stays

821 + # valid on every batch.

822 combined_topk = (

827 - (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)

823 + (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)

824 // _SPARSE_PREFILL_TOPK_ALIGNMENT

825 * _SPARSE_PREFILL_TOPK_ALIGNMENT

826 )

@@ -853,9 +849,6 @@ def combine_topk_swa_indices(

849 TOP_K=topk,

850 COMPRESS_RATIO=compress_ratio,

851 WINDOW_SIZE=window_size,

856 - left_visible=left_visible,

857 - right_visible=right_visible,

858 - max_image_tokens=max_image_tokens,

852 )

853 return combined_indices, combined_lens

854

@@ -924,7 +917,6 @@ class CombineTopkSwaIndicesKernel(

917 TOP_K: int

918 COMPRESS_RATIO: int

919 WINDOW_SIZE: int

927 - IMAGE_WIDTH: int

920 PADDED_TOP_K: int

921 input_variant: TritonPointerInputVariant

922

@@ -946,14 +938,11 @@ class CombineTopkSwaIndicesKernel(

938 query_start_loc_ptr,

939 seq_lens_ptr,

940 gather_lens_ptr,

949 - left_visible_ptr,

950 - right_visible_ptr,

941 M,

942 N,

943 TOP_K: tl.constexpr,

944 COMPRESS_RATIO: tl.constexpr,

945 WINDOW_SIZE: tl.constexpr,

956 - IMAGE_WIDTH: tl.constexpr,

946 PADDED_TOP_K: tl.constexpr,

947 ):

948 batch_idx = tl.program_id(0)

@@ -988,19 +977,8 @@ class CombineTopkSwaIndicesKernel(

977 # (integer div by constexpr 0 is UB and yields garbage, which

978 # turns the store index below negative).

979 topk_len = 0

991 - if IMAGE_WIDTH > 0:

992 - # In-image bidirectional visibility: the window starts up to

993 - # max(left - (window - 1), 0) positions earlier and extends

994 - # `right` positions past the query token (both 0 outside

995 - # image spans, reducing to the plain causal window).

996 - left = tl.load(left_visible_ptr + token_idx)

997 - right = tl.load(right_visible_ptr + token_idx)

998 - else:

999 - left = 0

1000 - right = 0

1001 - left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)

1002 - swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)

1003 - swa_len = pos + right - swa_start + 1

980 + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1), 0)

981 + swa_len = pos - swa_start + 1

982

983 offset = tl.arange(0, PADDED_TOP_K)

984 mask = offset < topk_len

@@ -1014,18 +992,17 @@ class CombineTopkSwaIndicesKernel(

992 mask=mask,

993 )

994 # Index into gathered buffer: N + (position - gather_start)

1017 - # For positions [swa_start, pos + right], the buffer indices are:

1018 - # [N + swa_start - gather_start, N + pos + right - gather_start]

1019 - for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE):

1020 - swa_offset = i + tl.arange(0, WINDOW_SIZE)

1021 - tl.store(

1022 - combined_indices_ptr

1023 - + token_idx * combined_indices_stride

1024 - + topk_len

1025 - + swa_offset,

1026 - M * batch_idx + N + swa_start + swa_offset - gather_start,

1027 - mask=swa_offset < swa_len,

1028 - )

995 + # For positions [swa_start, pos], the buffer indices are:

996 + # [N + swa_start - gather_start, N + pos - gather_start]

997 + swa_offset = tl.arange(0, WINDOW_SIZE)

998 + tl.store(

999 + combined_indices_ptr

1000 + + token_idx * combined_indices_stride

1001 + + topk_len

1002 + + swa_offset,

1003 + M * batch_idx + N + swa_start + swa_offset - gather_start,

1004 + mask=swa_offset < swa_len,

1005 + )

1006

1007 combined_len = topk_len + swa_len

1008 tl.store(combined_lens_ptr + token_idx, combined_len)

@@ -1041,7 +1018,6 @@ class CombineTopkSwaIndicesKernel(

1018 topk: int,

1019 compress_ratio: int,

1020 WINDOW_SIZE: int,

1044 - image_width: int = 0,

1021 ) -> CompileKey:

1022 padded_topk = next_power_of_2(topk_width)

1023 input_variant = TritonPointerInputVariant.from_alignment(

@@ -1054,7 +1030,6 @@ class CombineTopkSwaIndicesKernel(

1030 TOP_K=topk,

1031 COMPRESS_RATIO=compress_ratio,

1032 WINDOW_SIZE=WINDOW_SIZE,

1057 - IMAGE_WIDTH=image_width,

1033 PADDED_TOP_K=padded_topk,

1034 input_variant=input_variant,

1035 )

@@ -1064,19 +1039,10 @@ class CombineTopkSwaIndicesKernel(

1039 return []

1040

1041 window_size = _hf_config_int(vllm_config, "sliding_window", 128)

1067 - image_width = (

1068 - _hf_config_int(vllm_config, "vision_max_n_token", 0)

1069 - if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0

1070 - else 0

1071 - )

1072 - # Warm both the plain-window variant (batches without image spans)

1073 - # and the in-image bidirectional variant.

1074 - image_widths = [0, image_width] if image_width > 0 else [0]

1042 return self._trace_dispatch(self.dispatch)(

1043 _combine_topk_swa_warmup_inputs(vllm_config),

1044 _COMBINE_TOPK_SWA_POINTER_INPUTS,

1045 WINDOW_SIZE=window_size,

1079 - image_width=image_widths,

1046 )

1047

1048 def compile(self, compile_key: CompileKey) -> None:

@@ -1093,14 +1059,11 @@ class CombineTopkSwaIndicesKernel(

1059 input_variant.pointer("query_start_loc", torch.int32),

1060 input_variant.pointer("seq_lens", torch.int32),

1061 input_variant.pointer("gather_lens", torch.int32),

1096 - int32_ptr,

1097 - int32_ptr,

1062 1, # do not specialize M

1063 1, # do not specialize N

1064 TOP_K=compile_key.TOP_K,

1065 COMPRESS_RATIO=compile_key.COMPRESS_RATIO,

1066 WINDOW_SIZE=compile_key.WINDOW_SIZE,

1103 - IMAGE_WIDTH=compile_key.IMAGE_WIDTH,

1067 PADDED_TOP_K=compile_key.PADDED_TOP_K,

1068 grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),

1069 )

@@ -1119,13 +1082,8 @@ class CombineTopkSwaIndicesKernel(

1082 TOP_K: int,

1083 COMPRESS_RATIO: int,

1084 WINDOW_SIZE: int,

1122 - left_visible: torch.Tensor | None = None,

1123 - right_visible: torch.Tensor | None = None,

1124 - max_image_tokens: int = 0,

1085 ) -> None:

1086 num_reqs = seq_lens.shape[0]

1127 - has_image = left_visible is not None

1128 - image_width = max_image_tokens if has_image else 0

1087 self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](

1088 combined_indices,

1089 combined_indices.stride(0),

@@ -1135,14 +1093,11 @@ class CombineTopkSwaIndicesKernel(

1093 query_start_loc,

1094 seq_lens,

1095 gather_lens,

1138 - left_visible if has_image else topk_indices,

1139 - right_visible if has_image else topk_indices,

1096 M,

1097 N,

1098 TOP_K=TOP_K,

1099 COMPRESS_RATIO=COMPRESS_RATIO,

1100 WINDOW_SIZE=WINDOW_SIZE,

1145 - IMAGE_WIDTH=image_width,

1101 PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),

1102 )

1103

@@ -1169,9 +1124,6 @@ def build_flashinfer_mixed_sparse_indices(

1124 decode_is_valid_token: torch.Tensor | None = None,

1125 swa_block_span: int | None = None,

1126 compressed_block_span: int | None = None,

1172 - prefill_left_visible: torch.Tensor | None = None,

1173 - prefill_right_visible: torch.Tensor | None = None,

1174 - max_image_tokens: int = 0,

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

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

1129

@@ -1181,21 +1133,11 @@ def build_flashinfer_mixed_sparse_indices(

1133 per token). Decode tokens read precomputed SWA/compressed indices; prefill

1134 tokens derive their SWA window from the position and translate local

1135 compressed indices to global slots via the block tables.

1184 -

1185 - When ``prefill_left_visible``/``prefill_right_visible`` are given (vision

1186 - variant), the SWA column region widens by ``max_image_tokens`` and prefill

1187 - tokens inside an image span get a bidirectionally widened window; decode

1188 - rows are padded with -1 across the extra columns.

1136 """

1137 assert decode_swa_indices.dtype == torch.int32

1138 assert decode_swa_indices.dim() == 2

1192 - swa_index_width = decode_swa_indices.shape[-1]

1193 - assert swa_index_width >= window_size

1194 - has_image = prefill_left_visible is not None

1195 - image_width = max_image_tokens if has_image else 0

1196 - swa_total_width = swa_index_width + image_width

1197 - if has_image:

1198 - assert prefill_right_visible is not None

1139 + swa_total_width = decode_swa_indices.shape[-1]

1140 + assert swa_total_width >= window_size

1141 if decode_compressed_topk_lens is not None:

1142 assert decode_compressed_topk_lens.dtype == torch.int32

1143 assert prefill_topk_indices.dtype == torch.int32

@@ -1254,7 +1196,7 @@ def build_flashinfer_mixed_sparse_indices(

1196 if num_tokens == 0:

1197 return sparse_indices, sparse_topk_lens

1198

1257 - window_block_size = triton.next_power_of_2(max(swa_index_width, 1))

1199 + window_block_size = triton.next_power_of_2(max(swa_total_width, 1))

1200 topk_block_size = triton.next_power_of_2(max(padded_topk, 1))

1201 max_block_size = max(window_block_size, topk_block_size)

1202 num_warps = 4 if max_block_size >= 256 else 1

@@ -1278,8 +1220,6 @@ def build_flashinfer_mixed_sparse_indices(

1220 decode_is_valid_token,

1221 prefill_topk_indices,

1222 prefill_topk_indices.stride(0),

1281 - prefill_left_visible if has_image else token_to_req_indices,

1282 - prefill_right_visible if has_image else token_to_req_indices,

1223 query_start_loc,

1224 seq_lens,

1225 token_to_req_indices,

@@ -1293,8 +1233,7 @@ def build_flashinfer_mixed_sparse_indices(

1233 compressed_span,

1234 NUM_DECODE_TOKENS=num_decode_tokens,

1235 WINDOW_SIZE=window_size,

1296 - SWA_INDEX_WIDTH=swa_index_width,

1297 - IMAGE_WIDTH=image_width,

1236 + SWA_INDEX_WIDTH=swa_total_width,

1237 COMPRESS_RATIO=compress_ratio,

1238 TOP_K=topk,

1239 PADDED_TOP_K=padded_topk,

@@ -1349,8 +1288,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1288 decode_is_valid_token_ptr,

1289 prefill_topk_indices_ptr,

1290 prefill_topk_stride,

1352 - left_visible_ptr,

1353 - right_visible_ptr,

1291 query_start_loc_ptr,

1292 seq_lens_ptr,

1293 token_to_req_indices_ptr,

@@ -1365,7 +1302,6 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1302 NUM_DECODE_TOKENS,

1303 WINDOW_SIZE: tl.constexpr,

1304 SWA_INDEX_WIDTH: tl.constexpr,

1368 - IMAGE_WIDTH: tl.constexpr,

1305 COMPRESS_RATIO: tl.constexpr,

1306 TOP_K: tl.constexpr,

1307 PADDED_TOP_K: tl.constexpr,

@@ -1377,18 +1313,14 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1313 TOPK_BLOCK_SIZE: tl.constexpr,

1314 ):

1315 token_idx = tl.program_id(0)

1380 - # Total SWA column count; > SWA_INDEX_WIDTH only for in-image

1381 - # bidirectional visibility (vision variant), where the extra columns are

1382 - # -1-padded for decode rows.

1383 - SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH

1316

1317 if token_idx < NUM_DECODE_TOKENS:

1386 - for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):

1318 + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):

1319 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)

1388 - mask = offset < SWA_TOTAL_WIDTH

1320 + mask = offset < SWA_INDEX_WIDTH

1321 values = tl.load(

1322 decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,

1391 - mask=offset < SWA_INDEX_WIDTH,

1323 + mask=mask,

1324 other=-1,

1325 )

1326 values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)

@@ -1431,7 +1363,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1363 tl.store(

1364 sparse_indices_ptr

1365 + token_idx * sparse_indices_stride

1434 - + SWA_TOTAL_WIDTH

1366 + + SWA_INDEX_WIDTH

1367 + offset,

1368 values,

1369 mask=mask,

@@ -1445,7 +1377,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1377 else:

1378 compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)

1379

1448 - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)

1380 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len)

1381 return

1382

1383 prefill_idx = token_idx - NUM_DECODE_TOKENS

@@ -1457,26 +1389,17 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1389 start_pos = seq_len - query_len

1390 token_idx_in_query = token_idx - query_start

1391 pos = start_pos + token_idx_in_query

1460 - if IMAGE_WIDTH > 0:

1461 - # In-image bidirectional visibility: window starts up to

1462 - # max(left - (window - 1), 0) earlier and extends `right` past pos.

1463 - left = tl.load(left_visible_ptr + token_idx)

1464 - right = tl.load(right_visible_ptr + token_idx)

1465 - else:

1466 - left = 0

1467 - right = 0

1468 - left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)

1469 - swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)

1470 - swa_len = pos + right - swa_start_pos + 1

1392 + swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1), 0)

1393 + swa_len = pos - swa_start_pos + 1

1394 if COMPRESS_RATIO > 0:

1395 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)

1396 else:

1397 # SWA-only layers (ratio 0): no indexer topk; div by constexpr 0 is UB.

1398 topk_len = 0

1399

1477 - for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):

1400 + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):

1401 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)

1479 - mask = offset < SWA_TOTAL_WIDTH

1402 + mask = offset < SWA_INDEX_WIDTH

1403 pos_offset = swa_start_pos + offset

1404 block_indices = pos_offset // swa_block_size

1405 block_numbers = tl.load(

@@ -1520,10 +1443,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel(

1443 tl.store(

1444 sparse_indices_ptr

1445 + token_idx * sparse_indices_stride

1523 - + SWA_TOTAL_WIDTH

1446 + + SWA_INDEX_WIDTH

1447 + offset,

1448 slot_ids,

1449 mask=mask,

1450 )

1451

1529 - tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)

1452 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len)

vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py

+2 -17

@@ -311,7 +311,7 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):

311 else:

312 assert self.topk_indices_buffer is not None

313 top_k = self.topk_indices_buffer.shape[-1]

314 - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)

314 + combined_topk = round_up(top_k + self.window_size, 128)

315 current_workspace_manager().get_simultaneous(

316 ((self.PREFILL_CHUNK_SIZE, m, q.shape[-1]), torch.bfloat16),

317 ((self.max_num_batched_tokens, combined_topk), torch.int32),

@@ -411,7 +411,7 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):

411 )

412 assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"

413 workspace_manager = current_workspace_manager()

414 - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)

414 + combined_topk = round_up(top_k + self.window_size, 128)

415 for chunk_start, chunk_end, chunk_n, chunk_m in chunk_plan:

416 chunk_size = chunk_end - chunk_start

417 kv_ws, idx_ws, lens_ws = workspace_manager.get_simultaneous(

@@ -457,21 +457,6 @@ class DeepseekV4MegaAttnAttention(DeepseekV4FlashMLAAttention):

457 chunk_m,

458 chunk_n,

459 out=(idx_ws[: qe - qs], lens_ws[: qe - qs]),

460 - left_visible=(

461 - swa_metadata.prefill_left_visible[

462 - num_decode_tokens + qs : num_decode_tokens + qe

463 - ]

464 - if swa_metadata.prefill_left_visible is not None

465 - else None

466 - ),

467 - right_visible=(

468 - swa_metadata.prefill_right_visible[

469 - num_decode_tokens + qs : num_decode_tokens + qe

470 - ]

471 - if swa_metadata.prefill_right_visible is not None

472 - else None

473 - ),

474 - max_image_tokens=self.max_image_tokens,

460 )

461 chunk_out = _token_slice(out, token_base + qs, token_base + qe)

462 # Mega attention over the gathered non-paged bf16 KV (RoPE already

vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py

-4

@@ -417,10 +417,6 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):

417 decode_is_valid_token=decode_is_valid_token,

418 swa_block_span=swa_block_span,

419 compressed_block_span=compressed_block_span,

420 - prefill_left_visible=swa_metadata.prefill_left_visible,

421 - prefill_right_visible=swa_metadata.prefill_right_visible,

422 - # getattr for tests that bypass __init__ via object.__new__.

423 - max_image_tokens=getattr(self, "max_image_tokens", 0),

420 )

421 if swa_only:

422 swa_metadata.flashinfer_sparse_index_cache["swa_only"] = (

vllm/models/deepseek_v41/nvidia/flashmla.py

+2 -19

@@ -120,9 +120,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

120 else:

121 assert self.topk_indices_buffer is not None

122 top_k = self.topk_indices_buffer.shape[-1]

123 - combined_topk = round_up(

124 - top_k + self.window_size + self.max_image_tokens, 128

125 - )

123 + combined_topk = round_up(top_k + self.window_size, 128)

124 current_workspace_manager().get_simultaneous(

125 ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),

126 ((self.max_num_batched_tokens, combined_topk), torch.int32),

@@ -303,7 +301,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

301 )

302 assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"

303 workspace_manager = current_workspace_manager()

306 - combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)

304 + combined_topk = round_up(top_k + self.window_size, 128)

305 for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:

306 chunk_size = chunk_end - chunk_start

307 workspace = workspace_manager.get_simultaneous(

@@ -361,21 +359,6 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):

359 chunk_M,

360 chunk_N,

361 out=(combined_indices_out, combined_lens_out),

364 - left_visible=(

365 - swa_metadata.prefill_left_visible[

366 - num_decode_tokens + query_start : num_decode_tokens + query_end

367 - ]

368 - if swa_metadata.prefill_left_visible is not None

369 - else None

370 - ),

371 - right_visible=(

372 - swa_metadata.prefill_right_visible[

373 - num_decode_tokens + query_start : num_decode_tokens + query_end

374 - ]

375 - if swa_metadata.prefill_right_visible is not None

376 - else None

377 - ),

378 - max_image_tokens=self.max_image_tokens,

362 )

363 flash_mla_sparse_fwd(

364 q=q[query_start:query_end],

vllm/transformers_utils/configs/deepseek_v41.py

-8

@@ -60,11 +60,3 @@ class DeepseekV41Config(PretrainedConfig):

60 self.vision_max_n_token = vision_config.get("max_image_tokens", 1024)

61 self.vision_min_pixels = vision_config.get("min_pixels", 295936)

62 self.vision_max_wh_ratio = vision_config.get("max_wh_ratio")

63 - # The vision variant needs the mm-prefix plumbing (atomic image-span

64 - # prefill + sparse-SWA window widening); the base arch-config

65 - # convertor reads this config-level attribute.

66 - self.is_mm_prefix_lm = vision_n_layers > 0

67 - # The sparse-SWA index kernels widen the window within image spans

68 - # in-kernel, so mm-prefix ranges longer than sliding_window must be

69 - # kept (they are only consumed by those kernels).

70 - self.mm_prefix_clamp_sliding_window = vision_n_layers > 0

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

+6 -4

@@ -452,13 +452,15 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

452 assert hasattr(hf_config, "sliding_window")

453 self.window_size = hf_config.sliding_window

454

455 - # Vision variant: image spans (up to vision_max_n_token tokens) are

455 + # V4 vision variant: image spans (up to vision_max_n_token tokens) are

456 # visible bidirectionally, so prefill index rows widen from

457 - # window_size to window_size + max_image_tokens. Text-only models keep

458 - # max_image_tokens == 0 and take the original code paths everywhere.

457 + # window_size to window_size + max_image_tokens. The V4 config sets

458 + # mm_prefix_clamp_sliding_window exactly for these in-kernel-widened

459 + # ranges; V4.1 image tokens use the plain causal window, and text-only

460 + # models keep max_image_tokens == 0 everywhere.

461 self.max_image_tokens = (

462 getattr(hf_config, "vision_max_n_token", 0)

461 - if getattr(hf_config, "vision_n_layers", 0) > 0

463 + if getattr(hf_config, "mm_prefix_clamp_sliding_window", False)

464 else 0

465 )

466 self.prefill_index_width = self.window_size + self.max_image_tokens