[zohaibm-amd]

Overview

Follow-up to #1410. With kernel_config.rdna_hybrid_w4a16_prefill_dequant=enable, the per-layer budget check _fits_in_memory_budget reads device-wide free memory. That reading counts blocks the torch caching allocator holds after load-time temporaries (checkpoint loading, int4 unpack/dequant scratch) as used. Those blocks are reusable by the next allocation, and the allocator releases them itself if an allocation would otherwise fail. vLLM's KV-cache sizing already excludes them (memory_profiling measures after empty_cache()). As a result the check over-reserves, and a model that fits fails at startup with Not enough room in the gpu_memory_utilization budget ....

This change counts reserved - allocated (already captured by MemorySnapshot as torch_memory - torch_allocated) as free. It adds no flush or device sync, so neither load time nor runtime changes, and the result no longer depends on allocator state.

An earlier revision of this PR flushed the cache with torch.accelerator.empty_cache() when the check came up short. That also worked, but it costs a sync plus fresh device allocations, and it would also run if weights were re-processed at runtime (weight reload, sleep/wake).

Why keep the check at all

The dense copies are 4x their int4 weights and come out of the KV budget. Without the check, a config that does not fit fails later, either with the generic KV-sizing error or with an OOM mid-load, and neither names this option. The copy set must also be all-or-nothing, because the compile cache key does not record which layers hold a copy.

Measurements (gfx1151, 96 GiB device memory, CI wheel + this change)

Gemma-4-31B AWQ-4bit, MTP-3 draft, max_model_len=4096, gpu_memory_utilization=0.99:

	Before	After
Peak need at a layer's check	104.4 GiB, of which 15.4 GiB is allocator cache	89.4 GiB
Budget	95.0 GiB	95.0 GiB
Dense-cache checks failing	21–36 of 240	0 of 240
Engine startup	fails	OK

Peak need with allocator cache excluded, for the other prefill-dequant configs we run (all fit their budgets):

Model	Need
Qwen3.6-35B-A3B W4A16 (8192)	26.4 GiB
Qwen3.6-35B-A3B W4A16 MTP	29.3 GiB
Gemma-4-E4B AWQ	23.4 GiB
Gemma-4-E4B QAT	23.5 GiB

Rejected alternatives:

• PYTORCH_ALLOC_CONF=expandable_segments:True is worse: 55.7 GiB cached-free, 55 failures.
• garbage_collection_threshold has no effect, because it needs a per-process memory fraction that vLLM does not set.

Serving performance is unaffected; the check only runs while weights load. A/B against the pre-#1410 hard build, with the same torch/triton ROCm nightly, 10 MMStar prompts and max_num_seqs=1: mean TTFT 939.3 vs 942.7 ms, mean TPOT 35.30 vs 35.45 ms, MTP acceptance 2.82 in both.

Tests

tests/kernels/quantization/test_rdna_hybrid_w4a16.py -k prefill_dequant: 20 passed.

• The snapshot stubs now carry torch_memory/torch_allocated.
• New test: allocator-cached blocks count toward the budget (40 GiB free + 20 GiB cached fits where 40 GiB free alone does not), and empty_cache is never called.