Skip to content

[Bugfix][SM70] Release Flash-V100 prefill workspaces before growing them - #660

Open
Leonccaa wants to merge 1 commit into
1CatAI:mainfrom
Leonccaa:fix/qsa-workspace-growth-20260919
Open

Leonccaa wants to merge 1 commit into
1CatAI:mainfrom
Leonccaa:fix/qsa-workspace-growth-20260919

Conversation

@Leonccaa

Copy link
Copy Markdown
Contributor

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_workspace

This 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 WorkspaceManager already 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:

  1. A transient double-residency spike of the previous buffer's size.
  2. The growing allocation can fail on the memory its own predecessor is holding.
    The except torch.OutOfMemoryError path then returns None, the one-shot
    warning 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_workspace already pops the entry before
allocating; this makes the other three match it.

Change

  • Drop the cache entry (and the local reference) before allocating, so only one
    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_cache is what actually recovers the fragmented headroom. Only runs on
    the 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.
  • Cap gather-dense doubling at the block table width, i.e. the pages
    max_model_len can ever address. Beyond that the overshoot is pure loss
    against the KV cache.
  • Make the two bridge getters CPU-safe the same way their siblings already are
    (is_cuda guard on device/stream, tensor-aware capture check), so the growth
    path 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:

seq held before after
32k 32 MiB +32 MiB +32 MiB
64k 64 MiB +64 MiB +32 MiB
96k 128 MiB +128 MiB +64 MiB
128k 128 MiB +0 (hit) +0

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 has
released the previous buffer before allocating, one covers the capacity cap, one
covers the OOM retry.

Run inside the ct252-e4m3-offload-c4 image on a V100 host, CPU-only
(CUDA_VISIBLE_DEVICES=-1), patch applied to that build:

  • baseline: 123 passed / 34 failed
  • patched: 128 passed / 29 failed

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 main against that image's build, plus no visible
GPU) 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 enters
FlashAttnV100Impl._flash_v100_prefill_with_prefix; all eleven Flash-V100
workspace caches stayed empty after 4k/60k/120k prefills under both fp8_e4m3
and fp16 KV. Exercising gather-dense needs a model on the standard
FLASH_ATTN_V100 impl with head_dim == 256 and fp16 KV (or fp8_e5m2 for the
bridge); I had no such deployment available. The change is allocation-ordering
only and does not alter kernel inputs or results.

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants