[Bugfix][Model] Restore causal image SWA for DeepSeek V4.1

[Juntian777]

Summary

Restore causal sliding-window attention for image tokens in DeepSeek V4.1 VL.

V4.1 inherited the image visibility rules used by DeepSeek V4 Vision-Exp: tokens inside an image could attend to earlier image tokens outside the sliding window and to future image tokens. This differs from the official V4.1 reference implementation, where _window_kv uses get_window_topk_idxs and both image and text tokens attend to the same causal window. With a window of 128, query position p should see only [max(0, p - 127), p] through SWA.

The incorrect widening affects full-image prefill as well as continuation chunks. For example, at the beginning of image span [1900, 3948), the inherited rule exposes future image tokens through position 3947; the V4.1 reference permits only positions 1773 through 1900.

The key line in the official get_window_topk_idxs is:

idxs = torch.where(idxs > end, -1, idxs)

The reference masks every index greater than the query position with -1; this is a causal 128-token window, including the current token. Executing the published get_window_topk_idxs at query position 1900 produces exactly 128 indices, from 1773 through 1900, with no future positions.

Changes

• Disable multimodal prefix-LM/bidirectional-range configuration for V4.1.
• Make the shared SWA metadata builder respect an explicit is_mm_prefix_lm=False, keeping V4.1 metadata and paged indices causal.
• Set the V4.1 attention-side image window extension to zero. Image preprocessing, the vision encoder, and vision_max_n_token remain unchanged.
• Keep the existing causal gather lengths and workspace sizes. The final diff does not change NVIDIA or AMD workspace allocation.
• Add six regression cases: full prefill and continuation prefill, each with compression ratios 0/1/2. Verify causal key positions in both gathered and paged indices, gather/CPU-plan consistency, and absence of image visibility extensions.

V4 Vision-Exp retains its existing image visibility behavior.

Validation

Run on GB200:

CUDA_VISIBLE_DEVICES=1 .venv/bin/python -m pytest \
  tests/v1/attention/test_deepseek_v4_swa_visible.py \
  tests/kernels/attention/test_flashmla_sparse.py \
  -k 'not test_builder_' -q --tb=short

Result: 30 passed, 3 deselected. The three existing builder tests require an inaccessible gated Hugging Face configuration. The new regression uses a local configuration and requires no model download.

Changed-file pre-commit checks passed. A local FlashMLA numerical sanity check also matched the causal reference, retaining a gather length of 227 and a first-query SWA length of 128 for seq_len=4000, query_len=100, and image span [1900, 3948).

Full-model VL accuracy, generation-length evaluation, and AMD runtime validation were not run. The evidence here establishes attention-window correctness at the metadata/kernel level.

Related work

Open-PR searches for SWA/gather, V4.1/prefill, and V4.1/causal/image found no other matching fix. #56986 concerns cross-layer compressed-KV gather reuse; it does not correct the image attention-window semantics addressed here.

AI assistance

AI assistance was used for analysis, implementation, and validation.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @Juntian777, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[Isotr0py]

/ci run

[github-actions]

❌ This PR is 7 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[mergify]

Hi @Juntian777, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

Hi @Juntian777, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[Isotr0py]

/ci run --allow-stale

[github-actions]

✅ Triggered Buildkite CI #89358 for commit e62b3c4abbae.