[Frank-whw]

Motivation

Sparse-prefill scratch replacement allocates its new tensor while the backend still owns the old one. A larger workspace (or a different dtype) can therefore fail even when the replacement alone fits the available allocator budget.

This is a follow-up to #43921, which left workspace retention unchanged. The release-before-replacement change was originally proposed by @divyvasal in #43726. Refs #41076; this addresses the transient replacement peak, not the independently unbounded size of a batch of long histories.

Modifications

Release the workspace owner's old tensor reference before allocating a replacement. Scratch contents are overwritten by the existing caller. Requests that fit the current capacity still reuse the buffer, and views retained by a caller continue owning their original storage.

Add a single-GPU regression with fresh subprocesses so the allocator cap cannot affect other tests or a resident server.

Accuracy Tests

On an H200 with Torch 2.14.1+cu130, starting from upstream a3abea35609eeeff7e165486dd614963af211161:

python test/registered/unit/layers/test_dsv4_sparse_prefill_workspace.py -v

Isolated 128 MiB allocator limit	Unmodified main	This change
Grow 64 MiB BF16 scratch to 96 MiB	OOM	Pass
Replace 48 MiB byte scratch with 96 MiB BF16 scratch	OOM	Pass
Retained-view values and smaller-request reuse	Pass	Pass

Both replacement cases pass the below-100-MiB peak-allocation assertion after the fix. These are controlled allocator regressions using the actual workspace class, not a replay of the reported production OOM. No model weights or full serving workload are used.

Existing DSV4 compressed sparse-prefill attention (C4/C128) and workspace reuse tests also pass on H200. The initial numerical-test attempt failed to link libcudart with the pip CUDA path; using the complete /usr/local/cuda-13.0 toolkit passed without source changes.

Speed Tests and Profiling

No kernel, gather layout or attention math changes. No throughput claim; the measured improvement is successful scratch replacement under the same allocator cap.

Checklist

• Applicable pre-commit checks pass.
• CUDA regression fails on unmodified main and passes after the fix.
• Retained views and existing capacity reuse remain valid.
• Upstream GPU CI and maintainer review.

---

CI States

Latest PR Test (Base): ❌ Run #39123421376
Latest PR Test (Extra): ❌ Run #39123421591
Latest PR Test (AMD ROCm 10): ❌ Run #39123421440