From 410225b71969e80f1f232c6f06230bb9ed6697db Mon Sep 17 00:00:00 2001
From: sjtu_shenhai <33784657+sjtushenhai@users.noreply.github.com>
Date: Fri, 31 Oct 2025 21:30:37 +0800
Subject: [PATCH] [Bug fix] Fix severe memory waste issue with torch.empty
 pin_memory (#12266)

---
 python/sglang/srt/mem_cache/memory_pool_host.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py
index edfae2cfe..d3e9d97ea 100644
--- a/python/sglang/srt/mem_cache/memory_pool_host.py
+++ b/python/sglang/srt/mem_cache/memory_pool_host.py
@@ -238,12 +238,16 @@ class MHATokenToKVPoolHost(HostKVCache):
             raise ValueError(f"Unsupported layout: {self.layout}")
         self.token_stride_size = self.head_num * self.head_dim * self.dtype.itemsize
         self.layout_dim = self.token_stride_size * self.layer_num
-        return torch.empty(
+        buffer = torch.empty(
             dims,
             dtype=self.dtype,
             device=self.device,
-            pin_memory=self.pin_memory,
         )
+        if self.pin_memory:
+            torch.cuda.cudart().cudaHostRegister(
+                buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
+            )
+        return buffer
 
     @property
     def k_buffer(self):
@@ -551,13 +555,16 @@ class MLATokenToKVPoolHost(HostKVCache):
             self.kv_lora_rank + self.qk_rope_head_dim
         ) * self.dtype.itemsize
         self.layout_dim = self.token_stride_size * self.layer_num
-
-        return torch.empty(
+        buffer = torch.empty(
             dims,
             dtype=self.dtype,
             device=self.device,
-            pin_memory=self.pin_memory,
         )
+        if self.pin_memory:
+            torch.cuda.cudart().cudaHostRegister(
+                buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
+            )
+        return buffer
 
     def load_to_device_per_layer(
         self, device_pool, host_indices, device_indices, layer_id, io_backend