Conversation
Adds a config-gated attention variant where the prompt/prefix stays globally attended while generated tokens additionally attend only a fixed window of recent tokens. Bounds per-request decode-side KV at O(prefix + window) instead of growing linearly with generation length: gap blocks between the prefill tail and the live decode window are evicted each step and masked out element-wise in the kernel. Off by default; enabled per model via `decode_sliding_window` in the model config. Validated on DeepSeek-OCR. - New `PrefixAnchoredSWASpec` / `PrefixAnchoredSWAManager` (gap-block eviction) in the v1 KV-cache layer. - Mask support in the FLASH_ATTN_V100 paged kernels via a template parameter: the non-masked instantiations are SASS-byte-identical to the previous build, so existing models run unchanged machine code. - Backend guard: windowed models require FLASH_ATTN_V100 and refuse silently-unsupported backends at startup. - Prefix caching is force-disabled when the window is active (windowed decode KV is not reusable across requests). - Full CUDA-graph capture downgrades to PIECEWISE when the window is set (full-graph capture faults with mid-stream block eviction). - Tests: mask formula vs torch reference, masked-kernel device tests, KV-manager gap-eviction unit tests, spec merge/registration, guard behavior. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Addresses the PR-attributable pre-commit failures; no behavioral change. - vllm/config/vllm.py: rewrap an over-long docstring line (E501). - tests/kernels/attention/test_sm70_flash_v100_anchored_swa.py: type the module-level `flash_ops` fallback and unpack it through a narrowing helper so mypy no longer indexes `tuple | None`; ruff-format. - vllm/v1/attention/backends/flash_attn_v100.py: in `build_for_cudagraph_capture`, cast to `FlashAttnV100Metadata` before assigning the anchor fields (mypy attr-defined); ruff-format. - flash-attention-v100/kernel/flash_decode_paged.cu: clang-format the two hunks this PR added (whitespace only). The pre-existing whole-file clang-format drift in fused_mha.h and fused_mha_forward_paged.cu (present at the merge-base) is left untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
yangzhuxinyzx
left a comment
There was a problem hiding this comment.
Thanks for the careful kernel and manager work. I found two correctness blockers in the integration path:
-
The anchored metadata is only attached in FlashAttnV100MetadataBuilder.build_for_cudagraph_capture(). The normal build() path calls _attach_common_flash_metadata(), but that helper does not copy common_attn_metadata.prefix_anchor_lens or set decode_sliding_window; build_for_drafting() has the same omission. TritonAttentionMetadataBuilder.build() does not propagate arbitrary CommonAttentionMetadata fields. A direct source-level probe with CommonAttentionMetadata(prefix_anchor_lens=[16]) produced regular Flash metadata with neither prefix_anchor_lens nor decode_sliding_window. Therefore _anchored_swa_params() returns (None, 0) on ordinary eager/PIECEWISE runtime metadata. This PR itself downgrades FULL graphs to PIECEWISE, so the KV manager can evict the gap while the attention path runs unmasked. Please attach the persistent anchor buffer/window in every runtime metadata path (preferably through one shared helper), including drafting if supported, and add an integration test that exercises builder.build(), not only direct kernel calls.
-
decode_sliding_window is activated on non-None but never validated as positive. PrefixAnchoredSWAManager uses 0/negative values in its eviction arithmetic, while _anchored_swa_params() explicitly disables the mask for values <= 0. Repro with block_size=4, prompt=16, computed=28: window=0 evicts logical blocks [4,5,6] and window=-8 evicts [4,5,6,7,8], yet the backend reports mask inactive for both. Please reject non-integer/non-positive values before constraints/spec construction and add zero/negative validation tests.
Local audit results otherwise: merge against current main was clean; git diff --check passed; Ruff 0.14.0 check/format passed on all 26 changed Python files; the new mask and single-type KV-manager suites passed 14/14. I did not repeat the expensive model E2E evidence.
|
Current-main revalidation on exact head The focused anchor/mask manager tests pass (8 passed), but the integration defect remains: anchor/window fields are attached only in There are now two additional integration blockers against |
f8358be to
e0c9095
Compare
|
Superseded by #281, now merged as |
|
公开开发已恢复:本 PR 已取消归档、解锁并重新打开。为避免把旧基线上的问题直接带入最新 |
|
Latest-main deduplication audit completed at The actual PR head
The architecture-keyed serving recipe/default files from the rejected identity-default stack are absent from the current tree; |
What
Serving-quality improvements for long-generation OCR workloads, plus a
correctness fix for the prefix-anchored SWA capability merged in #281:
scheduled-token basis; under async scheduling this can free blocks the
worker has not yet computed past. Free on the processed-token basis
instead. Regression tests simulate the async gap.
architecture-keyed registry applied at the chat endpoint. A model can
declare its recommended sampling/extra-args recipe; requests may
override within declared bounds. DeepSeek-OCR registers its checkpoint
README recipe (n-gram blocker 30/90 with whitelist tokens, temperature
0, max_tokens 8192) so an unconfigured client gets the model's intended
behavior via plain chat completions.
model_type previously fell through to the DeepSeek-VL2 template, which
injects role markers the OCR checkpoint was not trained with; the
fallback now pins the raw-concatenation prompt (golden test included).
exact-pattern detector as a safety net for degenerate table pages,
request-overridable.
official modes (tiny/small/base/large/gundam) selectable via
mm_processor_kwargs: {"image_mode": ...}; multi-image servingprofile;
max_cropsplumbing. The serving default remains thecheckpoint's gundam configuration.
Why
On document OCR, the model's own recipe is load-bearing: without the
n-gram blocker and the correct raw prompt, long pages degenerate into
loops. These changes make the endpoint serve the checkpoint's intended
recipe by default while keeping every knob per-request-overridable within
bounds. The SWA fix closes an async-scheduling correctness gap in the
capability added by #281.
Validation
All device work on Tesla V100-SXM2 32 GB (sm_70), fp16, on current main.
(proves coexistence with the merged SWA capability), the new
eviction-basis regression tests, resolution-mode suite, template golden,
defaults/bounds and repetition-detection suites.
resolution-mode discrimination (default ≡ gundam as identity control;
base/large produce distinct image-token counts); document-page scoring
against known ground truth (per-mode table in tests receipt).
Serving DeepSeek-OCR (validated configuration)
Launch line used for every endpoint validation in this PR (Tesla
V100-SXM2 32 GB, fp16):
--logits-processorsloads the n-gram blocker class; the serving-defaultsregistry supplies its per-request parameters and fails startup with the
exact flag text if the class is missing (never silent looping).
--no-enable-prefix-caching/--mm-processor-cache-gb 0are thecheckpoint recipe's engine flags.
--limit-mm-per-prompt '{"image":8}'(the SM70 baseline defaults to 1 image per prompt).
--enforce-eageris the validated configuration; CUDA graphs measured~4x faster decode at ~1% window cost and can be enabled after a
graph-mode validation pass.
--tensor-parallel-sizemust divide 10 (1 or 2 on a 4-GPU host); 4fails at startup with the head-divisibility error.
Merge base
Branch applies and merges cleanly onto current
mainat34403018d([Kernel][SM70] Optimize E4M3 page800 long decode, #290); all validation in
this PR ran on that base.