From c4b5280a550086ba00d96c38c78cbec8552d8b6e Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Fri, 25 Sep 2026 20:49:16 -0500 Subject: [PATCH 1/2] [Bugfix][Kimi-K3] Refresh DSpark context KV cache pointers after the KV cache is re-bound K3DSparkModel cached the per-layer KV cache base pointers (and the uniform-block-layout flag) on first use and never rebuilt them. Since #57632 captures the context K/V precompute in the draft CUDA graph, the first use now happens during the CUDA graph memory-profiling capture, which runs against a temporary KV cache that is freed afterwards. The real KV cache is then allocated at different addresses, but the draft kept the stale pointers, so every later context K/V insert wrote into freed memory. On 8x MI355X with Kimi-K3 TP8 + DSpark, this faults with "HSA_STATUS_ERROR_MEMORY_FAULT / Memory access fault ... Reason: Unknown" at max-num-seqs 8 and 16 (FULL and FULL_AND_PIECEWISE draft graphs), during warmup, right after graph capture, or on the first request. Key both cached values on the live kv_cache.data_ptr() tuple and rebuild them when it changes. The pointer array is updated in place, so graphs captured against it keep reading the live pointers. Rebuilds only happen on eager calls: the JIT warmup and the eager pass before each capture both run before any recording. Verified on hardware with the same image: on every rank the array is built once against the profiling cache and rebuilt once against the real cache (different base address), and the configuration that faulted 3/3 now serves cleanly. Signed-off-by: Oxana Korzh Co-Authored-By: Claude Opus 5.5 (1M context) --- vllm/models/kimi_k3/nvidia/dspark_mla.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 0baff884ccaa..0edff6ffc520 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -376,7 +376,9 @@ def _has_uniform_block_layout( self, cache_layers: list[MultiHeadLatentAttention], ) -> bool: - if not hasattr(self, "_layers_share_kv_block_layout"): + key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) + if getattr(self, "_kv_block_layout_key", None) != key: + self._kv_block_layout_key = key ref_cache = cache_layers[0].kv_cache self._layers_share_kv_block_layout = all( cl.kv_cache.size(1) == ref_cache.size(1) @@ -390,16 +392,18 @@ def _get_context_kv_cache_ptrs( self, cache_layers: list[MultiHeadLatentAttention], ) -> torch.Tensor: - # The per-layer KV cache base pointers are stable after allocation, so - # build the pointer array once and return it on every call. - if not hasattr(self, "_context_cache_ptrs"): - ref_cache = cache_layers[0].kv_cache + key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) + if getattr(self, "_context_cache_ptrs_key", None) != key: + assert not torch.cuda.is_current_stream_capturing() cache_ptrs = torch.tensor( - [cl.kv_cache.data_ptr() for cl in cache_layers], - dtype=torch.int64, - device=ref_cache.device, + key, dtype=torch.int64, device=cache_layers[0].kv_cache.device ) - self._context_cache_ptrs = cache_ptrs + current = getattr(self, "_context_cache_ptrs", None) + if current is not None and current.shape == cache_ptrs.shape: + current.copy_(cache_ptrs) + else: + self._context_cache_ptrs = cache_ptrs + self._context_cache_ptrs_key = key return self._context_cache_ptrs def forward( From d44cf803a4856762f5b1073677ead11fe152425d Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Sat, 26 Sep 2026 15:32:34 -0500 Subject: [PATCH 2/2] Invalidate the cached context KV pointers instead of updating them in place Per review: the profiling graphs are released in _teardown_profiling_state and the runtime graphs are captured after the real KV cache is bound, so nothing keeps a reference to the old pointer tensor. Drop the in-place copy_() and invalidate both cached values in _has_uniform_block_layout when the kv_cache.data_ptr() tuple changes; _get_context_kv_cache_ptrs is back to upstream. Signed-off-by: Oxana Korzh Co-Authored-By: Claude Opus 5.5 (1M context) --- vllm/models/kimi_k3/nvidia/dspark_mla.py | 25 ++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 0edff6ffc520..37207d7bc9cf 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -377,8 +377,11 @@ def _has_uniform_block_layout( cache_layers: list[MultiHeadLatentAttention], ) -> bool: key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) - if getattr(self, "_kv_block_layout_key", None) != key: - self._kv_block_layout_key = key + if getattr(self, "_kv_cache_ptrs_key", None) != key: + assert not torch.cuda.is_current_stream_capturing() + self._kv_cache_ptrs_key = key + if hasattr(self, "_context_cache_ptrs"): + del self._context_cache_ptrs ref_cache = cache_layers[0].kv_cache self._layers_share_kv_block_layout = all( cl.kv_cache.size(1) == ref_cache.size(1) @@ -392,18 +395,16 @@ def _get_context_kv_cache_ptrs( self, cache_layers: list[MultiHeadLatentAttention], ) -> torch.Tensor: - key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) - if getattr(self, "_context_cache_ptrs_key", None) != key: - assert not torch.cuda.is_current_stream_capturing() + # The per-layer KV cache base pointers are stable after allocation, so + # build the pointer array once and return it on every call. + if not hasattr(self, "_context_cache_ptrs"): + ref_cache = cache_layers[0].kv_cache cache_ptrs = torch.tensor( - key, dtype=torch.int64, device=cache_layers[0].kv_cache.device + [cl.kv_cache.data_ptr() for cl in cache_layers], + dtype=torch.int64, + device=ref_cache.device, ) - current = getattr(self, "_context_cache_ptrs", None) - if current is not None and current.shape == cache_ptrs.shape: - current.copy_(cache_ptrs) - else: - self._context_cache_ptrs = cache_ptrs - self._context_cache_ptrs_key = key + self._context_cache_ptrs = cache_ptrs return self._context_cache_ptrs def forward(