[Bug]: TurboQuant _continuation_prefill workspace allocation fails at long context — v0.20.0 regression

Summary

On vLLM 0.20.0 with TurboQuant enabled, any prefill request whose cached KV grows beyond ~6-8K tokens triggers an assertion failure in _continuation_prefill's workspace allocation. The workspace is sized to 0.26 MB at startup but _continuation_prefill needs ~2 MB at long context, and the workspace is locked against growth after engine init.

The same workloads run cleanly on the v0.18-era TQ fork (vibhavagarwal5/PR #38479 base), so this is a v0.20-specific regression introduced when the workspace manager (PR #40941, merged Apr 27) was wired into the TQ continuation prefill path.

Reproduction

Setup: 8× RTX A4000 SM86, CUDA 13.0, driver 580.76.05, vLLM 0.20.0 from the GA tag.

vllm serve /path/to/Nemotron-3-Super-120B-AWQ-4bit \
  --tensor-parallel-size 8 --gpu-memory-utilization 0.92 \
  --max-num-seqs 4 --max-model-len 131072 \
  --trust-remote-code --enable-expert-parallel \
  --mamba-ssm-cache-dtype float16 \
  --kv-cache-dtype turboquant_3bit_nc

Send a long prompt:

import requests
prompt = 'In computational science, ' * 2000  # ~8K tokens after tokenization
r = requests.post("http://localhost:8001/v1/chat/completions", json={
    "model": "/path/to/Nemotron-3-Super-120B-AWQ-4bit",
    "messages": [{"role": "user", "content": prompt + " Reply OK."}],
    "max_tokens": 50,
})

Expected: response.

Actual:

AssertionError: Workspace is locked but allocation from
'turboquant_attn.py:757:_continuation_prefill' requires 2.00 MB,
current size is 0.26 MB. Workspace growth is not allowed after locking.

Engine dies; subsequent requests return HTTP 500 EngineCore encountered an issue.

Threshold sweep on the same hardware

Same model, same hardware, same TQ preset (tq-t3nc ↔ turboquant_3bit_nc). Only difference: vLLM revision.

Independent of TheTom's sparse-V PR

We hit the same crash with the sparse-V PR (#41422) applied AND with stock v0.20.0. The bug is in v0.20's _continuation_prefill workspace usage, not in any open PR.

Root cause analysis (best-effort)

vllm/v1/attention/backends/turboquant_attn.py:748-760:

# Allocate slightly over to align to block_size for the grid.
# Reuse cached buffers to avoid per-call allocation (~16MB at 8K).
alloc_len = math.ceil(cached_len / block_size) * block_size
buf_shape = (1, Hk, alloc_len, D)
# Use WorkspaceManager for dequant buffers.
# Shared across all layers — saves 60× memory at long context.
# Required for CUDA Graph capture (per-layer growth incompatible with CG).
k_buf, v_buf = current_workspace_manager().get_simultaneous(
    (buf_shape, torch.float16),
    (buf_shape, torch.float16),
)

The buffer shape scales with cached_len. The workspace manager (PR #40941) locks the workspace at engine init based on the high-water mark observed during the dummy run. The dummy run uses short chunked-prefill, so the locked-in workspace is sized for short cached lengths only. At inference time, the first request that exceeds the dummy run's cached_len triggers the lock-violation assertion.

Two interacting issues:

1. Profile undercount: the dummy run (is_profile=True, force_eager=True) doesn't exercise long-context continuation prefill. Whatever schedule the profile runs, it produces a cached_len smaller than what real workloads encounter when prefill chunking kicks in past max_num_batched_tokens.

2. Lock-vs-growth: the workspace manager is correct to lock for CUDA Graph compatibility, but a TQ-aware sizing heuristic is needed so the locked size accommodates the worst case the user actually configured (max_model_len).

This is the same conceptual failure mode jhsmith409 reported on PR #39931 for hybrid GDN models (chunk_fwd_o OOM at o = torch.empty_like(v)), where their proposed mitigation was a 1.25 GiB safety-margin overlay in gpu_worker.py. Different op, same upstream cause.

Suggested fixes

In rough order of invasiveness:

1. Profile-time long-context dummy run: have the engine init explicitly run a continuation prefill at max_model_len or a representative cached length so the workspace high-water mark covers production workloads. Adds ~5-10 seconds to startup. Symmetrical to the existing short-prefill profile pass.

2. TQ-aware static workspace sizing: compute the worst-case _continuation_prefill buffer from (num_kv_heads, max_model_len, head_dim) at engine init and reserve it ahead of locking, regardless of what the dummy run touches.

3. Allow late-growth with replay-aware lock: relax the workspace lock to permit growth that re-records affected CUDA Graphs. More complex, broader implications.

4. Per-call fallback: when the workspace can't satisfy the request, fall back to direct torch.empty allocation with a one-shot warning. Pessimizes long-context decode but keeps the engine alive instead of asserting.

We applied #1 locally as a stopgap (overrode the dummy run to include a worst-case continuation prefill) and confirmed it resolves the crash on Super-120B. Happy to share the diff.

Environment

• vLLM: 0.20.0 (vllm==0.20.0+precompiled, also reproduced on rebuilt-from-source 0.20.1.dev with the same line numbers)
• PyTorch: 2.11.0+cu130
• Driver/CUDA: 580.76.05 / 13.0
• Hardware: 8× RTX A4000 (SM86)
• Models that reproduce: Nemotron-3-Super-120B-AWQ-4bit, Nemotron-Cascade-2-30B-A3B (BF16). Likely any TQ-enabled model whose _continuation_prefill is invoked with cached_len larger than the engine's profile run.
• Models that do NOT reproduce: same models on PR #38479 fork (pre-workspace-manager).

Related

• PR #40941 (merged Apr 27): introduced the WorkspaceManager and wired it into TQ continuation prefill. Likely where the regression originated.
• PR #41422 (TheTom sparse-V): independent change; its long-context test plan is blocked by this bug on Ampere/Ada.
• PR #39931 #41185 #41123 (hybrid model TQ work): unrelated but shares the same family of "dummy run undercounts the real workload" failure mode.

— MidasMining, 8× RTX A4000 SM86 / vLLM 0.20.0