[2/N] [Kernel] Fuse padding-preserving HiSparse slot translation (#39837)

python/sglang/kernels/ops/kvcache/hisparse_slot_mapping.py

@@ -0,0 +1,59 @@
+"""Padding-preserving logical-to-physical HiSparse slot translation."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _translate_padded_hisparse_locations(
+    mapping,
+    locations,
+    output,
+    count,
+    stride,
+    BLOCK: tl.constexpr,
+):
+    offset = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    loc = tl.load(locations + offset * stride, offset < count, other=-1)
+    physical = tl.load(
+        mapping + tl.maximum(loc, 0), (offset < count) & (loc >= 0), other=0
+    )
+    tl.store(output + offset, tl.where(loc >= 0, physical, loc), offset < count)
+
+
+def translate_padded_hisparse_locations(
+    mapping: torch.Tensor, locations: torch.Tensor
+) -> torch.Tensor:
+    """Translate logical token locations into physical GPU cache rows.
+
+    Tensor layout (all tensors are 1D):
+      mapping: contiguous [num_mapping_entries].
+        mapping[logical_slot] = physical GPU cache row for that logical slot.
+      locations: possibly strided [num_tokens].
+        locations[i] = logical slot for token i; negative values mean padding.
+      output: contiguous [num_tokens].
+        output[i] = mapping[locations[i]] for a nonnegative location;
+        otherwise output[i] = locations[i], preserving the padding value.
+
+    There is no layer axis: each layer uses the row numbers in its own KV buffer.
+    For example, mapping[17] = 3 and mapping[18] = 5 translate locations
+    [17, 18, -1] into [3, 5, -1]. Nonnegative locations must index within mapping.
+
+    Inputs are int32/int64 tensors on the same device and remain unchanged.
+    Output uses that device and the promoted integer dtype of both inputs.
+    """
+    assert mapping.ndim == locations.ndim == 1 and mapping.is_contiguous()
+    assert mapping.device == locations.device
+    assert mapping.dtype in (torch.int32, torch.int64)
+    assert locations.dtype in (torch.int32, torch.int64)
+    output = torch.empty(
+        locations.shape,
+        device=locations.device,
+        dtype=torch.promote_types(mapping.dtype, locations.dtype),
+    )
+    if locations.numel():
+        _translate_padded_hisparse_locations[(triton.cdiv(locations.numel(), 128),)](
+            mapping, locations, output, locations.numel(), locations.stride(0), 128
+        )
+    return output

python/sglang/srt/mem_cache/hisparse_memory_pool.py

@@ -5,6 +5,9 @@ from typing import Optional
 
 import torch
 
+from sglang.kernels.ops.kvcache.hisparse_slot_mapping import (
+    translate_padded_hisparse_locations,
+)
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
 from sglang.srt.utils import is_cuda, is_hip
@@ -74,7 +77,18 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
             full_to_hisparse_device_index_mapping
         )
 
-    def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
+    def translate_loc_to_hisparse_device(
+        self, compressed_indices: torch.Tensor
+    ) -> torch.Tensor:
+        """Map logical locations to physical slots with the same shape.
+
+        CUDA and ROCm use a fused kernel for 1D GPU slot lists, preserving
+        negative padding. Page tables and CPU inputs keep the direct gather.
+        """
+        if compressed_indices.is_cuda and compressed_indices.ndim == 1:
+            return translate_padded_hisparse_locations(
+                self.full_to_hisparse_device_index_mapping, compressed_indices
+            )
         return self.full_to_hisparse_device_index_mapping[compressed_indices]
 
     def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py

@@ -42,6 +42,7 @@ from sglang.srt.lora.deepseek_mla_correction import (
 from sglang.srt.lora.deepseek_mla_correction import (
     is_kv_b_lora_active,
 )
+from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
@@ -350,12 +351,18 @@ def _fused_rope_cat_and_cache(
         and attn.current_attention_backend == "aiter"
         else kv_cache_dtype
     )
+    kv_pool = get_token_to_kv_pool()
+    if isinstance(kv_pool, HiSparseDSATokenToKVPool):
+        # The fused write bypasses set_mla_kv_buffer()'s logical-to-device mapping.
+        out_cache_loc = kv_pool.translate_loc_to_hisparse_device(out_cache_loc)
+    # AITER reads slot_mapping with stride 1, including on the resident path.
+    out_cache_loc = out_cache_loc.contiguous()
     return fused_qk_rope_cat_and_cache_mla(
         q_nope_out,
         q_pe,
         k_nope,
         k_pe,
-        get_token_to_kv_pool().get_key_buffer(attn.attn_mqa.layer_id),
+        kv_pool.get_key_buffer(attn.attn_mqa.layer_id),
         out_cache_loc,
         positions,
         attn.rotary_emb.cos_cache,