Conversation
The fp8 prefill bridge, its tail workspace and the gather-dense KV workspace all allocated the grown buffer while the previous one was still referenced by the cache entry and the local, so both were resident across the growth. On a 4x V100 32GB long-context deployment that is a 134 MiB (128k) to 268 MiB (256k) transient spike per rank, and the allocation can fail on the memory its own predecessor is holding: the OOM path then silently falls back to the slow route for the rest of the process. This is the same shape as the upstream FlashMLA workspace manager fix, which vLLM addressed in #56902 by dropping the old buffer first. - drop the cache entry before allocating, so only one buffer is live - retry the grown allocation once after empty_cache, since the segment freed by the drop is smaller than the request and cannot be reused - cap gather-dense doubling at the block table width, i.e. the pages max_model_len can ever address, so the overshoot is bounded - make the two bridge getters CPU-safe like their siblings so the growth path is reachable from host-only tests Verified on llm252 in the ct252-e4m3-offload-c4 image: the patch applies to that build and the policy suite goes 123 passed / 34 failed to 128 passed / 29 failed, the delta being exactly the five new regression tests. Remaining failures are pre-existing in that environment. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The three Flash-V100 prefill workspace caches allocate the grown buffer while the
previous one is still referenced by both the cache entry and a local, so old and
new are resident at the same time across a growth:
_get_fp8_prefill_bridge_workspace_get_fp8_prefill_bridge_tail_workspace_get_prefill_gather_dense_workspaceThis is the same shape as the upstream FlashMLA workspace-manager bug
(vllm-project/vllm#56902), where the fix was to drop the old buffer first. Our
shared
WorkspaceManageralready does this (vllm/v1/worker/workspace.py);these three caches predate it and do not.
Two consequences on a 4x V100 32GB long-context deployment:
The
except torch.OutOfMemoryErrorpath then returnsNone, the one-shotwarning fires, and the request silently falls back to the slow route — for
the rest of the process, since the cache keeps the old undersized buffer and
every longer prefill retries the same doubled allocation.
_get_prefill_dense_splitkv3_workspacealready pops the entry beforeallocating; this makes the other three match it.
Change
buffer is live during a growth.
_allocate_growing_workspace: on OOM, release cached blocks and retry once.The segment freed by the drop is smaller than the grown request, so
empty_cacheis what actually recovers the fragmented headroom. Only runs onthe OOM path. The stale entry is deliberately not restored on failure, so a
later smaller request re-allocates from scratch instead of inheriting a
doubled capacity that already failed.
max_model_lencan ever address. Beyond that the overshoot is pure lossagainst the KV cache.
(
is_cudaguard on device/stream, tensor-aware capture check), so the growthpath is reachable from host-only tests.
Measurements (Tesla V100-PCIE-32GB,
head_dim=256, 1 KV head, block 16)Workspace footprint is 1 KiB/token; peak-over-entry during a growth:
The saving equals the previous workspace's size, so a 128k -> 256k growth saves
128 MiB. For scale, a measured 4x V100 config sizes KV at 8.26 KiB/token
(fp8_e4m3) / 14.4 KiB/token (fp16), i.e. 100 MiB is 7.1k-12.4k context tokens.
Tests
Five new regression tests in
tests/v1/attention/test_sm70_flash_v100_policy.py:three assert (via a weakref probe inside
torch.empty) that each growth hasreleased the previous buffer before allocating, one covers the capacity cap, one
covers the OOM retry.
Run inside the
ct252-e4m3-offload-c4image on a V100 host, CPU-only(
CUDA_VISIBLE_DEVICES=-1), patch applied to that build:The delta is exactly the five new tests (all red on baseline, all green after).
No other test changes state. The 29 remaining failures are pre-existing in that
environment (test file from
mainagainst that image's build, plus no visibleGPU) and identical in both runs.
Not covered
No end-to-end run exercises these paths. Both engine runs I made were a
Qwen4-exp QSA model, whose attention calls the C++ kernels directly
(
vllm/models/qwen4_exp/nvidia/ops/qsa.py) and never entersFlashAttnV100Impl._flash_v100_prefill_with_prefix; all eleven Flash-V100workspace caches stayed empty after 4k/60k/120k prefills under both
fp8_e4m3and fp16 KV. Exercising gather-dense needs a model on the standard
FLASH_ATTN_V100 impl with
head_dim == 256and fp16 KV (orfp8_e5m2for thebridge); I had no such deployment available. The change is allocation-ordering
only and does not alter kernel inputs or results.