Repository navigation
[Bugfix][V1] Null sliding-window block-table entries below the window - #60767
Conversation
|
HiSparse added |
TheEpicDolphin
left a comment
There was a problem hiding this comment.
I like the general approach, but it seems like KpoolTailSpec, which inherits from SlidingWindowSpec, would break from this change. It uses a one block ring buffer, and when first_pos gets large enough it will null that one block table entry.
Maybe we can add a num_retained_tokens method to both, and have it return None for KpoolTailSpec so that it opts out of the optimization? Just an idea.
| query_start_loc: torch.Tensor | None = None, | ||
| seq_lens: torch.Tensor | None = None, |
There was a problem hiding this comment.
RequestStates has a num_computed_tokens property that already contains the first token pos for each sequence. You may be able to pass that in instead of query_start_loc and seq_lens together.
There was a problem hiding this comment.
Fixed in new revision. num_computed_tokens is used here which got passed via num_computed_tokens=self.req_states.num_computed_tokens.gpu
| block_ids = tl.load(src_row_ptr + offset, mask=offset < num_blocks) | ||
| block_ids = tl.where(offset < num_null, 0, block_ids) |
There was a problem hiding this comment.
nit: can simplify this down to:
block_ids = tl.load(
src_row_ptr + offset,
mask=(offset >= num_null) & (offset < num_blocks),
other=0,
)
There was a problem hiding this comment.
Thanks! Addressed in new revision
|
@TheEpicDolphin thanks for the review, all three addressed in 7a659e9.
|
|
This pull request has merge conflicts that must be resolved before it can be |
The sliding-window manager frees blocks that fall wholly outside a request's window, but the worker's persistent block table keeps the freed block ids. Attention kernels load KV in tiles (FA4: 128 keys on SM100); with 64-token pages the tile that holds the window start can include a page that lies wholly below the window. Its positions are masked, but masked NaN still propagates (0 * NaN). Once the freed block is reused, its bytes can decode as NaN in the layer that reads it, and that request's attention output becomes NaN. When gathering the forward-pass block tables, replace every sliding-window entry that lies wholly below first_scheduled_position - SlidingWindowSpec.num_retained_tokens (sliding_window - 1 + extra_retained_tokens) with the null block. No query of the step can attend those positions, and the cutoff is never below the manager's free boundary, so every block it may have freed is covered. The first scheduled position is RequestState.num_computed_tokens, which the runner already holds per request slot, so that tensor is passed to gather_block_tables and the nulling is folded into the load mask. Only the forward-pass copy changes; slot mappings and the persistent tables are untouched. Disabled under context parallelism. Fixes vllm-project#60560. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Stu Cao <stu@primeintellect.ai>
7a659e9 to
16d9ae0
Compare
TheEpicDolphin
left a comment
There was a problem hiding this comment.
Thanks for addressing the feedback, this looks good to me. I think it's a relatively safe short-term fix.
|
✅ @stu-cao, CI is now available for this PR.
|
|
hey @stu-cao, please run the CI on this PR when you get the chance |
|
/ci run |
|
✅ Triggered Buildkite CI #95000 for commit ℹ️ This PR is 98 commits behind upstream |
|
✅ Triggered Buildkite CI #95559 for commit ℹ️ This PR is 63 commits behind upstream |
|
This pull request has merge conflicts that must be resolved before it can be |
Resolves the overlap with vllm-project#61174 (stale shared block-table rows): the gather keeps main's num_rows / _num_dirty_rows handling and this branch's sliding-window masking arguments; both sides' new tests are kept. Signed-off-by: Stu Cao <stu@primeintellect.ai>
|
/ci run |
|
✅ Triggered Buildkite CI #95604 for commit ℹ️ This PR is 1 commit behind upstream |
|
/ci run |
|
❌ CI already ran on this commit: https://buildkite.com/vllm/ci/builds/95604. No new CI build was started. Push new changes to run CI again, use |
Picks up vllm-project#62213 (order-dependent ROCm aiter sampler test), which was the only failure in CI build 95604. Signed-off-by: Stu Cao <stu@primeintellect.ai>
|
/ci run |
|
✅ Triggered Buildkite CI #95657 for commit |
Picks up vllm-project#62221 (Rust frontend version mismatch) and vllm-project#62274 (ROCm inductor batch_invariant guard), the only failures in CI build 95657. Signed-off-by: Stu Cao <stu@primeintellect.ai>
|
/ci run |
|
✅ Triggered Buildkite CI #95780 for commit |
|
@TheEpicDolphin Could you please help merge? CI is all green |
BlockTables.gather_block_tables now nulls entries below a sliding window and tracks dirty rows (upstream vllm-project#60767); the fake tables carry the fields of a table with no sliding-window group.
### What this PR does / why we need it?
Port `vllm-ascend` to vLLM main `544f182d` (bumped from `c0b6f74`) while
keeping the
`v0.31.0` release lane green. The upstream window drops
`replace_parameter(..., prefer_copy=...)`,
fuses the Kimi-K3 output gate into the QKV-A projection, adds the ROCm
Kimi-K3
specialized-decode hooks, widens the EPLB expert-weight and
rejection-sample contracts, and
adds the DP-drafter `num_reqs_override`, the MRv2 `skip_sample` drafting
skip, the HiSparse
`topk_buffer_width`, the SimpleCPUOffload CPU-block RPC and the abstract
`AsyncModelRunnerOutput.get_kv_connector_output`. Every gate below is
either a
`vllm_version_is("0.31.0")` split or an attribute/signature probe.
| Files | Upstream vLLM change (why it breaks vllm-ascend) | vllm-ascend
adaptation |
|-------|--------------------------------------------------|------------------------|
| `.github/vllm-main-verified.commit` | — (main-lane pin only) |
Verified main pin bumped `c0b6f74` -> `544f182d` |
| `tests/ut/attention/test_mla_prolog_v3_route.py` |
[vllm#60805](vllm-project/vllm#60805) makes
`MLAAttention.get_kv_cache_spec` read `topk_buffer_width` (set in
`__init__`);
[vllm#59634](vllm-project/vllm#59634) makes
`_mla_preprocess` return a 3-tuple | Set `owner.topk_buffer_width =
None` on the stub; mock `_mla_preprocess` as `(None, None, None)` |
| `tests/ut/attention/test_mla_v1.py` |
[vllm#59634](vllm-project/vllm#59634) fuses the
Kimi-K3 output gate as a third QKV-A shard | Add tests for
`qkv_a_gate_fused`/`use_output_gate` and `_split_fused_qkv_lora` gate
peeling; unpack the 3-tuple |
| `tests/ut/models/minimax_m3/test_minimax_m3_vit.py` |
[vllm#60238](vllm-project/vllm#60238) makes
`_mark_tower_model` read `use_v2_model_runner` | Add
`use_v2_model_runner=False` to the constructed runner stub |
| `tests/ut/models/test_kimi_k3_adapter.py` |
[vllm#60683](vllm-project/vllm#60683) adds the
ROCm specialized-decode hooks (`init_specialized_decode`,
`get_model_state_cls`, the `step` forward argument, and `layer.kv_cache`
reassignment via `clear_layer_kv_caches`) | Add main-only tests (`skipif
vllm_version_is("0.31.0")`) covering the inherited hooks |
| `tests/ut/ops/test_routed_experts.py` |
[vllm#61786](vllm-project/vllm#61786) registers
expert weights via `w.view(torch.uint8) if w.element_size() == 1 else w`
| New UT asserting `EplbExpertTensorList.element_size()/view()` keep the
container and its aliasing semantics |
| `tests/ut/patch/platform/test_patch_structured_output.py` |
[vllm#49625](vllm-project/vllm#49625) makes
`StructuredOutputManager.grammar_init` read `backend_name` | Set
`manager.backend_name = None` on the stub |
| `tests/ut/runtime_guard/test_unified_wiring.py` |
[vllm#56430](vllm-project/vllm#56430) adds
`AsyncModelRunnerOutput.get_kv_connector_output` | Add a stub
`get_kv_connector_output` to the fake async output |
| `tests/ut/simple_kv_offload/test_simple_kv_offload.py` |
[vllm#52921](vllm-project/vllm#52921) adds the
`get_simple_cpu_offload_num_cpu_blocks` worker RPC | New main-only UTs
for the aligned / no-group / no-connector paths |
| `tests/ut/spec_decode/test_dflash2_v2.py` |
[vllm#58984](vllm-project/vllm#58984) adds
`num_reqs_override` to `prepare_dflash_inputs_triton` | Update the
expected kwarg; new test for positional forwarding |
| `tests/ut/worker/test_model_runner_v1.py` |
[vllm#60805](vllm-project/vllm#60805)
`topk_buffer_width` | Set `layer.topk_buffer_width = None` on the layer
stub |
| `tests/ut/worker/test_model_runner_v2.py` | Same
[vllm#60805](vllm-project/vllm#60805) | Same
stub attribute |
| `tests/ut/worker/v2/test_speculator.py` |
[vllm#62784](vllm-project/vllm#62784) adds
`skip_sample` to the drafter `_prefill` | Expect/forward `skip_sample`
on main (7 -> 8 positional args) |
| `tests/ut/worker/v2/test_v2_block_table.py` |
[vllm#60767](vllm-project/vllm#60767) makes
`init_block_table_layout_tensors` read `_sliding_window_reach` | Set
`tables._sliding_window_reach = [0, 0]` |
| `vllm_ascend/attention/mla_v1.py` |
[vllm#59634](vllm-project/vllm#59634) fuses the
Kimi-K3 output gate into the QKV-A projection (three shards, no
standalone `g_proj`);
[vllm#62326](vllm-project/vllm#62326) drops
`replace_parameter(..., prefer_copy=)` | Detect the gate-fused layout
from `fused_qkv_a_proj.output_sizes`, disable MLAPO/PROLOG_V3 for it,
bound the KV-A slice with `_fused_qkv_a_kv_end`, peel the gate in
`_split_fused_qkv_lora`; version-split `replace_parameter` |
| `vllm_ascend/attention/sfa_v1.py` |
[vllm#62326](vllm-project/vllm#62326) drops
`replace_parameter(..., prefer_copy=)` | Version-split
`replace_parameter` (keep `prefer_copy=True` on v0.31.0) |
| `vllm_ascend/models/deepseek_v4/indexer.py` |
[vllm#52670](vllm-project/vllm#52670) keeps
`STR_DTYPE_TO_TORCH_DTYPE["fp8"]` as the uint8 storage dtype, so
`kv_cache_dtype_str_to_dtype` cannot validate the indexer's fp8 compute
dtype | Validate `indexer_kv_dtype` against the platform default string
and return the compute dtype directly |
| `vllm_ascend/models/kimi_k3.py` |
[vllm#60683](vllm-project/vllm#60683) adds the
ROCm specialized-decode hooks and reassigns `layer.kv_cache` on teardown
| Add a `kv_cache` setter, call `init_specialized_decode` + override
`get_model_state_cls` to the Ascend hybrid state on main, forward `step`
|
| `vllm_ascend/observability/runtime_guard/hooks.py` |
[vllm#56430](vllm-project/vllm#56430) makes
`AsyncModelRunnerOutput.get_kv_connector_output` required | Implement it
on `AscendAsyncOutput` (main only) |
| `vllm_ascend/ops/fused_moe/routed_experts.py` |
[vllm#61786](vllm-project/vllm#61786) moves
1-byte expert tensors as uint8 views | Add `element_size()`/`view()` to
`EplbExpertTensorList`; widen the `get_expert_weights` return type |
| `vllm_ascend/ops/gdn.py` |
[vllm#62326](vllm-project/vllm#62326) drops
`replace_parameter(..., prefer_copy=)` | Version-split
`replace_parameter` for the packed conv weight |
| `vllm_ascend/ops/kimi_kda.py` | Same
[vllm#62326](vllm-project/vllm#62326) | Same
version-split for the packed conv weight |
| `vllm_ascend/ops/linear.py` | Same
[vllm#62326](vllm-project/vllm#62326) | Same
version-split for `weight_fp32` |
| `vllm_ascend/ops/triton/v2/spec_decode/prepare_dflash_inputs.py` |
[vllm#58984](vllm-project/vllm#58984) forwards
the DP-synced request count during dummy runs | Accept
`num_reqs_override`, cap `num_reqs` and slice `num_scheduled_tokens` |
| `vllm_ascend/worker/v2/spec_decode/autoregressive/speculator.py` |
[vllm#62784](vllm-project/vllm#62784) adds
`skip_sample` to `_prefill` | Version-split `_prefill`, forward
`skip_sample` on main |
| `vllm_ascend/worker/v2/spec_decode/dflash/speculator.py` |
[vllm#58984](vllm-project/vllm#58984) adds
`num_reqs_override` | Forward `num_reqs_override` through both
`prepare_dflash_inputs` wrappers |
| `vllm_ascend/worker/v2/spec_decode/gemma4/speculator.py` |
[vllm#62784](vllm-project/vllm#62784) adds
`skip_sample` | Accept and forward `skip_sample` |
| `vllm_ascend/worker/v2/spec_decode/rejection_sampler_utils.py` |
[vllm#53987](vllm-project/vllm#53987) adds the
FLy entropy-gated verification mode; upstream also threads the
watermarking kwargs | Accept both kwarg groups, but raise
`NotImplementedError` for FLy and watermarking on NPU |
| `vllm_ascend/worker/worker.py` |
[vllm#52921](vllm-project/vllm#52921) adds the
`get_simple_cpu_offload_num_cpu_blocks` RPC to the worker | Implement it
(main only) via
`_find_simple_cpu_offload_connector().get_aligned_num_cpu_blocks()` |
### Does this PR introduce _any_ user-facing change?
### How was this patch tested?
---
### Checklist
#### Testing
- [ ] **Boundary / edge-case coverage**: I reviewed whether this change
involves
boundary conditions (e.g. empty input, single/zero element, min/max
values,
off-by-one, sequence-length limits, integer overflow). Where applicable,
I
added unit tests covering these edge cases in `tests/ut/` or
`tests/e2e/`.
- [ ] **Regression test**: If this PR fixes a bug, I added a regression
test
that would fail without the fix.
- [ ] **Error / failure path coverage**: I covered abnormal paths
(exceptions,
timeouts, out-of-range, division by zero) in addition to the happy path.
- [ ] **Meaningful & reliable tests**: New tests assert real behavior
and are
not written merely to satisfy CI.
- [ ] **Performance impact**: If this PR touches a performance-sensitive
path,
I attached benchmark results and added performance regression tests.
- vLLM main:
vllm-project/vllm@c0b6f74
---------
Signed-off-by: zhangxinyuehfad <starmoon_zhang@163.com>
Signed-off-by: hfadzxy <starmoon_zhang@163.com>
Both accuracy bugs in ROCM_ATTN's Triton fallback that this avoided are fixed on main: the sliding-window NaN (vllm-project#62828) no longer occurs since vllm-project#60767 nulls below-window block-table entries, and the gfx942 unwritten output heads were fixed by vllm-project#63268. What remains is a performance change, which belongs in vllm-project#61433. Keep only the gfx950 rule for vllm-project#62606. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Chen, Todd <zhenchen@amd.com>
…vllm-project#60767) Signed-off-by: Stu Cao <stu@primeintellect.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: vllm-agent <claw@inferact.ai>
Purpose
Fixes #60560.
The sliding-window KV manager frees blocks once they fall wholly outside a request's window, but the worker's persistent block table keeps the freed block ids. Attention kernels load KV in tiles (FA4 uses 128-key tiles on SM100). With 64-token pages, the tile that holds the window start can include one page lying wholly below the window. That page is masked, but masked NaN still propagates (0 × NaN). After the freed block is reused, its bytes can decode as NaN in the layer that reads it, and that request's attention output turns NaN.
#60560 has the full analysis. In short, on B300 with a DFlash2 drafter (all 6 layers sliding window, 2048) under the hybrid KV cache manager and
--block-size 64:This PR makes the worker never hand those entries to attention. When
gather_block_tablesbuilds the forward-pass tables, every sliding-window entry wholly belowfirst_scheduled_position - (sliding_window - 1 + extra_retained_tokens)is replaced by the null block:num_computed_tokens, already on the device, so there is no host sync.SlidingWindowSpec.num_retained_tokensgives the reach. (KpoolTailSpec, which an earlier revision opted out, was removed onmainin [KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots #57169; its successorCircularBufferSpecis not aSlidingWindowSpecand is unaffected.)cp_size > 1, where a block-table column is not a contiguous token range.Speculator steps that reuse the gathered tables (e.g. DFlash) are covered too, since their queries sit at or after the target's first scheduled position.
Scope
The change is not specific to any model or drafter. It applies to every KV cache group whose spec is a
SlidingWindowSpec(includingSlidingWindowMLASpec), target or draft, with any attention backend. Why it is safe:compute_slot_mappings), or, for DFlash, from this step's positions, which are above the cutoff.BlockTablesholds one entry per KV block, so the cutoff ismax(first_pos - reach, 0) // block_sizeentries.Not covered:
Assumption: the null block holds no NaN. The KV cache is zero-allocated, and the null entries the tables already contain (sliding-window prefix hits, P/D tail clipping) rely on the same thing.
Validation: end-to-end validation so far covers a DFlash2 drafter under the hybrid KV cache manager. A check on a target model with sliding-window layers (e.g. gpt-oss, Gemma 3) at high concurrency is still to do.
Relation to existing PRs
SchedulerOutputchange, and no table resend. If maintainers prefer Rewrite worker block tables after SWA eviction #42276's direction, I'm happy to close this one in its favour, and the regression test here should carry over.Test Plan
pytest tests/v1/worker/test_gpu_block_table.py: the newtest_gather_block_tables_nulls_blocks_below_sliding_windowchecks:VLLM_KV_CACHE_LAYOUT=BLHNC,--block-size 64, hybrid KV cache manager on, decode prefix caching off. Agentic multi-turn traffic (AgentX), 32 concurrent sessions, 900 s. Acceptance = accepted / drafted tokens over the profiling window.pre-commit run ruff-checkandpre-commit run mypyon the changed files.Test Result
main791f92135a([KV Cache] GLM-5.3-Flash: use the generic packed KV layout; CircularBufferSpec tail with speculative ring slots #57169 removedKpoolTailSpecand kernel block sizes; the two review suggestions from @TheEpicDolphin are in): 16 passed, 0 failed on B300 (2026-10-14;test_gpu_block_table.pyincl. the new test, plus the selectedtest_kv_cache_utilscases).The end-to-end run used this exact diff on our serving build (vLLM 0.30.1-based), where it applies unchanged, because that stack (NVFP4 KV, NIXL whole-row registration) is not on
mainyet.AI assistance (Claude) was used to investigate the bug and to draft this change and its test. I reviewed every changed line and the test results.
🤖 Generated with Claude Code