iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Juntian Liu
2026-09-16 18:09:00 +00:00
committed by GitHub
co-authored by Codex Isotr0py
parent 651a88c09a
commit 9f9e1dac26
8 changed files with 122 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_deepseek_v4_swa_visible.py
+81 -6
View File
@@ -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))
vllm/models/deepseek_v41/attention.py
-8
View File
@@ -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:
vllm/models/deepseek_v41/common/ops/cache_utils.py
+31 -108
View File
@@ -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)
vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py
+2 -17
View File
@@ -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
vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py
-4
View File
@@ -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"] = (
vllm/models/deepseek_v41/nvidia/flashmla.py
+2 -19
View File
@@ -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],
vllm/transformers_utils/configs/deepseek_v41.py
-8
View File
@@ -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
vllm/v1/attention/backends/mla/sparse_swa.py
+6 -4
View File
@@ -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