[Bugfix][Qwen4Exp] Release the profiling KV cache held by QSA key views (#58961)
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Gabriel Wu committed
Sep 28, 2026 at 16:11 UTC
a2be4d3cc7b978a6720dab32d28ae310775acb16
2 files changed
+34
-10
tests/models/qwen4_exp/test_qsa_reference.py
+22
-1
@@ -17,6 +17,7 @@ from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops
17
from vllm.models.qwen4_exp.nvidia.ops import qsa_indexer as qsa_indexer_ops
18
from vllm.platforms import current_platform
19
from vllm.triton_utils import HAS_TRITON
20
+from vllm.v1.worker.utils import clear_layer_kv_caches
21
22
requires_qsa_kernels = pytest.mark.skipif(
23
not current_platform.is_cuda() or not HAS_TRITON,
@@ -365,7 +366,9 @@ def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) ->
366
assert set(slots.tolist()).isdisjoint((committed % capacity).tolist())
367
368
368
-def _qsa_key_cache(block_size: int, compress_ratio: int) -> qsa_cache.QSAKeyStateCache:
369
+def _qsa_key_cache(
370
+ block_size: int, compress_ratio: int, **kwargs
371
+) -> qsa_cache.QSAKeyStateCache:
372
return qsa_cache.QSAKeyStateCache(
373
head_size=64,
374
dtype=torch.bfloat16,
@@ -375,6 +378,7 @@ def _qsa_key_cache(block_size: int, compress_ratio: int) -> qsa_cache.QSAKeyStat
378
compilation_config=SimpleNamespace(static_forward_context={})
379
),
380
compress_ratio=compress_ratio,
381
+ **kwargs,
382
)
383
384
@@ -402,6 +406,23 @@ def test_qsa_state_caches_adapt_the_unified_logical_layout() -> None:
406
assert compressed_cache.kv_cache.data_ptr() == compressed_view.data_ptr()
407
408
409
+def test_clearing_qsa_key_cache_releases_its_storage() -> None:
410
+ """Profiling teardown clears `kv_cache`; no derived view may keep it alive."""
411
+ raw_cache = _qsa_key_cache(
412
+ block_size=32, compress_ratio=4, cache_rope_positions=True
413
+ )
414
+ kv = torch.zeros(2, 1, 8, raw_cache.head_size, dtype=torch.bfloat16)
415
+ references_before_bind = torch._C._storage_Use_Count(kv.untyped_storage()._cdata)
416
+ raw_cache.bind_kv_cache(kv)
417
+ assert raw_cache.key_cache.shape[-1] == 64
418
+ assert raw_cache.rope_position_cache.dtype == torch.int64
419
+
420
+ clear_layer_kv_caches([raw_cache])
421
+
422
+ references = torch._C._storage_Use_Count(kv.untyped_storage()._cdata)
423
+ assert references == references_before_bind
424
+
425
+
426
@pytest.mark.parametrize(
427
("compress_ratio", "num_spec", "expected"),
428
[(4, 0, 4), (4, 1, 8), (4, 3, 8), (4, 4, 8), (4, 5, 12), (2, 3, 6)],
vllm/models/qwen4_exp/common/qsa_cache.py
+12
-9
@@ -825,15 +825,18 @@ class QSAKeyStateCache(_QSAStateCache):
825
)
826
super().__init__(head_size=storage_head_size, **kwargs)
827
828
- def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
829
- super().bind_kv_cache(kv_cache)
830
- qsa_cache = self.kv_cache
831
- self.key_cache = qsa_cache[..., : self.key_head_size]
832
- if self.cache_rope_positions:
833
- position_tail = qsa_cache[..., self.rope_position_offset :]
834
- self.rope_position_cache = position_tail.view(torch.int64)
835
- else:
836
- self.rope_position_cache = None
828
+ # Derived on access so `kv_cache` stays the only reference to the bound
829
+ # storage: clearing it (e.g. after CUDA graph memory profiling) must free
830
+ # the cache, or the freed block stays pinned in the allocator.
831
+ @property
832
+ def key_cache(self) -> torch.Tensor:
833
+ return self.kv_cache[..., : self.key_head_size]
834
+
835
+ @property
836
+ def rope_position_cache(self) -> torch.Tensor | None:
837
+ if not self.cache_rope_positions:
838
+ return None
839
+ return self.kv_cache[..., self.rope_position_offset :].view(torch.int64)
840
841
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
842
# Hold the open group's committed keys plus every row a speculative