_load_autosplit double-allocates dequant buffer (placeholder + forward simulation) causing OOM with large quantized cache sizes

[Nakanokensetsu]

Problem

When using quantized KV cache (cache_mode: "2,2" etc.) with large cache_size, _load_autosplit in model_ls.py temporarily allocates the dequantization buffer twice per attention layer:

1. get_kv_alloc_placeholder() — allocates FP16 K+V tensors of full cache shape (simulating dequant overhead)
2. module.forward(dummy_state, params) — runs actual attention forward, which calls cache.get_kv() internally and allocates the same-shaped FP16 K+V tensors again

At cache_size=409600 with Qwen3.6-27B (K2V2 quant, 8 KV heads, head_dim=128):

• Each tensor shape: (1600, 256, 8, 128) × FP16 = 800 MB
• Per attention layer peak: 800 MB × 4 (K placeholder + V placeholder + K forward + V forward) = 3.2 GB temporary

This causes GPU OOM during loading even though model+cache fits at inference time:

• Model + cache (balanced): ~9.1 GB/GPU — well within 12 GB
• Autosplit peak with double allocation: 9.1 + 3.2 = 12.3 GB → OOM on GPU with 12 GB physical limit

At cache_size=368640 the same double allocation yields 2.88 GB → peak = 9.0 + 2.88 = 11.88 GB → fits. So the failure is specific to larger cache sizes where the placeholder shape grows enough to push the combined peak over the physical limit.

Reproduction

model:
  backend: exllamav3
  model_name: Qwen3.6-27B-exl3-3.08bpw
  cache_mode: "2,2"
  cache_size: 409600   # 5 × max_seq_len
  max_seq_len: 81920
  gpu_split: [12, 12]  # 2× RTX 5070 (12 GB each)
draft_model:
  draft_mode: mtp
  draft_cache_mode: Q4

Expected: model loads (total model+cache ≈ 18.3 GB fits across 24 GB VRAM)
Actual: RuntimeError: Insufficient VRAM in split for model and cache

Root cause

model_ls.py lines 151–160 (v0.0.43):

# Placeholder — correctly simulates dequant overhead
qcache_overhead = []
for cm in module.all_cache_modules():
    for cl in cm.cache_layers:
        qcache_overhead.append(cl.get_kv_alloc_placeholder())  # allocates ~1.6 GB

# Forward — ALSO calls cache.get_kv() internally → allocates another ~1.6 GB
if self.caps.get("autosplit_load_fwd", True) and not autosplit_no_forward:
    dummy_state = module.prepare_for_device(dummy_state, params)
    dummy_state = module.forward(dummy_state, params)  # double-counts dequant peak!

The placeholder already accounts for the dequant overhead. The forward simulation then allocates the same buffers again while the placeholder is still live, doubling the counted peak.

Workaround

Pass autosplit_no_forward=True to model.load_gen(). This skips the forward simulation, leaving only the placeholder (1.6 GB). The autosplit then distributes correctly:

- GPU0: 10,111 MB | GPU1: 7,992 MB (cold)
- GPU0: 11,791 MB | GPU1: 9,624 MB (after first inference — PyTorch caches the dequant buffer for reuse)

Suggested fix

Option A — only allocate placeholder when skipping forward (they serve the same purpose, pick one):

qcache_overhead = []
if autosplit_no_forward:
    for cm in module.all_cache_modules():
        for cl in cm.cache_layers:
            qcache_overhead.append(cl.get_kv_alloc_placeholder())
# forward simulation already exercises the actual peak when it runs

Option B — free the placeholder before running forward, so peak is measured once:

qcache_overhead = []
for cm in module.all_cache_modules():
    for cl in cm.cache_layers:
        qcache_overhead.append(cl.get_kv_alloc_placeholder())

if self.caps.get("autosplit_load_fwd", True) and not autosplit_no_forward:
    qcache_overhead = None  # free placeholder first, then forward measures real peak
    dummy_state = module.prepare_for_device(dummy_state, params)
    dummy_state = module.forward(dummy_state, params)

Option B is probably more accurate since the forward simulation captures the real peak including any activation overhead beyond just K+V tensors.

---

[turboderp]

Thanks. That was in fact a redundant placeholder allocation. It's moot now, though, since cache quantization is fused into the new attention kernel, so those staging buffers won't be needed anymore.