diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index aa3814030..815490dc9 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -702,7 +702,6 @@ class Envs: SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(False) # SWA radix cache - SGLANG_OPT_CACHE_SWA_TRANSLATION = EnvBool(True) # TODO(DSV4): @ispobock this has bug on main branch when retract SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False) SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 363051600..aa80dfefd 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -116,6 +116,9 @@ class DSV4AttnMetadata: swa_topk_lengths: torch.Tensor c4_sparse_topk: int + # SWA KV-store write target (out_cache_loc translated to SWA space), computed + # once per iteration in make_core_attn_metadata and read by the store path. + swa_out_cache_loc: Optional[torch.Tensor] = None c4_out_loc: Optional[torch.Tensor] = None c4_topk_lengths_raw: Optional[torch.Tensor] = None c4_topk_lengths_clamp1: Optional[torch.Tensor] = None @@ -172,6 +175,9 @@ class DSV4AttnMetadata: "c4_sparse_raw_indices", ], assign_fields=[ + # Recomputed by the recorded init_forward_metadata_in_graph op + # each forward; not copied across replays. + "swa_out_cache_loc", "c1_flashmla_metadata", "c4_flashmla_metadata", "c128_flashmla_metadata", @@ -218,6 +224,7 @@ class DSV4AttnMetadata: ] _CP_GLOBAL_FIELDS = [ "raw_out_loc", + "swa_out_cache_loc", "c4_out_loc", "c128_out_loc", ] @@ -699,6 +706,36 @@ class DeepseekV4AttnBackend( raw_metadata=self.forward_metadata, ) + # Compute the SWA KV-store write target once per forward and cache it on + # the metadata for every layer's store. This is recorded inside the cuda + # graph, so replay re-reads the live out_cache_loc buffer (spec-v2 and DP + # padding rebind out_cache_loc after out-graph metadata prep). flash_mla + # kernels require int32 indices. + metadata = self.forward_metadata + if ( + isinstance(metadata, DSV4Metadata) + and forward_batch.out_cache_loc is not None + ): + out_cache_loc = forward_batch.out_cache_loc + if ( + forward_batch.forward_mode.is_decode_or_idle() + and self.topk > 0 + and self.speculative_num_steps > 1 + ): + # Multi-step draft decode shares one out_cache_loc buffer across + # steps; mirror the eager init's per-step slice. + out_cache_loc = per_step_draft_out_cache_loc( + out_cache_loc, + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, + )[self.speculative_step_id] + metadata.core_attn_metadata.swa_out_cache_loc = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int32 + ) + ) + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, @@ -789,12 +826,22 @@ class DeepseekV4AttnBackend( ) elif bucket == _GraphBucket.DRAFT_EXTEND: num_tokens_per_bs = self.draft_extend_num_tokens_per_bs + if out_cache_loc is not None: + # Pad the real write locations to the captured token count so + # raw_out_loc reflects the actual replay out_cache_loc. + out_cache_loc = torch.nn.functional.pad( + out_cache_loc, + pad=(0, num_tokens_per_bs * bs - len(out_cache_loc)), + mode="constant", + value=0, + ) temp_metadata = self.init_forward_metadata_draft_extend( max_seq_len=chosen_max_seq_len, req_pool_indices=req_pool_indices, seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu.tolist(), num_tokens_per_bs=num_tokens_per_bs, + out_cache_loc=out_cache_loc, use_prefill_cuda_graph=True, ) else: @@ -934,21 +981,47 @@ class DeepseekV4AttnBackend( if current_raw is not None: self.forward_metadata = current_raw + def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: + """Resolve the SWA KV-store write target for the current forward. + + Fast path: the per-forward value cached by init_forward_metadata_in_graph + (recorded inside cuda graphs, so replay re-reads live buffers). Fallback: + translate at store time, matching the pre-cache behavior, for paths that + never run the in-graph init — eager idle (forward_idle skips attn init), + runners that only run the out-graph prep (e.g. + EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after + init (shape mismatch). Idle always falls back: its metadata is absent or + left over from a previous forward, and translating the zero-padded + out_cache_loc writes to the dummy slot. + """ + out_cache_loc = forward_batch.out_cache_loc + core = getattr(self.forward_metadata, "core_attn_metadata", None) + cached = core.swa_out_cache_loc if core is not None else None + if ( + cached is not None + and not forward_batch.forward_mode.is_idle() + and cached.shape[0] == out_cache_loc.shape[0] + ): + return cached + return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int32 + ) + def store_cache( self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch ) -> None: - raw_loc = forward_batch.out_cache_loc + swa_loc = self.get_swa_out_cache_loc(forward_batch) if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get(): self.token_to_kv_pool.set_swa_key_buffer_radix_fused( layer_id=layer_id, - raw_loc=raw_loc, + swa_loc=swa_loc, cache_k=swa_k, ) else: swa_k_pack = quant_to_nope_fp8_rope_bf16_pack_triton(swa_k) self.token_to_kv_pool.set_swa_key_buffer_radix( layer_id=layer_id, - raw_loc=raw_loc, + swa_loc=swa_loc, cache_nope_fp8_rope_bf16_pack=swa_k_pack, ) @@ -1322,7 +1395,8 @@ class DeepseekV4AttnBackend( assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW) raw_indices.masked_fill_(invalid_offset_mask, -1) swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices) - return swa_indices + # flash_mla attention requires int32 page indices. + return swa_indices.to(torch.int32) class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend): diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index c42c8d678..3571813af 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -106,6 +106,9 @@ class DSV4AttnMetadata: swa_topk_lengths: torch.Tensor c4_sparse_topk: int + # SWA KV-store write target (out_cache_loc translated to SWA space), computed + # once per iteration in make_core_attn_metadata and read by the store path. + swa_out_cache_loc: Optional[torch.Tensor] = None c4_out_loc: Optional[torch.Tensor] = None c4_topk_lengths_raw: Optional[torch.Tensor] = None c4_topk_lengths_clamp1: Optional[torch.Tensor] = None @@ -160,6 +163,9 @@ class DSV4AttnMetadata: "c4_sparse_page_indices", ], assign_fields=[ + # Recomputed by the recorded init_forward_metadata_in_graph op + # each forward; not copied across replays. + "swa_out_cache_loc", "c1_flashmla_metadata", "c4_flashmla_metadata", "c128_flashmla_metadata", @@ -206,6 +212,7 @@ class DSV4AttnMetadata: ] _CP_GLOBAL_FIELDS = [ "raw_out_loc", + "swa_out_cache_loc", "c4_out_loc", "c128_out_loc", ] @@ -672,6 +679,36 @@ class DeepseekV4HipRadixBackend( raw_metadata=self.forward_metadata, ) + # Compute the SWA KV-store write target once per forward and cache it on + # the metadata for every layer's store. This is recorded inside the cuda + # graph, so replay re-reads the live out_cache_loc buffer (spec-v2 and DP + # padding rebind out_cache_loc after out-graph metadata prep). flash_mla + # kernels require int32 indices. + metadata = self.forward_metadata + if ( + isinstance(metadata, DSV4Metadata) + and forward_batch.out_cache_loc is not None + ): + out_cache_loc = forward_batch.out_cache_loc + if ( + forward_batch.forward_mode.is_decode_or_idle() + and self.topk > 0 + and self.speculative_num_steps > 1 + ): + # Multi-step draft decode shares one out_cache_loc buffer across + # steps; mirror the eager init's per-step slice. + out_cache_loc = per_step_draft_out_cache_loc( + out_cache_loc, + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, + )[self.speculative_step_id] + metadata.core_attn_metadata.swa_out_cache_loc = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int32 + ) + ) + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, @@ -760,12 +797,22 @@ class DeepseekV4HipRadixBackend( ) elif bucket == _GraphBucket.DRAFT_EXTEND: num_tokens_per_bs = self.draft_extend_num_tokens_per_bs + if out_cache_loc is not None: + # Pad the real write locations to the captured token count so + # raw_out_loc reflects the actual replay out_cache_loc. + out_cache_loc = torch.nn.functional.pad( + out_cache_loc, + pad=(0, num_tokens_per_bs * bs - len(out_cache_loc)), + mode="constant", + value=0, + ) temp_metadata = self.init_forward_metadata_draft_extend( max_seq_len=chosen_max_seq_len, req_pool_indices=req_pool_indices, seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu.tolist(), num_tokens_per_bs=num_tokens_per_bs, + out_cache_loc=out_cache_loc, use_prefill_cuda_graph=True, ) else: @@ -905,21 +952,47 @@ class DeepseekV4HipRadixBackend( if current_raw is not None: self.forward_metadata = current_raw + def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: + """Resolve the SWA KV-store write target for the current forward. + + Fast path: the per-forward value cached by init_forward_metadata_in_graph + (recorded inside cuda graphs, so replay re-reads live buffers). Fallback: + translate at store time, matching the pre-cache behavior, for paths that + never run the in-graph init — eager idle (forward_idle skips attn init), + runners that only run the out-graph prep (e.g. + EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after + init (shape mismatch). Idle always falls back: its metadata is absent or + left over from a previous forward, and translating the zero-padded + out_cache_loc writes to the dummy slot. + """ + out_cache_loc = forward_batch.out_cache_loc + core = getattr(self.forward_metadata, "core_attn_metadata", None) + cached = core.swa_out_cache_loc if core is not None else None + if ( + cached is not None + and not forward_batch.forward_mode.is_idle() + and cached.shape[0] == out_cache_loc.shape[0] + ): + return cached + return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int32 + ) + def store_cache( self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch ) -> None: - raw_loc = forward_batch.out_cache_loc + swa_loc = self.get_swa_out_cache_loc(forward_batch) if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get(): self.token_to_kv_pool.set_swa_key_buffer_radix_fused( layer_id=layer_id, - raw_loc=raw_loc, + swa_loc=swa_loc, cache_k=swa_k, ) else: swa_k_pack = quant_to_nope_fp8_rope_bf16_pack_triton(swa_k) self.token_to_kv_pool.set_swa_key_buffer_radix( layer_id=layer_id, - raw_loc=raw_loc, + swa_loc=swa_loc, cache_nope_fp8_rope_bf16_pack=swa_k_pack, ) @@ -1165,7 +1238,8 @@ class DeepseekV4HipRadixBackend( assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW) raw_indices.masked_fill_(invalid_offset_mask, -1) swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices) - return swa_indices + # flash_mla attention requires int32 page indices. + return swa_indices.to(torch.int32) class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend): diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index e3b8eebd5..254afaec8 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -1486,9 +1486,10 @@ class DeepseekSparseAttnBackend( # todo hisparse: to cover more backends if self.hisparse_coordinator is not None: + # flash_mla_sparse_fwd / tilelang require int32 page indices. page_table_1 = self.token_to_kv_pool.translate_loc_to_hisparse_device( page_table_1 - ) + ).to(torch.int32) if dsa_impl == "tilelang": if q_rope is not None: diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index b48bd985c..fdbe68cc7 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -509,7 +509,10 @@ class CompressorBackendMixin: if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"): # The v2 compressor writes directly into the raw C4 KV tensor. # HiSparse C4 therefore needs the physical C4 location here. - out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc) + # The compress kernel requires an int32 write location. + out_loc = compress_kv_pool.translate_loc_to_hisparse_device( + out_loc + ).to(torch.int32) self._forward_compress_all_in_one( kv_score_buffer=state_pool.kv_score_buffer.kv_score, kv_score_input=kv_score_input, diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 5f4786520..e14bdc648 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -600,10 +600,11 @@ class C4IndexerBackendMixin: ) ) else: + # flash_mla C4 attention requires int32 page indices. core_metadata.c4_sparse_page_indices = ( token_to_kv_pool.c4_kv_pool.translate_loc_to_hisparse_device( core_metadata.c4_sparse_page_indices - ) + ).to(torch.int32) ) if capture_enabled: diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 351dc90ee..2a9b1e854 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -693,10 +693,11 @@ class FlashAttentionBackend(AttentionBackend): ] if self.use_sliding_window_kv_pool: + # FA3 requires an int32 page_table. metadata.swa_page_table = ( self.token_to_kv_pool.translate_loc_from_full_to_swa( metadata.page_table - ) + ).to(torch.int32) ) # Convert the page table to a strided format which is needed by FA3 API @@ -886,7 +887,7 @@ class FlashAttentionBackend(AttentionBackend): else: page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa( metadata.page_table - ) + ).to(torch.int32) cu_seqlens_q = metadata.cu_seqlens_q cache_seqlens = metadata.cache_seqlens_int32 max_seqlen_q = metadata.max_seq_len_q @@ -1365,7 +1366,7 @@ class FlashAttentionBackend(AttentionBackend): page_table = ( self.token_to_kv_pool.translate_loc_from_full_to_swa( metadata.page_table - ) + ).to(torch.int32) ) cache_seqlens = metadata.cache_seqlens_int32 max_seqlen_q = metadata.max_seq_len_q @@ -2424,7 +2425,7 @@ class FlashAttentionBackend(AttentionBackend): if self.use_sliding_window_kv_pool: page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa( metadata.page_table - ) + ).to(torch.int32) else: page_table = metadata.page_table if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None: @@ -2551,7 +2552,7 @@ class FlashAttentionBackend(AttentionBackend): if self.use_sliding_window_kv_pool: sliced_page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa( metadata.page_table[:bs, :max_seq_len] - ) + ).to(torch.int32) else: sliced_page_table = metadata.page_table[:bs, :max_seq_len] @@ -2628,10 +2629,10 @@ class FlashAttentionBackend(AttentionBackend): if self.use_sliding_window_kv_pool: page_table_a = self.token_to_kv_pool.translate_loc_from_full_to_swa( page_table_a - ) + ).to(torch.int32) page_table_b = self.token_to_kv_pool.translate_loc_from_full_to_swa( page_table_b - ) + ).to(torch.int32) prepare_swa_spec_page_table_triton( page_table, diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 4af1724aa..7e2f694d9 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1519,12 +1519,9 @@ def update_sliding_window_buffer( ) if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"): kv_last_index = window_kv_indptr[-1] - # Flush before+after: window_kv_indices is a different tensor than out_cache_loc. - token_to_kv_pool.invalidate_loc_cache() window_kv_indices[:kv_last_index] = ( token_to_kv_pool.translate_loc_from_full_to_swa( window_kv_indices[:kv_last_index] ) ) - token_to_kv_pool.invalidate_loc_cache() return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 78ca325f8..de3ca31ad 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -164,9 +164,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): if self._swa_kv_pool is None: return None shape = token_indices.shape - return self._swa_kv_pool.translate_loc_from_full_to_swa( - token_indices.reshape(-1) - ).reshape(shape) + # trtllm-gen SWA attention kernels require int32 page indices. + return ( + self._swa_kv_pool.translate_loc_from_full_to_swa(token_indices.reshape(-1)) + .reshape(shape) + .to(torch.int32) + ) def _alloc_swa_page_table( self, max_bs: int, max_num_pages: int diff --git a/python/sglang/srt/mem_cache/base_swa_memory_pool.py b/python/sglang/srt/mem_cache/base_swa_memory_pool.py index 717ba05d4..06ff43984 100644 --- a/python/sglang/srt/mem_cache/base_swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/base_swa_memory_pool.py @@ -16,9 +16,6 @@ class BaseSWAKVPool(KVCache): swa_kv_pool: KVCache - def invalidate_loc_cache(self) -> None: - pass - @abc.abstractmethod def register_mapping(self, full_to_swa_index_mapping: torch.Tensor) -> None: raise NotImplementedError() diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index e36b1f3e8..77a5a1387 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -513,15 +513,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): else: self._init_paged_compress_states(enable_memory_saver) - self._should_cache_swa = envs.SGLANG_OPT_CACHE_SWA_TRANSLATION.get() - self.cached_loc = None - def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping - self.cached_loc = None # mapping replaced; discard any cached translation - - def invalidate_loc_cache(self) -> None: - self.cached_loc = None def get_ring_size(self, compress_ratio: int) -> int: server_args = get_global_server_args() @@ -530,15 +523,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): assert self.full_to_swa_index_mapping is not None - - return self.full_to_swa_index_mapping[kv_indices].to(torch.int32) - - def get_cached_swa_loc(self, raw_loc: torch.Tensor, layer_id: int) -> torch.Tensor: - if self._should_cache_swa: - if layer_id == self.start_layer or self.cached_loc is None: - self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc) - return self.cached_loc - return self.translate_loc_from_full_to_swa(raw_loc) + return self.full_to_swa_index_mapping[kv_indices] def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]: data_ptrs: List[int] = [] @@ -768,10 +753,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def set_swa_key_buffer_radix( self, layer_id: int, - raw_loc: torch.Tensor, + swa_loc: torch.Tensor, cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack, ) -> None: - swa_loc = self.translate_loc_from_full_to_swa(raw_loc) self.swa_kv_pool.set_key_buffer( self._swa_local_layer_id(layer_id), swa_loc, cache_nope_fp8_rope_bf16_pack ) @@ -783,10 +767,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def set_swa_key_buffer_radix_fused( self, layer_id: int, - raw_loc: torch.Tensor, + swa_loc: torch.Tensor, cache_k: torch.Tensor, ) -> None: - swa_loc = self.get_cached_swa_loc(raw_loc, layer_id) return self.swa_kv_pool.set_key_buffer_fused( self._swa_local_layer_id(layer_id), swa_loc, cache_k ) @@ -794,14 +777,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def set_swa_key_buffer_radix_fused_norm_rope( self, layer_id: int, - raw_loc: torch.Tensor, + swa_loc: torch.Tensor, kv: torch.Tensor, kv_weight: torch.Tensor, eps: float, freqs_cis: torch.Tensor, positions: torch.Tensor, ) -> None: - swa_loc = self.get_cached_swa_loc(raw_loc, layer_id) fused_k_norm_rope_flashmla( kv=kv, kv_weight=kv_weight, diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index b08fe0db1..d80d73697 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/python/sglang/srt/mem_cache/hisparse_memory_pool.py @@ -77,9 +77,7 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool): ) def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor): - return self.full_to_hisparse_device_index_mapping[compressed_indices].to( - torch.int32 - ) + return self.full_to_hisparse_device_index_mapping[compressed_indices] def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor): return self.full_to_hisparse_device_index_mapping[compressed_indices] diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 21c9f692b..550bb53ac 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -92,8 +92,6 @@ class SWAKVPool(BaseSWAKVPool): for swa_layer_id, global_layer_id in enumerate(swa_attention_layer_ids): self.layers_mapping[global_layer_id] = (swa_layer_id, True) self.full_to_swa_index_mapping: Optional[torch.Tensor] = None - self._cached_swa_loc: Optional[torch.Tensor] = None - self._cached_loc_key: Optional[tuple] = None k_size, v_size = self.get_kv_size_bytes() self.mem_usage = (k_size + v_size) / GB @@ -103,11 +101,6 @@ class SWAKVPool(BaseSWAKVPool): def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): self.full_to_swa_index_mapping = full_to_swa_index_mapping - self.invalidate_loc_cache() - - def invalidate_loc_cache(self) -> None: - self._cached_swa_loc = None - self._cached_loc_key = None def register_layer_transfer_counter(self, layer_transfer_counter): # Wait happens at this wrapper. Inner pools must not wait again. @@ -167,21 +160,8 @@ class SWAKVPool(BaseSWAKVPool): def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor: assert self.full_to_swa_index_mapping is not None - # data_ptr() (not untyped_storage().data_ptr()) encodes the offset, so - # views at different positions within the same storage get distinct keys. # -1 in kv_indices maps to -1 via the sentinel appended to the mapping. - key = (kv_indices.data_ptr(), kv_indices.numel()) - if key != self._cached_loc_key: - if self._cached_loc_key is not None: - logger.debug( - "translate_loc_from_full_to_swa: loc tensor changed mid-forward " - "without invalidate_loc_cache() — possible missing call site" - ) - self._cached_swa_loc = self.full_to_swa_index_mapping[kv_indices].to( - torch.int32 - ) - self._cached_loc_key = key - return self._cached_swa_loc + return self.full_to_swa_index_mapping[kv_indices] def set_kv_buffer( self, @@ -425,7 +405,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return self._kvcache.translate_loc_from_full_to_swa(kv_indices) def alloc(self, need_size: int): - self._kvcache.invalidate_loc_cache() assert self.page_size == 1 if need_size > self.full_attn_allocator.available_size(): return None @@ -454,7 +433,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): last_loc: torch.Tensor, # last_loc for full layers extend_num_tokens: int, ): - self._kvcache.invalidate_loc_cache() assert self.page_size > 1 num_new_pages = get_num_new_pages( @@ -507,7 +485,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): extend_num_tokens: int, swa_tail_len: int, ): - self._kvcache.invalidate_loc_cache() """Allocate full KV for the whole extend and SWA KV only for the tail. This is used by disaggregated decode preallocation: decode receives full @@ -571,7 +548,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): seq_lens_cpu: torch.Tensor, last_loc: torch.Tensor, # last_loc for full layers ): - self._kvcache.invalidate_loc_cache() assert self.page_size > 1 swa_last_loc = self.translate_loc_from_full_to_swa(last_loc) @@ -619,7 +595,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if full_indices.numel() == 0: return assert full_indices.numel() == swa_indices.numel() - self._kvcache.invalidate_loc_cache() if _is_npu: self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = ( swa_indices.to(torch.int64) @@ -628,7 +603,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_swa_index_mapping[full_indices] = swa_indices def free_swa(self, free_index: torch.Tensor): - self._kvcache.invalidate_loc_cache() swa_indices = self.full_to_swa_index_mapping[free_index] swa_indices = swa_indices[swa_indices > 0] self.swa_attn_allocator.free(swa_indices) @@ -646,7 +620,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.swa_attn_allocator.restore_state(state[1]) def clear(self): - self._kvcache.invalidate_loc_cache() self.swa_attn_allocator.clear() self.full_attn_allocator.clear() # Note: the last item is -1, we don't clear it, see the comment in __init__ diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py b/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py index 6dada9c31..15681ea13 100644 --- a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py @@ -414,9 +414,6 @@ class BreakableCudaGraphRunner: self.model_runner.attn_backend.init_forward_metadata(forward_batch) def run_once(): - # Invalidate SWA loc cache — same fix as in cuda_graph_runner.run_once. - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() return self._run_forward(forward_batch, num_tokens) with forward_context( diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 300cd9a84..cf40e18a2 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -1047,12 +1047,6 @@ class CudaGraphRunner: attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) def run_once(): - # Without this, warmup-1 caches the translation; the capture - # run hits the cache, skips the gather, and replay reuses - # stale SWA locations. - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() - # Must run inside the capture block: warmup mutations here are # undone by on_after_cuda_graph_warmup so capture starts clean. attn_backend.init_forward_metadata_in_graph(forward_batch) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 12c83090b..ce056571c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -3362,9 +3362,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.hisparse_coordinator.wait_for_pending_backup() self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size) - if self.is_hybrid_swa: - self.token_to_kv_pool.invalidate_loc_cache() - # Replay cuda graph if applicable if can_run_graph: ret = self.graph_runner.replay( diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index b6ed9ac07..e455efaf2 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -623,10 +623,6 @@ class PiecewiseCudaGraphRunner: # Run and capture def run_once(): - # Invalidate SWA loc cache — same fix as in cuda_graph_runner.run_once. - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() - # Clean intermediate result cache for DP attention forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = ( None diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 29174ffe6..c5fc558c9 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -471,6 +471,7 @@ class MQALayer(nn.Module): x: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, + attn_backend, qkv_a: Optional[torch.Tensor] = None, ) -> None: """Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache. @@ -487,7 +488,7 @@ class MQALayer(nn.Module): assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) token_to_kv_pool.set_swa_key_buffer_radix_fused_norm_rope( layer_id=self.layer_id, - raw_loc=forward_batch.out_cache_loc, + swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch), kv=kv, kv_weight=self.kv_norm.weight.data, eps=self.eps, @@ -562,7 +563,9 @@ class MQALayer(nn.Module): if qkv_a_ready is not None: stream_kv.wait_event(qkv_a_ready) # Fused norm + rope + cache write -- no bf16 KV intermediate. - self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a) + self._compute_kv_to_cache( + x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a + ) del qkv_a @@ -646,9 +649,7 @@ class MQALayer(nn.Module): ) token_to_kv_pool = get_token_to_kv_pool() - swa_loc = token_to_kv_pool.get_cached_swa_loc( - forward_batch.out_cache_loc, self.layer_id - ) + swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch) swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id] swa_page_size = token_to_kv_pool.swa_kv_pool.page_size @@ -672,7 +673,9 @@ class MQALayer(nn.Module): else: q_lora = self.q_norm(q_lora) q = self._compute_q_b(q_lora, positions, q_out) - self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a) + self._compute_kv_to_cache( + x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a + ) del qkv_a @@ -736,9 +739,7 @@ class MQALayer(nn.Module): ) token_to_kv_pool = get_token_to_kv_pool() - swa_loc = token_to_kv_pool.get_cached_swa_loc( - forward_batch.out_cache_loc, self.layer_id - ) + swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch) swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id] swa_page_size = token_to_kv_pool.swa_kv_pool.page_size @@ -790,7 +791,7 @@ class MQALayer(nn.Module): ) else: self._compute_kv_to_cache( - x_linear, positions, forward_batch, qkv_a=qkv_a + x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a ) kv = None diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index e8aa57741..fb4c6173e 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -352,9 +352,6 @@ class EAGLEDraftCudaGraphRunner: ) def run_once(): - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() - self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch) forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 4c7ee8cd7..422c2b26e 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -376,9 +376,6 @@ class EAGLEDraftExtendCudaGraphRunner: ) def run_once(): - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() - # Clean intermediate result cache for DP attention forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None set_dp_buffer_len( diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index da44c6342..a3ca1ea54 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -270,9 +270,6 @@ class FrozenKVMTPCudaGraphRunner: ) def run_once(): - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() - forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None set_dp_buffer_len( global_dp_buffer_len, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index a4ffe2c4a..ec465e3b3 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -401,9 +401,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner: attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step] def run_once(): - if self.model_runner.is_hybrid_swa: - self.model_runner.token_to_kv_pool.invalidate_loc_cache() - # Clean intermediate result cache for DP attention forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None set_dp_buffer_len( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 830756db1..b47c0e0f5 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -518,9 +518,12 @@ class ProjectedDSV4Attention(nn.Module): # `[num_tokens, 1, hidden_dim]`. k_flat = k.reshape(k.shape[0], -1).to(torch.bfloat16) pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_flat) - attn_backend.token_to_kv_pool.set_swa_key_buffer_radix( + pool = attn_backend.token_to_kv_pool + pool.set_swa_key_buffer_radix( layer_id=self.attn.layer_id, - raw_loc=forward_batch.out_cache_loc.to(torch.int64), + swa_loc=pool.translate_loc_from_full_to_swa( + forward_batch.out_cache_loc.to(torch.int64) + ), cache_nope_fp8_rope_bf16_pack=pack, ) out = attn_backend.forward( @@ -546,7 +549,9 @@ def _write_swa_cache( pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_bf16.to(torch.bfloat16)) runner.token_to_kv_pool.set_swa_key_buffer_radix( layer_id=layer_id, - raw_loc=loc.to(torch.int64), + swa_loc=runner.token_to_kv_pool.translate_loc_from_full_to_swa( + loc.to(torch.int64) + ), cache_nope_fp8_rope_bf16_pack=pack, ) diff --git a/test/manual/core/test_dsv4_cached_loc_invalidation.py b/test/manual/core/test_dsv4_cached_loc_invalidation.py deleted file mode 100644 index ff7314fc4..000000000 --- a/test/manual/core/test_dsv4_cached_loc_invalidation.py +++ /dev/null @@ -1,180 +0,0 @@ -"""Regression test for PR #25889: DeepSeekV4TokenToKVPool.register_mapping() -must clear cached_loc. - -Bug scenario (pre-fix): - During a forward pass, the first SWA layer (layer_id == start_layer) computes - and caches `cached_loc` via translate_loc_from_full_to_swa(). If HiCache then - loads back SWA KV from host, it calls register_mapping() to install the new - full->swa index mapping. Before the fix, register_mapping() only stored the - new tensor but did NOT clear cached_loc. Subsequent SWA layers (layer_id > - start_layer) saw `cached_loc is not None` and returned the stale translation, - silently writing KV to wrong SWA pool slots. - -Fix (PR #25889): add `self.cached_loc = None` in register_mapping(). - -Test structure: - - test_stale_without_fix: shows the stale-return bug using a replica - of the pre-fix logic - - test_correct_with_fix: verifies the fixed logic returns fresh values - - test_actual_pool_register_mapping: exercises the real production method - directly (bypassing the full __init__) - -Run with: - python -m pytest test/manual/core/test_dsv4_cached_loc_invalidation.py -v -""" - -import unittest - -import torch - -from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool -from sglang.srt.utils import get_device -from sglang.test.test_utils import CustomTestCase - -# --------------------------------------------------------------------------- -# Minimal stub that replicates DeepSeekV4TokenToKVPool's caching pattern. -# Used to demonstrate both sides of the bug without constructing the full pool. -# --------------------------------------------------------------------------- - - -class _DSV4CacheStub: - """Stripped-down replica of the caching logic in DeepSeekV4TokenToKVPool.""" - - start_layer = 3 - - def __init__(self, device): - self.device = device - self.cached_loc = None - self.full_to_swa_index_mapping = None - - # --- pre-fix version --- - def register_mapping_buggy(self, mapping: torch.Tensor) -> None: - self.full_to_swa_index_mapping = mapping - # BUG: cached_loc not cleared → stale on next mid-forward call - - # --- post-fix version (PR #25889) --- - def register_mapping_fixed(self, mapping: torch.Tensor) -> None: - self.full_to_swa_index_mapping = mapping - self.cached_loc = None # THE FIX - - def _translate(self, raw_loc: torch.Tensor) -> torch.Tensor: - return self.full_to_swa_index_mapping[raw_loc] - - def get_swa_loc(self, layer_id: int, raw_loc: torch.Tensor) -> torch.Tensor: - """Exact replica of set_swa_key_buffer_radix_fused caching branch.""" - if layer_id == self.start_layer or self.cached_loc is None: - self.cached_loc = self._translate(raw_loc) - return self.cached_loc - - -def _make_mapping(indices, values, size=32, device="cpu"): - m = torch.zeros(size, dtype=torch.int64, device=device) - m[indices] = torch.tensor(values, dtype=torch.int64, device=device) - return m - - -class TestDSV4CachedLocBugAndFix(CustomTestCase): - """Shows the pre-fix bug and verifies the post-fix behaviour.""" - - def setUp(self): - self.device = get_device() - self.raw_loc = torch.tensor([0, 1, 2, 3], device=self.device) - self.mapping_v1 = _make_mapping( - [0, 1, 2, 3], [10, 11, 12, 13], device=self.device - ) - self.mapping_v2 = _make_mapping( - [0, 1, 2, 3], [20, 21, 22, 23], device=self.device - ) - - def test_stale_without_fix(self): - """Without the fix, register_mapping() mid-forward leaves a stale cached_loc. - - Sequence: - 1. start_layer: cached_loc = translate(mapping_v1) = [10..13] - 2. HiCache load-back: register_mapping(mapping_v2) ← buggy, no clear - 3. start_layer+1: layer_id != start_layer, cached_loc is not None - → returns stale [10..13], NOT the correct [20..23] - """ - stub = _DSV4CacheStub(self.device) - stub.register_mapping_buggy(self.mapping_v1) - - # Forward pass — first SWA layer primes the cache. - loc = stub.get_swa_loc(stub.start_layer, self.raw_loc) - self.assertEqual(loc.tolist(), [10, 11, 12, 13]) - - # HiCache load-back installs new mapping (buggy path). - stub.register_mapping_buggy(self.mapping_v2) - self.assertIsNotNone(stub.cached_loc, "Bug: cached_loc not cleared") - - # Next SWA layer — should use mapping_v2 but returns mapping_v1. - loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc) - self.assertEqual( - loc_next.tolist(), - [10, 11, 12, 13], - "Confirms the bug: stale cached_loc [10..13] returned instead of [20..23]", - ) - - def test_correct_with_fix(self): - """With the fix, register_mapping() clears cached_loc; next layer recomputes. - - Same sequence as test_stale_without_fix but using the fixed register_mapping. - """ - stub = _DSV4CacheStub(self.device) - stub.register_mapping_fixed(self.mapping_v1) - - # Forward pass — first SWA layer primes the cache. - loc = stub.get_swa_loc(stub.start_layer, self.raw_loc) - self.assertEqual(loc.tolist(), [10, 11, 12, 13]) - - # HiCache load-back installs new mapping (fixed path). - stub.register_mapping_fixed(self.mapping_v2) - self.assertIsNone( - stub.cached_loc, "Fix: cached_loc cleared by register_mapping" - ) - - # Next SWA layer — recomputes with mapping_v2. - loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc) - self.assertEqual( - loc_next.tolist(), - [20, 21, 22, 23], - "Fix works: fresh translation [20..23] from new mapping", - ) - - -class TestDSV4ActualPoolRegisterMapping(CustomTestCase): - """Exercises the production DeepSeekV4TokenToKVPool.register_mapping() directly. - - Uses __new__ to bypass the complex __init__ (which needs full GPU pool setup) - and tests only the register_mapping / cached_loc contract. - """ - - def test_register_mapping_clears_cached_loc(self): - device = get_device() - - # Bypass full __init__ — only the fields register_mapping touches matter. - pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool) - old_loc = torch.tensor([10, 11, 12, 13], device=device) - pool.cached_loc = old_loc - pool.full_to_swa_index_mapping = None - - new_mapping = torch.arange(64, dtype=torch.int64, device=device) - pool.register_mapping(new_mapping) - - self.assertIsNone(pool.cached_loc, "register_mapping must clear cached_loc") - self.assertIs(pool.full_to_swa_index_mapping, new_mapping) - - def test_register_mapping_clears_none_cached_loc(self): - """Idempotent when cached_loc is already None.""" - pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool) - pool.cached_loc = None - pool.full_to_swa_index_mapping = None - - mapping = torch.arange(16, dtype=torch.int64, device=get_device()) - pool.register_mapping(mapping) - - self.assertIsNone(pool.cached_loc) - self.assertIs(pool.full_to_swa_index_mapping, mapping) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/manual/core/test_dsv4_hicache_swa_translation_cache.py b/test/manual/core/test_dsv4_hicache_swa_translation_cache.py deleted file mode 100644 index b888149e8..000000000 --- a/test/manual/core/test_dsv4_hicache_swa_translation_cache.py +++ /dev/null @@ -1,114 +0,0 @@ -"""E2E regression for PR #25889: DSV4 cached_loc stale after HiCache load-back. - -Bug (pre-fix): - DeepSeekV4TokenToKVPool.register_mapping() replaces full_to_swa_index_mapping - but does NOT clear self.cached_loc. When SGLANG_OPT_CACHE_SWA_TRANSLATION=True, - set_swa_key_buffer_radix_fused() caches the full→SWA translation across SWA - layers. After a HiCache commit/load-back that calls register_mapping() with a - new mapping, subsequent SWA layers reuse the stale cached_loc and write KV to - wrong SWA pool slots — producing divergent logprobs. - -Fix (PR #25889): - register_mapping() sets self.cached_loc = None so the first SWA layer in the - next forward pass recomputes the translation from the fresh mapping. - -Test strategy: - Subclass the existing DSV4 Flash HiCache KL suite and activate the SWA - translation cache via SGLANG_OPT_CACHE_SWA_TRANSLATION=1. The mixin tests - (test_multiturn_logprobs_match, test_multiturn_prefill_cache_hit_branching, - test_multiturn_decode_cache_hit_branching) compare logprobs from cold and - warm radix-cache hits. Without the fix the stale translation corrupts SWA KV - data and the KL divergence exceeds the threshold; with the fix it stays within. -""" - -import unittest - -from test_unified_radix_cache_kl_hicache import ( - DSV4_FLASH_LAUNCH_TIMEOUT, - DSV4_FLASH_MODEL, - _assert_dsv4_decode_cached_tokens, -) - -from sglang.srt.utils import kill_process_tree -from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin -from sglang.test.kl_multiturn_utils import get_input_ids -from sglang.test.test_utils import ( - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - - -class TestDSV4HiCacheSWATranslationCache(UnifiedRadixTreeTestMixin, CustomTestCase): - """DSV4 Flash FP8 + HiCache + SWA translation cache enabled. - - Identical server config to TestUnifiedDeepSeekV4FlashHiCache but with - SGLANG_OPT_CACHE_SWA_TRANSLATION=1 to activate the cached_loc path. - Without PR #25889 the KL tests fail; with the fix they pass. - """ - - kl_threshold = 0.005 - sampling_temperature = 0 - decode_hit_request_batch_size = 3 - decode_hit_inter_batch_delay_s = 0.5 - decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens) - gsm8k_threshold = 0.90 - num_gsm8k_questions = 100 - - @unittest.skipIf(True, "Covered by test_multiturn_prefill_cache_hit_branching.") - def test_multiturn_logprobs_match(self): - pass - - @classmethod - def setUpClass(cls): - cls.model = DSV4_FLASH_MODEL - cls.base_url = DEFAULT_URL_FOR_TEST - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DSV4_FLASH_LAUNCH_TIMEOUT, - other_args=[ - "--trust-remote-code", - "--tp-size", - "4", - "--attention-backend", - "compressed", - "--page-size", - "256", - "--chunked-prefill-size", - "8192", - "--mem-fraction-static", - "0.9", - "--disable-shared-experts-fusion", - "--enable-hierarchical-cache", - "--hicache-ratio", - "4", - "--hicache-write-policy", - "write_through", - "--hicache-io-backend", - "direct", - "--hicache-mem-layout", - "page_first_direct", - "--swa-full-tokens-ratio", - "0.25", - "--max-total-tokens", - "20000", - "--max-running-requests", - "4", - ], - env={ - "SGLANG_DSV4_FP4_EXPERTS": "0", - "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - # Activate the SWA translation cache — the flag that exposes the bug. - "SGLANG_OPT_CACHE_SWA_TRANSLATION": "1", - }, - ) - cls.input_ids = get_input_ids(cls.model, num_samples=18) - - @classmethod - def tearDownClass(cls): - kill_process_tree(cls.process.pid) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/manual/core/test_dsv4_stale_loc_crash.py b/test/manual/core/test_dsv4_stale_loc_crash.py deleted file mode 100644 index 3b31ace55..000000000 --- a/test/manual/core/test_dsv4_stale_loc_crash.py +++ /dev/null @@ -1,164 +0,0 @@ -"""Crash regression for PR #25889: stale cached_loc after register_mapping(). - -Bug: - DeepSeekV4TokenToKVPool caches the full→SWA index translation in - self.cached_loc (when SGLANG_OPT_CACHE_SWA_TRANSLATION=True). - register_mapping() was replacing full_to_swa_index_mapping without - clearing cached_loc, so a subsequent set_swa_key_buffer_radix_fused - call would use stale SWA indices and, if those indices exceed the - current pool size, raise a RuntimeError (OOB tensor access). - -Crash scenario reproduced here: - Pass 1 (large SWA pool, size=8): first SWA layer primes - cached_loc = [4, 5, 6, 7]. - register_mapping() called with new mapping (pre-fix: cache not cleared). - Pass 2 (smaller SWA pool, size=4): same SWA layer finds cached_loc is - not None → uses stale [4, 5, 6, 7] → OOB write on size-4 pool → - RuntimeError. - -Fix: register_mapping() sets self.cached_loc = None so the next call -recomputes with the fresh mapping. - -Run with: - python -m pytest test/registered/unit/mem_cache/test_dsv4_stale_loc_crash.py -v -""" - -import unittest - -import torch - -from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=3, suite="base-a-test-cpu") - -_NUM_HEADS = 2 -_HEAD_DIM = 8 -_SWA_LARGE = 8 # size-8 pool used during pass 1 -_SWA_SMALL = 4 # size-4 pool used during pass 2 (simulates post-HiCache loadback) - - -class _SWAPoolMock: - """Minimal SWA pool mock whose set_key_buffer_fused does a real tensor write. - - A write with OOB swa_loc raises RuntimeError, reproducing the crash. - """ - - def __init__(self, size: int): - self.buf = torch.zeros(size, _NUM_HEADS, _HEAD_DIM) - - def set_key_buffer_fused( - self, local_layer_id: int, swa_loc: torch.Tensor, cache_k: torch.Tensor - ) -> None: - n = swa_loc.numel() - self.buf[swa_loc.long()] = cache_k[:n].reshape(n, _NUM_HEADS, _HEAD_DIM) - - -def _build_pool( - mapping: torch.Tensor, - swa_pool: _SWAPoolMock, - start_layer: int = 0, -) -> DeepSeekV4TokenToKVPool: - """Create a minimal DeepSeekV4TokenToKVPool via __new__, bypassing __init__.""" - pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool) - pool.cached_loc = None - pool._should_cache_swa = True - pool.start_layer = start_layer - pool.full_to_swa_index_mapping = mapping - pool.swa_kv_pool = swa_pool - # _swa_local_layer_id: map global SWA layer id → local index 0 for this test. - pool._swa_local_layer_id = lambda lid: 0 - return pool - - -def _mapping(indices: list, size: int = 16) -> torch.Tensor: - m = torch.full((size,), -1, dtype=torch.int64) - for i, v in enumerate(indices): - m[i] = v - return m - - -class TestDSV4StaleLocCrash(CustomTestCase): - """ - Two paired tests that together constitute the crash regression for #25889. - - test_crash_without_fix: reproduces the RuntimeError that occurs when - register_mapping() does NOT clear cached_loc. - test_fix_prevents_crash: same sequence with the fixed register_mapping() - → no error, correct SWA slots written. - """ - - def setUp(self): - # raw_loc: full-pool indices for 4 tokens. - self.raw_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int64) - # cache_k: synthetic key data for 4 tokens. - self.cache_k = torch.ones(4, _NUM_HEADS, _HEAD_DIM) - # SWA layer id > start_layer (0), so cached_loc is only reset by the - # "is None" branch, NOT by the "layer_id == start_layer" branch. - self.swa_layer_id = 1 - # mapping_v1: raw_loc [0,1,2,3] → SWA slots [4,5,6,7] (valid for size-8 pool). - self.mapping_v1 = _mapping([4, 5, 6, 7]) - # mapping_v2: same raw_loc → SWA slots [0,1,2,3] (valid for size-4 pool). - self.mapping_v2 = _mapping([0, 1, 2, 3]) - - def test_crash_without_fix(self): - """Without the fix, stale cached_loc [4,5,6,7] causes RuntimeError on - the size-4 pool used in pass 2.""" - large_pool = _SWAPoolMock(_SWA_LARGE) - pool = _build_pool(self.mapping_v1, large_pool, start_layer=0) - - # Pass 1: swa_layer_id=1, cached_loc is None → compute and cache [4,5,6,7]. - pool.set_swa_key_buffer_radix_fused( - self.swa_layer_id, self.raw_loc, self.cache_k - ) - self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7], "Pass 1: cache primed") - - # PRE-FIX register_mapping: replace mapping WITHOUT clearing cached_loc. - pool.full_to_swa_index_mapping = self.mapping_v2 - - # Pass 2 on a smaller pool (size=4). cached_loc is still [4,5,6,7]. - # OOB write → RuntimeError. - pool.swa_kv_pool = _SWAPoolMock(_SWA_SMALL) - with self.assertRaises((RuntimeError, IndexError)): - pool.set_swa_key_buffer_radix_fused( - self.swa_layer_id, self.raw_loc, self.cache_k - ) - - def test_fix_prevents_crash(self): - """With the fix, register_mapping() clears cached_loc. Pass 2 recomputes - [0,1,2,3] from mapping_v2 and writes to the correct size-4 pool slots.""" - large_pool = _SWAPoolMock(_SWA_LARGE) - pool = _build_pool(self.mapping_v1, large_pool, start_layer=0) - - # Pass 1: prime cache → cached_loc = [4,5,6,7]. - pool.set_swa_key_buffer_radix_fused( - self.swa_layer_id, self.raw_loc, self.cache_k - ) - self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7]) - - # FIXED register_mapping: clears cached_loc. - pool.register_mapping(self.mapping_v2) - self.assertIsNone( - pool.cached_loc, "Fix: cached_loc cleared by register_mapping" - ) - - # Pass 2 on the smaller pool: cached_loc is None → recompute with mapping_v2. - small_pool = _SWAPoolMock(_SWA_SMALL) - pool.swa_kv_pool = small_pool - pool.set_swa_key_buffer_radix_fused( - self.swa_layer_id, self.raw_loc, self.cache_k - ) - - self.assertEqual( - pool.cached_loc.tolist(), [0, 1, 2, 3], "Fresh indices after fix" - ) - # Verify data landed in the correct SWA slots [0-3], not the stale [4-7]. - self.assertTrue( - small_pool.buf[0:4].abs().sum().item() > 0, - "Correct SWA slots 0-3 received data", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/manual/core/test_swa_loc_translation_cache.py b/test/manual/core/test_swa_loc_translation_cache.py deleted file mode 100644 index 27564ecc6..000000000 --- a/test/manual/core/test_swa_loc_translation_cache.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Manual tests for SWAKVPool.translate_loc_from_full_to_swa cache behaviour. - -These tests cover three properties introduced by PR #25824: - - 1. Cache key uses data_ptr() — correctly distinguishes views at different - offsets within the same storage (untyped_storage().data_ptr() would not). - - 2. Allocator mutations invalidate the cache — alloc/free/clear/ - set_full_to_swa_mapping each call invalidate_loc_cache() so the next - translation sees the fresh mapping. - - 3. BaseSWAKVPool.invalidate_loc_cache is a no-op default — subclasses that - don't cache (e.g. DSV4) can be called safely without AttributeError. - -Run with: - python -m pytest test/manual/core/test_swa_loc_translation_cache.py -v -""" - -import unittest - -import torch - -from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool -from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator -from sglang.srt.utils import get_device -from sglang.test.test_utils import CustomTestCase - - -def _build_pool( - kv_size: int = 32, - kv_size_swa: int = 32, - page_size: int = 1, -): - device = get_device() - num_layers = 8 - full_layer_ids = [0, 4] - swa_layer_ids = [i for i in range(num_layers) if i not in set(full_layer_ids)] - - pool = SWAKVPool( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=torch.bfloat16, - head_num=4, - head_dim=64, - swa_attention_layer_ids=swa_layer_ids, - full_attention_layer_ids=full_layer_ids, - enable_kvcache_transpose=False, - device=device, - ) - allocator = SWATokenToKVPoolAllocator( - size=kv_size, - size_swa=kv_size_swa, - page_size=page_size, - dtype=torch.bfloat16, - device=device, - kvcache=pool, - need_sort=False, - ) - return pool, allocator, device - - -class TestCacheKeyDataPtr(CustomTestCase): - """Cache key uses data_ptr(), which encodes the storage offset.""" - - def test_same_offset_view_is_cache_hit(self): - """Two different Python objects pointing to the same base are a hit.""" - pool, allocator, device = _build_pool() - loc = allocator.alloc(4) - self.assertIsNotNone(loc) - - # Create two slice objects at offset 0 — same data_ptr, same numel. - view_a = loc[:4] - view_b = loc[:4] - self.assertIsNot(view_a, view_b) # different Python objects - self.assertEqual(view_a.data_ptr(), view_b.data_ptr()) - - result_a = pool.translate_loc_from_full_to_swa(view_a) - result_b = pool.translate_loc_from_full_to_swa(view_b) - # Both should return the identical tensor (cache hit). - self.assertIs(result_a, result_b) - - def test_different_offset_view_is_cache_miss(self): - """Views at different offsets produce different data_ptr → cache miss.""" - pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32) - loc = allocator.alloc(10) - self.assertIsNotNone(loc) - self.assertGreaterEqual(loc.numel(), 10) - - view_lo = loc[0:5] - view_hi = loc[5:10] - self.assertEqual(view_lo.numel(), view_hi.numel()) # same numel - # Different data_ptr (different storage offset). - self.assertNotEqual(view_lo.data_ptr(), view_hi.data_ptr()) - - # Prime the cache with view_lo. - result_lo = pool.translate_loc_from_full_to_swa(view_lo) - # view_hi should be a cache miss and produce a distinct translation. - result_hi = pool.translate_loc_from_full_to_swa(view_hi) - # They should NOT be the same object (different cache entries). - self.assertIsNot(result_lo, result_hi) - # And the content must differ (different full indices → different swa). - self.assertFalse(torch.equal(result_lo, result_hi)) - - def test_storage_base_ptr_would_collide(self): - """Demonstrate that untyped_storage().data_ptr() WOULD collide for the - two views above — confirming data_ptr() is the right key.""" - t = torch.arange(20, device=get_device()) - a, b = t[0:10], t[5:15] - # Same storage base — old key would collide. - self.assertEqual(a.untyped_storage().data_ptr(), b.untyped_storage().data_ptr()) - self.assertEqual(a.numel(), b.numel()) - # But data_ptr differs — new key is safe. - self.assertNotEqual(a.data_ptr(), b.data_ptr()) - - -class TestAllocatorMutationInvalidation(CustomTestCase): - """Each allocator method that writes the mapping calls invalidate_loc_cache.""" - - def _prime_and_check_invalidation(self, pool, allocator, mutate_fn): - """Helper: prime cache, mutate, assert fresh translation.""" - loc = allocator.alloc(4) - self.assertIsNotNone(loc) - # Prime the cache. - first = pool.translate_loc_from_full_to_swa(loc) - self.assertIsNotNone(pool._cached_loc_key) - - # Mutate — should invalidate. - mutate_fn(allocator, loc) - - # Cache must be cleared after mutation. - self.assertIsNone(pool._cached_loc_key) - self.assertIsNone(pool._cached_swa_loc) - - def test_alloc_invalidates(self): - pool, allocator, _ = _build_pool() - loc = allocator.alloc(4) - pool.translate_loc_from_full_to_swa(loc) - self.assertIsNotNone(pool._cached_loc_key) - # Another alloc should invalidate. - allocator.alloc(4) - self.assertIsNone(pool._cached_loc_key) - - def test_free_swa_invalidates(self): - pool, allocator, _ = _build_pool() - loc = allocator.alloc(4) - pool.translate_loc_from_full_to_swa(loc) - self.assertIsNotNone(pool._cached_loc_key) - allocator.free_swa(loc) - self.assertIsNone(pool._cached_loc_key) - - def test_clear_invalidates(self): - pool, allocator, _ = _build_pool() - loc = allocator.alloc(4) - pool.translate_loc_from_full_to_swa(loc) - self.assertIsNotNone(pool._cached_loc_key) - allocator.clear() - self.assertIsNone(pool._cached_loc_key) - - def test_set_full_to_swa_mapping_invalidates(self): - """HiCache load-back path: set_full_to_swa_mapping must invalidate.""" - pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32) - loc = allocator.alloc(4) - pool.translate_loc_from_full_to_swa(loc) - self.assertIsNotNone(pool._cached_loc_key) - - # Simulate HiCache rebuild with new swa indices. - new_swa = torch.arange(4, dtype=torch.int64, device=device) - allocator.set_full_to_swa_mapping(loc, new_swa) - - self.assertIsNone(pool._cached_loc_key) - # Translation after rebuild should reflect the new mapping. - result = pool.translate_loc_from_full_to_swa(loc) - self.assertEqual(result.tolist(), new_swa.tolist()) - - -class TestBaseClassNoOp(CustomTestCase): - """BaseSWAKVPool.invalidate_loc_cache is a no-op default — must not raise.""" - - def test_noop_does_not_raise(self): - # BaseSWAKVPool is abstract; instantiate via SWAKVPool which inherits. - pool, _, _ = _build_pool() - # Calling on the concrete class uses the override — that's fine. - pool.invalidate_loc_cache() # must not raise - pool.invalidate_loc_cache() # idempotent - - def test_base_class_noop_directly(self): - """Call the base-class method directly to verify it's a true no-op.""" - pool, _, _ = _build_pool() - # Prime the cache first. - loc = pool.full_to_swa_index_mapping # any tensor - pool._cached_loc_key = ("dummy", 1) - pool._cached_swa_loc = torch.zeros(1) - # Call the BASE class method directly — should not clear the cache - # (it's a no-op; the concrete override is what clears). - BaseSWAKVPool.invalidate_loc_cache(pool) - # base no-op: cache untouched - self.assertIsNotNone(pool._cached_loc_key) - - -class TestExplicitInvalidationCycle(CustomTestCase): - """Simulates the per-forward-pass invalidation done by model_runner.""" - - def test_fresh_translation_after_explicit_invalidation(self): - """After invalidate_loc_cache(), a new alloc produces the right mapping.""" - pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32) - - # First "forward pass": alloc 4 tokens, translate. - loc1 = allocator.alloc(4) - trans1 = pool.translate_loc_from_full_to_swa(loc1).clone() - - # Simulate start of next forward pass: model_runner calls invalidate. - pool.invalidate_loc_cache() - self.assertIsNone(pool._cached_loc_key) - - # Alloc 4 more (mapping changes), translate loc1 again. - loc2 = allocator.alloc(4) - # Alloc already invalidated; translate loc1 with fresh mapping. - trans1_after = pool.translate_loc_from_full_to_swa(loc1) - - # loc1's SWA mapping hasn't changed (same full→swa assignment), - # so result should be equal — but it must have been recomputed - # (cache key was None before this call). - self.assertEqual(trans1.tolist(), trans1_after.tolist()) - - # loc2 should have different translation than loc1. - trans2 = pool.translate_loc_from_full_to_swa(loc2) - # They have different indices, so translation differs. - self.assertFalse(torch.equal(trans1_after, trans2)) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index f2e07daba..0166025eb 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py @@ -16,6 +16,7 @@ import importlib.util import sys import unittest from pathlib import Path +from types import SimpleNamespace import torch @@ -337,5 +338,75 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase): run_dsv4_eagle_draft_extend_cuda_graph_runner_case(self, case) +class TestDSV4SwaOutCacheLocResolution(CustomTestCase): + """`get_swa_out_cache_loc`: cached fast path vs store-time fallback. + + The KV-store consumers run in paths that never invoke + `init_forward_metadata_in_graph` (eager idle, runners that only run the + out-graph prep) or whose batch is re-padded after init (DP attention). + The resolver must use the per-forward cached value only when it is + provably current and fall back to translating `out_cache_loc` otherwise. + """ + + def _make_backend(self, mapping: torch.Tensor): + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + ) + + backend = object.__new__(DeepseekV4AttnBackend) + backend.forward_metadata = None + backend.token_to_kv_pool = SimpleNamespace( + translate_loc_from_full_to_swa=lambda loc: mapping[loc] + ) + return backend + + @staticmethod + def _make_fb(out_cache_loc: torch.Tensor, forward_mode: ForwardMode): + return SimpleNamespace(out_cache_loc=out_cache_loc, forward_mode=forward_mode) + + @staticmethod + def _set_cached(backend, cached: torch.Tensor): + backend.forward_metadata = SimpleNamespace( + core_attn_metadata=SimpleNamespace(swa_out_cache_loc=cached) + ) + + def test_no_metadata_falls_back_to_translate(self): + mapping = torch.arange(10, dtype=torch.int64) * 2 + backend = self._make_backend(mapping) + fb = self._make_fb(torch.tensor([3, 4]), ForwardMode.DECODE) + out = backend.get_swa_out_cache_loc(fb) + self.assertEqual(out.dtype, torch.int32) + self.assertEqual(out.tolist(), [6, 8]) + + def test_current_cached_value_is_used(self): + mapping = torch.arange(10, dtype=torch.int64) * 2 + backend = self._make_backend(mapping) + cached = torch.tensor([6, 8], dtype=torch.int32) + self._set_cached(backend, cached) + fb = self._make_fb(torch.tensor([3, 4]), ForwardMode.DECODE) + self.assertIs(backend.get_swa_out_cache_loc(fb), cached) + + def test_stale_shape_falls_back_to_translate(self): + # DP padding rebinds out_cache_loc to a longer tensor after init; + # a pre-pad cached value must not be used. + mapping = torch.arange(10, dtype=torch.int64) * 2 + backend = self._make_backend(mapping) + self._set_cached(backend, torch.tensor([6, 8], dtype=torch.int32)) + fb = self._make_fb(torch.tensor([3, 4, 0, 0]), ForwardMode.DRAFT_EXTEND_V2) + out = backend.get_swa_out_cache_loc(fb) + self.assertEqual(out.tolist(), [6, 8, 0, 0]) + + def test_idle_never_uses_cached_value(self): + # Idle forwards skip attn init, so any metadata is left over from a + # previous forward; writing dummy tokens to its locations would + # corrupt live KV. Idle must translate the zero-padded out_cache_loc. + mapping = torch.arange(10, dtype=torch.int64) * 2 + backend = self._make_backend(mapping) + self._set_cached(backend, torch.tensor([6, 8], dtype=torch.int32)) + fb = self._make_fb(torch.tensor([0, 0]), ForwardMode.IDLE) + out = backend.get_swa_out_cache_loc(fb) + self.assertEqual(out.tolist(), [0, 0]) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py index 3de2fe5ee..15ff77ea3 100644 --- a/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py +++ b/test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py @@ -33,7 +33,6 @@ def _make_self(*, page_size: int, full_available: int, swa_available: int): ), translate_loc_from_full_to_swa=lambda last_loc: last_loc, full_to_swa_index_mapping=torch.zeros(64, dtype=torch.int64), - _kvcache=SimpleNamespace(invalidate_loc_cache=lambda: None), )