[Feat] KVShard: Add attention and indexer support for page-sharded DSA caches

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -152,6 +152,7 @@ def _decline_fused(reason: str) -> None:
     from aiter.ops.cache import indexer_k_quant_and_cache
 
 from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.layers.attention.kv_shard_hooks import get_kv_shard_pool
 from sglang.srt.layers.cp.base import get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_active
 from sglang.srt.layers.linear import ReplicatedLinear
@@ -240,6 +241,14 @@ def _make_eager_idle_topk_result(
     )
 
 
+def _index_k_write_targets(pool, layer_id: int, loc: torch.Tensor):
+    """(index-K buffer, token slots, selected token rows or None) to write."""
+    shard_pool = get_kv_shard_pool(pool)
+    if shard_pool is None:
+        return ((pool.get_index_k_with_scale_buffer(layer_id=layer_id), loc, None),)
+    return shard_pool.index_key_cache.write_targets(layer_id, loc)
+
+
 def rotate_activation(x: torch.Tensor) -> torch.Tensor:
     if _is_hip:
         from fast_hadamard_transform import hadamard_transform
@@ -1026,17 +1035,18 @@ def _fused_k_prepare_and_store(
             and out_cache_loc is not None
             and can_use_dsa_fused_store(torch.bfloat16, out_cache_loc.dtype, page_size)
         ):
-            fused_k_indexer_norm_rope_store(
-                key_raw,
-                pool.get_index_k_with_scale_buffer(layer_id=layer_id),
-                out_cache_loc,
-                self.k_norm.weight,
-                self.k_norm.bias,
-                self.k_norm.variance_epsilon,
-                self._indexer_cos_sin_cache,
-                positions,
-                page_size,
-            )
+            for buf, loc, rows in _index_k_write_targets(pool, layer_id, out_cache_loc):
+                fused_k_indexer_norm_rope_store(
+                    key_raw if rows is None else key_raw.index_select(0, rows),
+                    buf,
+                    loc,
+                    self.k_norm.weight,
+                    self.k_norm.bias,
+                    self.k_norm.variance_epsilon,
+                    self._indexer_cos_sin_cache,
+                    positions if rows is None else positions.index_select(0, rows),
+                    page_size,
+                )
             return
 
         # Fallback: separate K kernel + store kernel.
@@ -1510,7 +1520,7 @@ def _get_topk_ragged(
         if _is_hip and not _use_aiter_preshuffle:
             block_tables = metadata.get_page_table_1()
         else:
-            block_tables = metadata.get_page_table_64()
+            block_tables = metadata.get_logical_page_table()
 
         assert (
             forward_batch.seq_lens_cpu is not None
@@ -1830,13 +1840,13 @@ def _store_index_k_cache(
             )
         ):
             # NOTE: wrapper already normalizes shape/contiguity and asserts dtypes.
-            buf = pool.get_index_k_with_scale_buffer(layer_id=layer_id)
-            fused_store_index_k_cache(
-                key,
-                buf,
-                out_cache_loc,
-                pool.index_page_size,
-            )
+            for buf, loc, rows in _index_k_write_targets(pool, layer_id, out_cache_loc):
+                fused_store_index_k_cache(
+                    key if rows is None else key.index_select(0, rows),
+                    buf,
+                    loc,
+                    pool.index_page_size,
+                )
             return
 
         # Fast path: AITER fused quant + cache store

python/sglang/srt/layers/attention/dsa/dsa_indexer_metadata.py

@@ -32,6 +32,10 @@ def get_page_table_64(self) -> torch.Tensor:
                 The page size of the table is 64.
         """
 
+    def get_logical_page_table(self) -> torch.Tensor:
+        """Page IDs for pool-managed reads, before any pool address translation."""
+        return self.get_page_table_64()
+
     @abstractmethod
     def get_page_table_1(self) -> torch.Tensor:
         """
@@ -105,6 +109,10 @@ def get_seqlens_int32(self) -> torch.Tensor:
     def get_page_table_64(self) -> torch.Tensor:
         return self.attn_metadata.real_page_table
 
+    def get_logical_page_table(self) -> torch.Tensor:
+        logical = self.attn_metadata.indexer_logical_page_table
+        return self.get_page_table_64() if logical is None else logical
+
     def get_page_table_1(self) -> torch.Tensor:
         return self.attn_metadata.page_table_1
 

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

@@ -87,6 +87,10 @@
     resolve_paged_mqa_logits_metadata_fn,
     should_use_dsa_fused_topk,
 )
+from sglang.srt.layers.attention.kv_shard_hooks import (
+    get_kv_shard_pool,
+    prepare_kv_shard_forward,
+)
 from sglang.srt.layers.attention.trtllm_mla_backend import (
     grow_multi_ctas_kv_counter_buffer_if_needed,
     make_persistent_multi_ctas_kv_counter_buffer,
@@ -295,6 +299,9 @@ class DSAMetadata:
 
     # k_start and k_end in kv cache for each token.
     indexer_k_start_end: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    # Pool-managed index reads take logical page IDs and translate them once.
+    # Raw attention/indexer kernels instead use the scratch-addressed tables.
+    indexer_logical_page_table: Optional[torch.Tensor] = None
     # seq lens for each batch.
     indexer_seq_lens_cpu: Optional[torch.Tensor] = None
     # seq lens for each batch.
@@ -551,11 +558,16 @@ def __init__(
         assert model_runner.req_to_token_pool is not None
         self.req_to_token_pool = model_runner.req_to_token_pool
         self.token_to_kv_pool = model_runner.token_to_kv_pool
+        self._kv_shard_pool = get_kv_shard_pool(self.token_to_kv_pool)
         self.hisparse_coordinator = model_runner.hisparse_coordinator
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
 
         self.use_mha: bool = False
-        self.supports_mha_one_shot: bool = self.dcp_size == 1
+        # Dense MHA expands prefix KV outside the DSA metadata path. Keep
+        # sharded prefill on the sparse kernels with explicit scratch addresses.
+        self.supports_mha_one_shot: bool = (
+            self.dcp_size == 1 and self._kv_shard_pool is None
+        )
         self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend
         self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend
         if (
@@ -1106,6 +1118,11 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
             forward_batch.req_pool_indices, :max_seqlen_k
         ]
 
+        kv_shard_extend = self._kv_shard_pool is not None and prepare_kv_shard_forward(
+            self._kv_shard_pool, self.req_to_token, forward_batch
+        )
+        indexer_logical_page_table = None
+
         page_table_1_flattened = None
         topk_indices_offset = None
 
@@ -1264,6 +1281,14 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                 )
                 page_table = page_table[bs_idx, :max_seqlen_k]
 
+            if kv_shard_extend:
+                # Kernels read the gathered prefix/chunk scratch; the indexer
+                # pool read takes logical pages and translates them itself.
+                indexer_logical_page_table = self._transform_table_1_to_real(page_table)
+                page_table = self._kv_shard_pool.translate_loc_to_scratch(
+                    page_table
+                ).to(page_table.dtype)
+
             if any(forward_batch.extend_prefix_lens_cpu) or bs_idx_cpu is not None:
                 max_seqlen_q = (
                     max(extend_seq_lens_cpu) if len(extend_seq_lens_cpu) != 0 else 1
@@ -1302,7 +1327,11 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                 kv_cache_capacity = (
                     self.token_to_kv_pool.size + self.token_to_kv_pool.page_size
                 )
-                if forward_batch.seq_lens_sum > kv_cache_capacity:
+                if (
+                    # Sharded extends index scratch rows, bounded by construction.
+                    not kv_shard_extend
+                    and forward_batch.seq_lens_sum > kv_cache_capacity
+                ):
                     max_idx = page_table_1_flattened.max().item()
                     assert max_idx < kv_cache_capacity, (
                         f"Invalid page table index: max={max_idx}, "
@@ -1381,6 +1410,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
             dsa_max_seqlen_q=1,
             topk_indices_offset=topk_indices_offset,
             indexer_k_start_end=indexer_k_start_end,
+            indexer_logical_page_table=indexer_logical_page_table,
             indexer_seq_lens_cpu=indexer_seq_lens_cpu,
             indexer_seq_lens=indexer_seq_lens,
             token_to_batch_idx=token_to_batch_idx,

python/sglang/srt/mem_cache/page_interleave_pool.py

@@ -780,39 +780,58 @@ def get_k_and_scale(
             max_seq_len,
         )
 
-    def store_quantized(
-        self,
-        layer_id: int,
-        loc: torch.Tensor,
-        index_k: torch.Tensor,
-        index_k_scale: torch.Tensor,
-    ) -> None:
+    def write_targets(
+        self, layer_id: int, loc: torch.Tensor
+    ) -> List[Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]]:
+        """(buffer, token slots, selected token rows or None) for a write.
+
+        During an extend every token is staged at its scratch slot; owned
+        tokens also persist at their local slots. Writers apply their own
+        quantization to each target, so sharded bytes match unsharded ones.
+        """
         if loc.numel() == 0:
-            return
+            return []
         pool = self.pool
-        assert self.buffer[layer_id - pool.start_layer].shape[0], (
-            "cannot write indexer keys for a skip-topk layer"
-        )
+        local = self.buffer[layer_id - pool.start_layer]
+        assert local.shape[0], "cannot write indexer keys for a skip-topk layer"
+        targets = []
         if pool._shard_extend_active:
-            dsa_index_k_cache.SetKAndS.execute(
-                pool=pool,
-                buf=pool._slots[layer_id % 2].tensors["index_k"],
-                loc=pool.translate_loc_to_scratch(loc),
-                index_k=index_k,
-                index_k_scale=index_k_scale,
+            targets.append(
+                (
+                    pool._slots[layer_id % 2].tensors["index_k"],
+                    pool.translate_loc_to_scratch(loc),
+                    None,
+                )
             )
         # Indexer callers may reuse or mutate their location tensors. Compute
         # fresh indices instead of memoizing only their pointer and length.
         owned_idx = torch.nonzero(
             pool.placement.local_mask(loc, pool.shard_rank)
         ).squeeze(1)
-        local_rows = pool.placement.local_index(loc[owned_idx])
         if owned_idx.numel():
-            super().store_quantized(
-                layer_id,
-                local_rows,
-                index_k.index_select(0, owned_idx),
-                index_k_scale.index_select(0, owned_idx),
+            targets.append(
+                (local, pool.placement.local_index(loc[owned_idx]), owned_idx)
+            )
+        return targets
+
+    def store_quantized(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        index_k: torch.Tensor,
+        index_k_scale: torch.Tensor,
+    ) -> None:
+        for buf, slots, rows in self.write_targets(layer_id, loc):
+            dsa_index_k_cache.SetKAndS.execute(
+                pool=self.pool,
+                buf=buf,
+                loc=slots,
+                index_k=index_k if rows is None else index_k.index_select(0, rows),
+                index_k_scale=(
+                    index_k_scale
+                    if rows is None
+                    else index_k_scale.index_select(0, rows)
+                ),
             )
 
 
@@ -840,8 +859,8 @@ class PageInterleaveDSATokenToKVPool(PageInterleaveMLATokenToKVPool, DSATokenToK
     expose owner-local storage; ``index_key_cache.get_local_buffer`` always
     exposes that storage.
 
-    This is a pool-level implementation. DSA backend metadata, fused indexer
-    stores, and model-runner selection still require separate integration.
+    This is a pool-level implementation. Model-runner selection still
+    requires separate integration.
     """
 
     def __init__(