[randomvariable]

Stacked on #51 → #50. Review bottom-up; base retargets as each lands.

Purpose

A CUDA graph captures the raw device pointer of whichever workspace buffer was current at capture time. Anything that moves a workspace after capture turns every later replay into a use-after-free: the caching allocator has already handed the freed block to other tensors, so kernels read scratch from — and write scratch over — live data. It surfaces as an illegal memory access or misaligned address, or a wedged GPU, minutes into serving rather than at capture.

Two workspaces could move:

1. SM100Workspace.ensure_size grew its buffer with self._workspace_buf.resize_(), which frees the old storage.
2. FlashInfer's cuDNN GEMM execute paths grow their shared 32 MiB cache workspace the same way. This is what kills the SM120 NVFP4 configuration in upstream #52540 — a 256-byte overflow past the default during capture was enough. Fixed in flashinfer-ai/flashinfer#4553, which is not in our pinned release, so vLLM needs a guard.

This fork's SM100Workspace matched upstream's pre-state exactly, resize_() included. Adapted from vllm-project/vllm#52553 (upstream issues vllm-project#52540, vllm-project#34948).

Approach

Retire, never resize. ensure_size allocates a new buffer and keeps the old one alive in _retired_bufs. Graphs captured against a retired buffer stay correct because its capture-time size was sufficient for the shapes baked into those graphs.

Pre-grow FlashInfer's shared buffer in kernel_warmup, before any capture, so it never has to move.

Both allocate with zeros, not empty: split-KV and cuDNN split-K plans keep semaphore and accumulator regions in the workspace and expect them zeroed. A fresh cudaMalloc happens to return zeroed pages, but a recycled dirty block from the caching allocator does not, and kernels then spin or corrupt.

Two deliberate deviations from upstream

• _grow_to extracted. Upstream inlines the grow in ensure_size, which can only be exercised by constructing attention metadata and calling into ops.sm100_cutlass_mla_get_workspace_size. Splitting the grow out makes the retirement invariant directly testable; ensure_size is otherwise unchanged.
• presize_flashinfer_gemm_workspaces takes torch.device | None and no-ops on None. Worker.device is optional here and mypy rejects upstream's unguarded worker.device.

Test Plan

.venv/bin/python -m pytest tests/v1/attention/test_gemm_workspace_stability.py -q
.venv/bin/python -m pytest tests/v1/attention/ -q \
  --ignore=tests/v1/attention/test_b12x_dots3.py \
  --ignore=tests/v1/attention/test_flashinfer_decode_qlen_guard.py
.venv/bin/pre-commit run --files <the five changed files>

CPU-runnable coverage for the FlashInfer half, with a stubbed module: no-op when FlashInfer is absent, zero_init=True requested, and a clean fallback for the older _get_cache_buf signature that does not accept it.

The workspace-retirement invariant — old buffer still alive at its original data_ptr() after a grow, new buffer zeroed — is GPU-gated, because SM100Workspace allocates on "cuda" in its constructor (kept as upstream has it).

Test Result

• tests/v1/attention/test_gemm_workspace_stability.py: 3 passed, 1 skipped (the skip is the GPU-gated invariant)
• tests/v1/attention/: 779 passed, 167 failed, 2978 skipped, 4 errors. Baseline with the three source files stashed: 776 passed, 167 failed, 2977 skipped, 4 errors — identical failure and error counts, +3 passes from the new CPU tests. No regressions.
• Two files cannot be collected in this venv at all (No module named 'vllm.third_party.deep_gemm', No module named 'flashinfer'), confirmed identical with the change stashed
• pre-commit: exit 0 (ruff, ruff-format, mypy 3.10, markdownlint)

Red phase verified: ImportError: cannot import name 'presize_flashinfer_gemm_workspaces' before the helper existed.

Not verified here

The crash itself needs SM120 NVFP4 serving with CUDA graphs; only the address-stability contract is covered locally. The bmm_fp8_workspace name list is upstream's and unverified against our pinned FlashInfer — if a lane uses a different cache name, that lane stays unguarded, which is no worse than today.

Also included

notes/features/b12x-w4a16-prepared-weights-blocked.md records why vllm-project/vllm#52604 cannot be incorporated yet: FlashInfer vllm-project#4560 (prepare_weights/run_prepared) is still open and unreleased, the PR is an explicitly temporary draft stack, and it reads a source_format config field this fork does not define.

AI assistance was used for this change.

[randomvariable]

Landed on homelabs-main as part of the full stack merge at b3ded080fa — every commit from this branch is now an ancestor of homelabs-main. GitHub declines to retarget the base ("There are no new commits between base branch 'homelabs-main' and head branch"), which is itself confirmation the content is in, so closing manually rather than leaving a stale open PR.