Clarify post-capture KV reservation logs (#33598)

This commit is contained in:
cctry
2026-08-04 15:06:11 -07:00
committed by GitHub
parent 19d3f86895
commit c8822fd990
3 changed files with 5 additions and 4 deletions
@@ -152,7 +152,8 @@ class KvVmmArena:
# no_split so the caching allocator hands our bump pointers back verbatim. # no_split so the caching allocator hands our bump pointers back verbatim.
self.pool = torch.cuda.MemPool(self._allocator, no_split=True) self.pool = torch.cuda.MemPool(self._allocator, no_split=True)
logger.info( logger.info(
"KvVmmArena[%s] ready: device=%d reserved=%.1f GiB granularity=%d KiB", "KvVmmArena[%s] ready: device=%d reserved_va=%.1f GiB "
"granularity=%d KiB",
self._sfx, self._sfx,
self.device_id, self.device_id,
self.reserved / (1024**3), self.reserved / (1024**3),
+2 -2
View File
@@ -1667,7 +1667,7 @@ class KVCache(abc.ABC):
k_size_GB = k_size / GB k_size_GB = k_size / GB
v_size_GB = v_size / GB v_size_GB = v_size / GB
logger.info( logger.info(
f"{cache_name} is allocated. dtype: {self.dtype}, " f"{cache_name} {'VA upper bound' if self.post_capture_active else 'is allocated'}. dtype: {self.dtype}, "
f"#tokens: {num_tokens}, K size: {k_size_GB:.2f} GB, " f"#tokens: {num_tokens}, K size: {k_size_GB:.2f} GB, "
f"V size: {v_size_GB:.2f} GB" f"V size: {v_size_GB:.2f} GB"
) )
@@ -1675,7 +1675,7 @@ class KVCache(abc.ABC):
else: else:
kv_size_GB = kv_size_bytes / GB kv_size_GB = kv_size_bytes / GB
logger.info( logger.info(
f"{cache_name} is allocated. dtype: {self.dtype}, " f"{cache_name} {'VA upper bound' if self.post_capture_active else 'is allocated'}. dtype: {self.dtype}, "
f"#tokens: {num_tokens}, KV size: {kv_size_GB:.2f} GB" f"#tokens: {num_tokens}, KV size: {kv_size_GB:.2f} GB"
) )
self.mem_usage = kv_size_GB self.mem_usage = kv_size_GB
@@ -86,7 +86,7 @@ class SWAKVPool(BaseSWAKVPool):
k_size, v_size = self.get_kv_size_bytes() k_size, v_size = self.get_kv_size_bytes()
self.mem_usage = (k_size + v_size) / GB self.mem_usage = (k_size + v_size) / GB
logger.info( logger.info(
f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}" f"SWAKVPool {'VA upper bound' if self.post_capture_active else 'mem usage'}: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}"
) )
@property @property