[Bugfix][Kimi-K3] Refresh DSpark context KV cache pointers after the KV cache is re-bound (#58814)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Oxana Korzh committed
Sep 28, 2026 at 10:05 UTC
94d14629247955c69d40bdf2a9d7e3bd45954452
1 file changed
+6
-1
vllm/models/kimi_k3/nvidia/dspark_mla.py
+6
-1
@@ -376,7 +376,12 @@ class K3DSparkModel(nn.Module):
376
self,
377
cache_layers: list[MultiHeadLatentAttention],
378
) -> bool:
379
- if not hasattr(self, "_layers_share_kv_block_layout"):
379
+ key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers)
380
+ if getattr(self, "_kv_cache_ptrs_key", None) != key:
381
+ assert not torch.cuda.is_current_stream_capturing()
382
+ self._kv_cache_ptrs_key = key
383
+ if hasattr(self, "_context_cache_ptrs"):
384
+ del self._context_cache_ptrs
385
ref_cache = cache_layers[0].kv_cache
386
self._layers_share_kv_block_layout = all(
387
cl.kv_cache.size(1) == ref_cache.size(1)