[Bug]: workspace.py rejects post-lock growth in deepseek_v4_attention._forward_prefill (DSV4 #40991) — patch attached #41700

Summary

vllm/v1/worker/workspace.py:_ensure_workspace_size locks the per-rank attention workspace at the post-profile size and refuses growth. DeepSeek-V4-Flash's _forward_prefill (added in #40991) then fails with an AssertionError on real prompts that need slightly more workspace than the dummy profile run sized for. The locked size is structural and not influenced by --max-num-batched-tokens, --max-num-seqs, or --gpu-memory-utilization. The currently published workaround is --enforce-eager, which costs ~4× decode throughput.

The smoking gun is in the source: deepseek_v4_attention.py:170-172 carries this comment:

# Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather
# workspace allocated at _forward_prefill (and the matching profile-time
# reservation in attention_impl's dummy-run branch).
PREFILL_CHUNK_SIZE = 4

The "matching profile-time reservation in attention_impl's dummy-run branch" implies a pre-allocation hook was always intended. It just isn't there: the dummy-run path (attention_impl, if not isinstance(attn_metadata, dict)) returns early without ever calling through to _forward_prefill, so warmup never sees prefill workspace requirements, and lock_workspace() fires at the post-decode-only size. The first real prefill request then crashes.

Repro

• vLLM: jasl/vllm@77bbc16 (current ds4-sm120 tip; same bug on 428e08e)
• Cherry-pick: neuralmagic/vllm@kylesayrs/deepseek-ct@f910a73a93 (PR [CompressedTensors] DeepSeek4 CT Quantization Support #41276 work)
• Patch: packed_modules_mapping injected on DeepseekV4ForCausalLM (see pasta-paul/dsv4-flash-w4a16-fp8/scripts/patch_v4_packed_mapping.py)
• transformers: 5.8.0.dev0 (HF main; PR [Bugfix] Skip torch-compile decoration for DiffusionGemma transformers fallback #45643 add-deepseek-v4 was merged 2026-05-02)
• compressed-tensors: 0.15.1.a20260428
• Hardware: 2× NVIDIA DGX Spark GB10 (SM 12.1a, 121 GiB UMA), TP=2 over QSFP RDMA
• Model: pastapaul/DeepSeek-V4-Flash-W4A16-FP8 (W4A16 INT4 GPTQ experts + FP8_BLOCK attention + BF16 shared experts)

vllm serve pastapaul/DeepSeek-V4-Flash-W4A16-FP8 \
  --served-model-name deepseek-v4-flash --trust-remote-code \
  --kv-cache-dtype fp8 --block-size 256 \
  --tokenizer-mode deepseek_v4 \
  --tool-call-parser deepseek_v4 --enable-auto-tool-choice \
  --reasoning-parser deepseek_v4 \
  --max-model-len 16384 \
  --max-num-seqs 4 --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.92 \
  -tp 2 --nnodes 2 \
  --master-addr <HEAD_IP> --master-port 29501 --node-rank 0
# rank 1 launches with --headless added

Then send a chat completion with a ~1.3K-token prompt and max_tokens=4000. Crashes during the first prefill:

AssertionError: Workspace is locked but allocation from
'deepseek_v4_attention.py:1457:_forward_prefill' requires 21.80 MB,
current size is 21.62 MB. Workspace growth is not allowed after locking.

The 21.62 MB locked size is identical across two builds 28 vLLM commits apart (only the file line number moved 1454→1457). Real-prompt sizes range 21.80 MB → 24.89 MB. Even when --max-num-batched-tokens was raised from default → 8192, the locked workspace did not grow accordingly. Lowering --max-num-seqs to 2 made the lock TIGHTER (smaller profile worst-case), not larger.

Working patch

Implementing what the comment describes — a _warmup_reserve_prefill_workspace() method on DeepseekV4MLAAttention that calls get_simultaneous() with worst-case shapes computed from max_model_len, max_num_batched_tokens, and config constants, called from the wrapper's dummy-run early-return:

# In attention_impl, dummy-run early-return:
if not isinstance(attn_metadata, dict):
    out.zero_()
    self.mla_attn._warmup_reserve_prefill_workspace()  # ← the hook
    return

Full patch (~30 lines, single-file, idempotent string-replacement script): pasta-paul/dsv4-flash-w4a16-fp8/scripts/patch_workspace_prereserve.py.

Validation results

Same en2zh_bus_001 1,304-token prompt that crashes without the patch:

Full harness run on :warmup (patched image), TP=2 on dual DGX Spark:

• chat-smoke quick: 4/4 PASS
• generation-matrix (18 prompts × 3 thinking modes): 50/54 (the 4 fails are a separate GENERATION_MAX_CASE_TOKENS ceiling issue unrelated to workspace)
• toolcall15: 41/45 (92%) — slight improvement over --enforce-eager baseline (89%)
• oracle_compare vs B200 TP=2: 5/5 ran, alignment numbers comparable
• gsm8k 8-shot: 95.37% (vs 92.87% on H200 reference rig)
• humaneval_instruct pass@1: 80.49% (vs 54.27% on H200 reference rig)
• 6+ hours continuous uptime, 0 workspace-lock errors

What an upstream fix could look like

The patch above works but is DSV4-specific. A more general fix in vllm/v1/worker/workspace.py would be one of:

1. Opt-in growth post-lock: keep the _locked assertion but expose _ensure_workspace_size(required_bytes, allow_growth=False). Layers that allocate variable-size scratch in non-graph-captured paths (like prefill) pass allow_growth=True. The allocator keeps the old tensor referenced (so captured graphs aren't invalidated) and allocates a new larger one alongside.
2. Documented warmup pre-allocation hook: add a method like WorkspaceManager.warmup_reserve(...) that layers can call from a known warmup phase before lock_workspace() runs. Layers that can't run their full forward in the dummy path (because metadata is required) get an explicit place to declare their max workspace.
3. Make attention_impl's dummy-run actually exercise prefill with synthetic but realistic attn_metadata. Heavier change but eliminates the entire class of "this layer's workspace was never sized" bugs.

Cross-references

• DSV4 main vLLM PR: [DSv4][Nvidia] SM12x DeepSeek V4 support #40991 (cc @jasl)
• compressed-tensors V4 attention PR: [CompressedTensors] DeepSeek4 CT Quantization Support #41276 (cc @kylesayrs)
• Related workspace allocation failure (different scenario, same workspace.py): [Bug]: Workspace allocation failure when combining Decode Context Parallelism (DCP) with EAGLE3 speculative decoding #40791
• Full Spark TP=2 validation report (build provenance, harness PASS rates, B200 oracle alignment, decode timings): https://github.com/pasta-paul/dsv4-flash-w4a16-fp8/blob/main/findings/spark_tp2_deployment.md
• The patch script: https://github.com/pasta-paul/dsv4-flash-w4a16-fp8/blob/main/scripts/patch_workspace_prereserve.py

Happy to refine the patch (the worst-case constants are heuristic — top_k=8192 upper bound, etc.) and submit a PR if maintainers prefer that path.