[XinyuJiangCMU]

Co-authored-by: @JessicaJiang-123

On gfx950, the DeepSeek-V4 compress-state pool allocated kv_score_buffer outside the torch_memory_saver region. As a result, pause() could not reclaim this memory during colocated RL training.

For a 16k, 4-node DeepSeek-V4-Flash run, this left approximately 52 GiB per GPU allocated throughout the training phase. This PR moves the HIP allocation into region(GPU_MEMORY_TYPE_KV_CACHE), matching the existing CUDA behavior.

The compute path is unchanged. This PR only changes how the buffer is allocated and registered with torch_memory_saver.

Root cause

_alloc_kv_score_buffer used a separate HIP path that allocated the buffer with a plain torch.empty(), outside any torch_memory_saver region:

if _is_hip:
    self.kv_score_buffer = KVAndScore(
        torch.empty((self._size, self.last_dim), dtype=dtype, device=device)
    )  # Not registered inside region(GPU_MEMORY_TYPE_KV_CACHE)
else:
    ... with region(...): with use_mem_pool(...): torch.empty(...)

torch_memory_saver can only unmap physical GPU memory for allocations registered inside a region. Because the HIP allocation was outside the region, the compress-state pool remained allocated after pause().

In addition, DeepSeekV4TokenToKVPool explicitly passed enable_memory_saver=False to _init_paged_compress_states on HIP, preventing the memory-saver allocation path from being used.

Change

Use the same allocation path for HIP and CUDA by allocating kv_score_buffer inside region(GPU_MEMORY_TYPE_KV_CACHE) through the custom memory pool and removing the HIP-specific enable_memory_saver=False behavior in DeepSeekV4TokenToKVPool. Both platforms now use the allocation path previously used only by CUDA.

Validation

Hardware: 8× MI355X GPUs (gfx950)
Model: DeepSeek-V4-Flash-FP8

Memory reclaimed

I constructed one ratio-128 compress-state pool with enable_memory_saver=True on gfx950 and measured physical GPU memory with mem_get_info before and after pause().

On the current main branch, DeepSeekV4TokenToKVPool forces enable_memory_saver=False on HIP. With this PR, the flag is enabled and the allocation is registered with torch_memory_saver.

Version	Pool allocation	Memory freed by pause()
main	2.02 GiB	0.000 GiB
This PR	1.87 GiB	1.873 GiB

With this PR, resume() correctly remaps the released memory.

Each ratio-128 pool uses approximately 1.87 GiB. Across roughly 20 ratio-128 attention pools, this accounts for the approximately 52 GiB per GPU that previously remained allocated during training.

In the full colocated run, the SGLang GPU-memory footprint during the training phase decreased from approximately 143 GiB to 87 GiB.

Correctness (GSM8K, 1319 questions, 8-shot, TP8)

Version	memory_saver=False	memory_saver=True
This PR	0.928	0.922
main	0.924	N/A

All three accuracy results are at least 0.91. The fix does not change model outputs. With memory_saver=True, the server starts successfully and serves requests correctly, validating the region and custom-memory-pool allocation path added by this PR.

Stability

A 4-node 16k colocate GRPO run on the fixed build: cold start and first resume with 0 SIGSEGV, and train/infer logprob drift ~0.04 matching the pre-fix baseline over 40+ steps.

Notes

• Scope: does not touch the compressor compute or the CUDA path.
• CUDA behavior is unchanged (both platforms now run the former else branch).

---

CI States

Latest PR Test (Base): ❌ Run #29718179112
Latest PR Test (Extra): ✅ Run #29723719042

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[1am9trash]

LGTM. Only effect hip code path.