Refactor MLA kv cache

python/sglang/srt/layers/attention/nsa_backend.py

@@ -35,6 +35,7 @@
     seqlens_expand_triton,
 )
 from sglang.srt.layers.dp_attention import get_attention_tp_size
+from sglang.srt.mem_cache.memory_pool import MLAKVCacheLayout
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.utils import is_cuda, is_hip
 
@@ -302,15 +303,13 @@ def __init__(
         )
         self.use_nsa = is_deepseek_nsa(model_runner.model_config.hf_config)
         assert self.use_nsa, "NSA backend only supports DeepSeek NSA"
-        self.nsa_kv_cache_store_fp8 = (
-            model_runner.token_to_kv_pool.nsa_kv_cache_store_fp8
-        )
+        self.kv_cache_layout = model_runner.token_to_kv_pool.kv_cache_layout
+        self.kv_cache_size = model_runner.token_to_kv_pool.kv_cache_size
         self.nsa_index_topk = get_nsa_index_topk(model_runner.model_config.hf_config)
         self.max_context_len = model_runner.model_config.context_len
         self.num_q_heads = (
             model_runner.model_config.num_attention_heads // get_attention_tp_size()
         )
-        self.kv_cache_dim = model_runner.token_to_kv_pool.kv_cache_dim
         self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
         self.kv_lora_rank = model_runner.model_config.kv_lora_rank
         self.qk_rope_head_dim = model_runner.model_config.qk_rope_head_dim
@@ -319,11 +318,12 @@ def __init__(
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
 
         self.use_mha: bool = False
+        nsa_prefill_backend = model_runner.server_args.nsa_prefill_backend
+        self.prefill_is_flashmla_auto = nsa_prefill_backend == "flashmla_auto"
         self.nsa_prefill_impl: _NSA_IMPL_T = (
-            model_runner.server_args.nsa_prefill_backend
+            None if self.prefill_is_flashmla_auto else nsa_prefill_backend
         )
         self.nsa_decode_impl: _NSA_IMPL_T = model_runner.server_args.nsa_decode_backend
-        self.enable_auto_select_prefill_impl = self.nsa_prefill_impl == "flashmla_auto"
 
         self._arange_buf = torch.arange(16384, device=self.device, dtype=torch.int32)
 
@@ -383,8 +383,22 @@ def _transform_table_1_to_real(self, page_table: torch.Tensor) -> torch.Tensor:
         )
         return page_table[:, strided_indices] // page_size
 
+    @staticmethod
+    def should_use_decode_backend(forward_mode: ForwardMode) -> bool:
+        """Whether the given forward mode dispatches to nsa_decode_backend.
+
+        This is answerable from static config alone (does not depend on
+        set_nsa_impl), so it is safe to call before init_forward_metadata.
+        """
+        return (
+            forward_mode.is_decode_or_idle()
+            or forward_mode.is_target_verify()
+            or forward_mode.is_draft_extend(include_v2=True)
+        )
+
     def init_forward_metadata(self, forward_batch: ForwardBatch):
-        """Init the metadata for a forward pass."""
+        """Sets the metadata as fields of NativeSparseAttnBackend for a single forward pass.
+        Using the forward metadata for more than one forward pass is a bug."""
         batch_size = forward_batch.batch_size
         device = forward_batch.seq_lens.device
 
@@ -403,10 +417,16 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
         ]
 
         page_table_1_flattened = None
+        # Topk indices needs to transform from indices in each request to indices
+        # in the flattened array of all requests by
+        #   topk_indices = topk_indices + topk_indices_offset
+        # Offset is repeated once for each token in the request in a flattened array.
+        # e.g, for three requests with lengths [2, 3, 4],
+        #   topk_indices_offset = [0, 0, 2, 2, 2, 5, 5, 5, 5]
+        # Only used for prefill with TopkTransformMethod.RAGGED.
         topk_indices_offset = None
 
-        # Centralized dispatch: decide all strategies for this batch
-        self.set_nsa_prefill_impl(forward_batch)
+        self.set_nsa_impl(forward_batch)
         topk_transform_method = self.get_topk_transform_method(
             forward_batch.forward_mode
         )
@@ -643,6 +663,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                     seq_len_q=1,
                 )
                 if self.nsa_decode_impl == "flashmla_kv"
+                or self.nsa_prefill_impl == "flashmla_kv"
                 else None
             ),
             paged_mqa_schedule_metadata=paged_mqa_schedule_metadata,
@@ -787,9 +808,8 @@ def init_forward_metadata_capture_cuda_graph(
         forward_mode: ForwardMode,
         spec_info: Optional[SpecInput],
     ):
-        self.set_nsa_prefill_impl(forward_batch=None)
+        self.set_nsa_impl(forward_batch=None)
 
-        """Initialize forward metadata for capturing CUDA graph."""
         if forward_mode.is_decode_or_idle():
             # Normal Decode
             # Get sequence information
@@ -944,10 +964,9 @@ def init_forward_metadata_replay_cuda_graph(
         seq_lens_cpu: Optional[torch.Tensor],
         out_cache_loc: Optional[torch.Tensor] = None,
     ):
-        """Initialize forward metadata for replaying CUDA graph."""
         assert seq_lens_cpu is not None
 
-        self.set_nsa_prefill_impl(forward_batch=None)
+        self.set_nsa_impl(forward_batch=None)
 
         seq_lens = seq_lens[:bs]
         seq_lens_cpu = seq_lens_cpu[:bs]
@@ -1104,16 +1123,15 @@ def init_forward_metadata_replay_cuda_graph_from_precomputed(
         precomputed: PrecomputedMetadata,
         forward_mode: ForwardMode,
     ):
-        """Fast path: copy precomputed metadata to this backend's metadata.
-
-        This function only performs copy operations, no computation.
+        """Compared to init_forward_metadata_replay_cuda_graph,
+        this function copies precomputed metadata instead of computing them.
 
         Args:
             bs: Batch size
             precomputed: Precomputed metadata to copy from
             forward_mode: Forward mode
         """
-        self.set_nsa_prefill_impl(forward_batch=None)
+        self.set_nsa_impl(forward_batch=None)
 
         metadata = self.decode_cuda_graph_metadata[bs]
 
@@ -1275,10 +1293,7 @@ def forward_extend(
 
         nsa_impl = (
             self.nsa_decode_impl
-            if (
-                forward_batch.forward_mode.is_target_verify()
-                or forward_batch.forward_mode.is_draft_extend(include_v2=True)
-            )
+            if self.should_use_decode_backend(forward_batch.forward_mode)
             else self.nsa_prefill_impl
         )
 
@@ -1399,18 +1414,18 @@ def forward_extend(
             if q_rope is not None:
                 q_all = concat_mla_absorb_q_general(q_nope, q_rope)
 
-            if topk_transform_method == TopkTransformMethod.RAGGED:
-                if any(forward_batch.extend_prefix_lens_cpu):
-                    page_table_1_flattened = (
-                        self.forward_metadata.page_table_1_flattened
-                    )
-                    assert page_table_1_flattened is not None
-                    kv_cache = dequantize_k_cache_paged(
-                        kv_cache, page_table_1_flattened
-                    )
-                else:
-                    kv_cache = _cat([k, k_rope], dim=-1)
-                page_table_1 = topk_indices
+            if topk_transform_method != TopkTransformMethod.RAGGED:
+                raise ValueError(
+                    "Internal error: Unexpected topk transform method for NSA backend flashmla_sparse."
+                )
+
+            if any(forward_batch.extend_prefix_lens_cpu):
+                page_table_1_flattened = self.forward_metadata.page_table_1_flattened
+                assert page_table_1_flattened is not None
+                kv_cache = dequantize_k_cache_paged(kv_cache, page_table_1_flattened)
+            else:
+                kv_cache = _cat([k, k_rope], dim=-1)
+            page_table_1 = topk_indices
 
             return self._forward_flashmla_sparse(
                 q_all=q_all,
@@ -1712,10 +1727,14 @@ def _forward_flashmla_kv(
 
         # TODO the 2nd dim is seq_len_q, need to be >1 when MTP
         q_all = q_all.view(-1, 1, layer.tp_q_head_num, layer.head_dim)
-        kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_dim)
+        kv_cache = kv_cache.view(-1, self.real_page_size, 1, self.kv_cache_size)
         assert self.real_page_size == 64, "only page size 64 is supported"
 
-        if not self.nsa_kv_cache_store_fp8:
+        if self.kv_cache_layout != MLAKVCacheLayout.FP8_NOPE_WITH_BLOCK_SCALE_BF16_ROPE:
+            assert (
+                self.kv_cache_layout != MLAKVCacheLayout.FP8_NOPE_FP8_ROPE
+            ), "Internal error: NSA backend flashmla_kv does not support FP8_NOPE_FP8_ROPE"
+
             # inefficiently quantize the whole cache
             kv_cache = quantize_k_cache(kv_cache)
 
@@ -1937,6 +1956,9 @@ def _forward_trtllm(
 
         merge_query = q_rope is not None
         if self.kv_cache_dtype == torch.float8_e4m3fn:
+            assert (
+                self.kv_cache_layout == MLAKVCacheLayout.FP8_NOPE_FP8_ROPE
+            ), "Internal error: trtllm mla backend only supports FP8_NOPE_FP8_ROPE"
             # For FP8 path, we quantize the query and rope parts and merge them into a single tensor
             # Note: rope application in deepseek_v2.py:forward_absorb_prepare is skipped for FP8 decode path of this trtllm_mla backend
             assert q_rope is not None, "For FP8 path q_rope should not be None."
@@ -1973,7 +1995,9 @@ def _forward_trtllm(
             )
 
         k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
-        kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_dim).unsqueeze(1)
+        kv_cache = k_cache.view(-1, self.real_page_size, self.kv_cache_size).unsqueeze(
+            1
+        )
 
         if merge_query:
             q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
@@ -2009,7 +2033,7 @@ def _forward_trtllm(
         _, num_heads, head_dim = q_all.shape
 
         q = q_all.view(batch_size, 1, num_heads, head_dim)
-        kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_dim)
+        kv = kv_cache.view(-1, 1, self.real_page_size, self.kv_cache_size)
         block_tables = page_table_1.unsqueeze(1)
         seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens
 
@@ -2056,9 +2080,10 @@ def get_cuda_graph_seq_len_fill_value(self):
         """Get the fill value for sequence length in CUDA graph."""
         return 1
 
-    def set_nsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None):
+    def set_nsa_impl(self, forward_batch: Optional[ForwardBatch] = None):
         """
-        Decide all attention prefill dispatch strategies for this batch.
+        Decide all attention dispatch strategies for this batch.
+        Sets nsa_prefill_impl, nsa_decode_impl and use_mha depending on forward mode.
         """
         from sglang.srt.utils import get_device_sm, is_blackwell
 
@@ -2087,9 +2112,17 @@ def set_nsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None):
         else:
             self.use_mha = False  # Decode/verify always use MLA
 
+        # forward_batch is None only for cudagraph
+        if forward_batch is None or forward_batch.forward_mode.is_decode_or_idle():
+            assert self.nsa_decode_impl != "flashmla_auto"
+            return
+
         # Set MLA implementation only if not using MHA
-        if not self.use_mha and self.enable_auto_select_prefill_impl:
-            if self.nsa_kv_cache_store_fp8:
+        if not self.use_mha and self.prefill_is_flashmla_auto:
+            if (
+                self.kv_cache_layout
+                == MLAKVCacheLayout.FP8_NOPE_WITH_BLOCK_SCALE_BF16_ROPE
+            ):
                 if (
                     is_blackwell()
                     and forward_batch is not None
@@ -2109,19 +2142,13 @@ def set_nsa_prefill_impl(self, forward_batch: Optional[ForwardBatch] = None):
     def get_topk_transform_method(
         self, forward_mode: Optional[ForwardMode] = None
     ) -> TopkTransformMethod:
-        """
-        SGLANG_NSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel.
-        This method is used to select the topk transform method which can be fused or unfused.
-        """
-        if (
-            # disable for MTP
-            self.nsa_kv_cache_store_fp8
+        if forward_mode is None or forward_mode.is_decode_or_idle():
+            return TopkTransformMethod.PAGED
+        elif (
+            self.kv_cache_layout == MLAKVCacheLayout.FP8_NOPE_WITH_BLOCK_SCALE_BF16_ROPE
             and self.nsa_prefill_impl == "flashmla_sparse"
         ):
             topk_transform_method = TopkTransformMethod.RAGGED
-
-            if forward_mode is not None and (forward_mode.is_decode_or_idle()):
-                topk_transform_method = TopkTransformMethod.PAGED
         else:
             topk_transform_method = TopkTransformMethod.PAGED
         return topk_transform_method
@@ -2133,11 +2160,12 @@ def get_indexer_metadata(
             forward_batch.hisparse_coordinator is not None
             and forward_batch.forward_mode.is_decode_or_idle()
         )
+        topk_transform_method = self.get_topk_transform_method(
+            forward_batch.forward_mode
+        )
         return NSAIndexerMetadata(
             attn_metadata=self.forward_metadata,
-            topk_transform_method=self.get_topk_transform_method(
-                forward_batch.forward_mode
-            ),
+            topk_transform_method=topk_transform_method,
             paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata,
             force_unfused_topk=force_unfused,
         )
@@ -2229,7 +2257,7 @@ def init_forward_metadata_replay_cuda_graph(
 
                     # Set nsa_prefill_impl for first 3 backends (required by the method)
                     for i in range(3):
-                        self.attn_backends[i].set_nsa_prefill_impl(forward_batch=None)
+                        self.attn_backends[i].set_nsa_impl(forward_batch=None)
 
                     # Prepare FlashMLA tensors if needed
                     flashmla_num_splits_src = None

python/sglang/srt/managers/hisparse_coordinator.py

@@ -53,7 +53,7 @@ def __init__(
             host_size=0,
             page_size=1,  # for simplicity, we set page size to 1 to enable backup one token at a time
             layout="layer_first",
-            override_kv_cache_dim=self.mem_pool_device.kv_cache_dim,
+            override_kv_cache_dim=self.mem_pool_device.kv_cache_size,
         )
 
         max_num_reqs = req_to_token_pool.size

python/sglang/srt/mem_cache/hisparse_memory_pool.py

@@ -10,7 +10,7 @@
     BaseTokenToKVPoolAllocator,
     PagedTokenToKVPoolAllocator,
 )
-from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
+from sglang.srt.mem_cache.memory_pool import MLAKVCacheLayout, NSATokenToKVPool
 from sglang.srt.utils import is_cuda, is_hip
 
 # sgl_kernel.kvcacheio is only available in CUDA/ROCm sgl-kernel builds (not XPU/MPS/NPU/CPU).
@@ -28,9 +28,10 @@ def transfer_kv_all_layer_mla(*args, **kwargs):
 
 
 class HiSparseNSATokenToKVPool(NSATokenToKVPool):
+
     def __init__(
         self,
-        size: int,
+        max_total_num_tokens: int,
         page_size: int,
         kv_lora_rank: int,
         dtype: torch.dtype,
@@ -39,13 +40,14 @@ def __init__(
         device: str,
         index_head_dim: int,
         enable_memory_saver: bool,
-        kv_cache_dim: int,
+        kv_cache_layout: MLAKVCacheLayout,
+        kv_cache_size: int,
         start_layer: Optional[int] = None,
         end_layer: Optional[int] = None,
         host_to_device_ratio: int = 2,
     ):
         super().__init__(
-            size=size,
+            max_total_num_tokens=max_total_num_tokens,
             page_size=page_size,
             kv_lora_rank=kv_lora_rank,
             dtype=dtype,
@@ -54,12 +56,13 @@ def __init__(
             device=device,
             index_head_dim=index_head_dim,
             enable_memory_saver=enable_memory_saver,
-            kv_cache_dim=kv_cache_dim,
+            kv_cache_layout=kv_cache_layout,
+            kv_cache_size=kv_cache_size,
             start_layer=start_layer,
             end_layer=end_layer,
-            index_buf_size=size * host_to_device_ratio,
+            index_buf_size=max_total_num_tokens * host_to_device_ratio,
         )
-        self.bytes_per_token = self.kv_cache_dim * self.dtype.itemsize
+        self.bytes_per_token = self.kv_cache_size * self.dtype.itemsize
 
     def register_mapping(self, full_to_hisparse_device_index_mapping: torch.Tensor):
         self.full_to_hisparse_device_index_mapping = (