[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,

test/registered/amd/test_rocm_hisparse_fused_kv_gpu.py

@@ -0,0 +1,136 @@
+"""Model-free AITER KV-write regression for MI35x, including graph replay."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=20, suite="stage-b-test-1-gpu-small-amd-mi35x")
+
+
+class TestRocmHiSparseFusedKVKernel(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        if not torch.version.hip or not torch.cuda.is_available():
+            raise AssertionError("The MI35x regression requires ROCm and an AMD GPU")
+        arch = torch.cuda.get_device_properties(0).gcnArchName
+        if not arch.startswith("gfx95"):
+            raise AssertionError("The MI35x regression requires gfx95 hardware")
+
+        from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
+        from sglang.srt.models.deepseek_common.attention_forward_methods import (
+            forward_mla_rocm,
+        )
+
+        if not forward_mla_rocm._use_aiter_gfx95:
+            raise AssertionError(
+                "This gfx95 regression must run with the fused MLA path; "
+                "set SGLANG_USE_AITER=1 before starting Python"
+            )
+
+        cls.forward = forward_mla_rocm
+        cls.pool_type = HiSparseDSATokenToKVPool
+
+    def setUp(self):
+        super().setUp()
+        # The owned guard catches out-of-range writes without corrupting other tensors.
+        self.backing = torch.full(
+            (24, 1, 576), -99.0, dtype=torch.bfloat16, device="cuda"
+        )
+        mapping = torch.zeros(65, dtype=torch.int64, device="cuda")
+        mapping[17], mapping[18] = 3, 5
+        mapping[-1] = -1
+        self.pool = self.pool_type.__new__(self.pool_type)
+        self.pool.register_mapping(mapping)
+        self.pool.kv_buffer = [self.backing[:8]]
+        self.pool.start_layer = 7
+        self.pool.layer_transfer_counter = None
+        self.pool.dtype = self.pool.store_dtype = torch.bfloat16
+        self.attn = SimpleNamespace(
+            kv_cache_dtype="bfloat16",
+            current_attention_backend="dsa",
+            attn_mqa=SimpleNamespace(layer_id=7, k_scale=torch.ones(1, device="cuda")),
+            rotary_emb=SimpleNamespace(
+                cos_cache=torch.ones((8, 64), dtype=torch.bfloat16, device="cuda"),
+                sin_cache=torch.zeros((8, 64), dtype=torch.bfloat16, device="cuda"),
+                is_neox_style=False,
+            ),
+        )
+        self.qn = torch.ones((3, 8, 512), dtype=torch.bfloat16, device="cuda")
+        self.qr = torch.ones((3, 8, 64), dtype=torch.bfloat16, device="cuda")
+        self.kn = (
+            torch.arange(1, 4, dtype=torch.bfloat16, device="cuda")[:, None, None]
+            .expand(3, 1, 512)
+            .contiguous()
+        )
+        self.kr = torch.full((3, 1, 64), 2.5, dtype=torch.bfloat16, device="cuda")
+        self.positions = torch.zeros(3, dtype=torch.int64, device="cuda")
+        self.locations = torch.tensor([17, 18, -1], device="cuda")
+        self.expected_rows = torch.cat((self.kn, self.kr), dim=-1)
+
+    def write(self):
+        return self.forward._fused_rope_cat_and_cache(
+            self.attn,
+            self.qn,
+            self.qr,
+            self.kn,
+            self.kr,
+            self.positions,
+            self.locations,
+        )
+
+    def check_cache(self, row_to_slot):
+        torch.cuda.synchronize()
+        expected = torch.full_like(self.backing, -99.0)
+        for row, slot in row_to_slot.items():
+            expected[slot] = self.expected_rows[row]
+        torch.testing.assert_close(self.backing, expected, rtol=0, atol=0)
+
+    def test_eager_and_graph_writes_use_physical_slots(self):
+        with patch.object(self.forward, "get_token_to_kv_pool", return_value=self.pool):
+            # The helper calls the real AITER kernel with the pool's device buffer.
+            self.write()
+            with self.subTest(mode="eager"):
+                self.check_cache({0: 3, 1: 5})
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                self.write()
+            self.backing.fill_(-99.0)
+            graph.replay()
+            with self.subTest(mode="graph"):
+                self.check_cache({0: 3, 1: 5})
+
+            # Inputs are mutable between graph replays, including padding.
+            self.locations.copy_(torch.tensor([18, -1, 17], device="cuda"))
+            self.backing.fill_(-99.0)
+            graph.replay()
+            with self.subTest(mode="graph-updated-inputs"):
+                self.check_cache({0: 5, 2: 3})
+
+            self.locations.copy_(torch.tensor([19, -1, 17], device="cuda"))
+            self.backing.fill_(-99.0)
+            graph.replay()
+            with self.subTest(mode="graph-unmapped-dummy-slot"):
+                self.check_cache({0: 0, 2: 3})
+
+    def test_resident_strided_locations(self):
+        """AITER must write selected slots rather than interleaved storage values."""
+        from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
+
+        resident = DSATokenToKVPool.__new__(DSATokenToKVPool)
+        resident.__dict__.update(self.pool.__dict__)
+        self.pool = resident
+        self.locations = torch.tensor([3, 21, 5, 22, -1, 23], device="cuda")[::2]
+        with patch.object(self.forward, "get_token_to_kv_pool", return_value=self.pool):
+            self.write()
+        with self.subTest(layout="strided"):
+            self.check_cache({0: 3, 1: 5})
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/kernels/ops/kvcache/test_hisparse_slot_mapping.py

@@ -0,0 +1,72 @@
+"""Portable HiSparse logical-to-physical slot translation kernel checks."""
+
+import unittest
+import weakref
+
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=20, stage="jit-kernel-unit", runner_config="amd")
+
+
+class TestHiSparseSlotMapping(CustomTestCase):
+    def test_accepts_weak_mapping_used_by_pool(self):
+        from sglang.kernels.ops.kvcache.hisparse_slot_mapping import (
+            translate_padded_hisparse_locations as translate,
+        )
+
+        mapping = torch.arange(32, dtype=torch.int64, device="cuda") + 100
+        locations = torch.tensor([17, -1, 18], device="cuda")
+
+        actual = translate(weakref.proxy(mapping), locations)
+
+        torch.testing.assert_close(actual, torch.tensor([117, -1, 118], device="cuda"))
+
+    def test_fused_slot_mapping_matches_padded_gather(self):
+        from sglang.kernels.ops.kvcache.hisparse_slot_mapping import (
+            translate_padded_hisparse_locations as translate,
+        )
+
+        for map_dtype in (torch.int32, torch.int64):
+            mapping = torch.arange(257, dtype=map_dtype, device="cuda") * 3 + 1
+            for loc_dtype in (torch.int32, torch.int64):
+                for count in (0, 1, 3, 127, 129, 1024):
+                    for stride in (1, 2):
+                        with self.subTest(
+                            map_dtype=map_dtype,
+                            loc_dtype=loc_dtype,
+                            count=count,
+                            stride=stride,
+                        ):
+                            locations = (
+                                torch.arange(
+                                    count * stride, dtype=loc_dtype, device="cuda"
+                                )[::stride]
+                                % 257
+                            )
+                            if stride == 2:
+                                storage = torch.zeros(
+                                    count * 2, dtype=loc_dtype, device="cuda"
+                                )
+                                storage[::2] = locations
+                                locations = storage[::2]
+                            locations[::3] = -1
+                            locations[1::7] = -2
+                            original = locations.clone()
+                            expected = torch.where(
+                                locations >= 0,
+                                mapping[locations.clamp_min(0)],
+                                locations,
+                            )
+                            actual = translate(mapping, locations)
+                            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+                            torch.testing.assert_close(
+                                locations, original, rtol=0, atol=0
+                            )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_hisparse_slot_translation.py

@@ -0,0 +1,68 @@
+"""Check portable HiSparse pool translation dispatch and gather fallbacks."""
+
+import unittest
+from unittest.mock import PropertyMock, patch
+
+import torch
+
+from sglang.srt.mem_cache import hisparse_memory_pool
+from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, stage="base-a", runner_config="cpu")
+
+
+class TestHiSparseSlotTranslation(CustomTestCase):
+    def make_pool(self, pool_type):
+        return pool_type.__new__(pool_type)
+
+    def test_pool_translation_selects_fused_kernel_for_gpu_slot_lists(self):
+        pool = self.make_pool(HiSparseDSATokenToKVPool)
+        mapping = torch.arange(32, dtype=torch.int64) + 100
+        pool.register_mapping(mapping)
+        storage = torch.tensor([17, 0, -1, 0, 18, 0])
+        locations = storage[::2]
+        expected = torch.tensor([117, -1, 118])
+        with (
+            patch.object(
+                torch.Tensor, "is_cuda", new_callable=PropertyMock, return_value=True
+            ),
+            patch.object(
+                hisparse_memory_pool,
+                "translate_padded_hisparse_locations",
+                return_value=expected,
+            ) as fused,
+        ):
+            self.assertIs(pool.translate_loc_to_hisparse_device(locations), expected)
+            self.assertIs(fused.call_args.args[0], mapping)
+            self.assertIs(fused.call_args.args[1], locations)
+            fused.assert_called_once()
+        torch.testing.assert_close(storage, torch.tensor([17, 0, -1, 0, 18, 0]))
+
+    def test_pool_translation_keeps_gather_for_page_tables_and_other_devices(self):
+        pool = self.make_pool(HiSparseDSATokenToKVPool)
+        mapping = torch.arange(32, dtype=torch.int64) + 100
+        pool.register_mapping(mapping)
+        for gpu, locations in (
+            (False, torch.tensor([17, -1, 18])),
+            (True, torch.tensor([[17, -1], [18, 0]])),
+            (True, torch.tensor(17)),
+        ):
+            with (
+                self.subTest(gpu=gpu, shape=locations.shape),
+                patch.object(
+                    torch.Tensor, "is_cuda", new_callable=PropertyMock, return_value=gpu
+                ),
+                patch.object(
+                    hisparse_memory_pool, "translate_padded_hisparse_locations"
+                ) as fused,
+            ):
+                actual = pool.translate_loc_to_hisparse_device(locations)
+                torch.testing.assert_close(actual, mapping[locations])
+                self.assertEqual(actual.shape, locations.shape)
+                fused.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/models/test_rocm_hisparse_fused_kv.py

@@ -0,0 +1,97 @@
+"""Exercise the ROCm writer's slot contract with a CPU kernel boundary."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
+from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
+from sglang.srt.models.deepseek_common.attention_forward_methods import forward_mla_rocm
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestRocmHiSparseFusedKV(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        self.cache = object()
+        self.pool = self.make_pool(DSATokenToKVPool)
+        self.attn = SimpleNamespace(
+            kv_cache_dtype="bfloat16",
+            current_attention_backend="dsa",
+            attn_mqa=SimpleNamespace(layer_id=7, k_scale=1.0),
+            rotary_emb=SimpleNamespace(
+                cos_cache=None, sin_cache=None, is_neox_style=False
+            ),
+        )
+
+    def make_pool(self, pool_type):
+        # Supply storage without allocating a model's cache; keep real accessors.
+        pool = pool_type.__new__(pool_type)
+        pool.kv_buffer = [self.cache]
+        pool.start_layer = 7
+        pool.layer_transfer_counter = None
+        pool.dtype = pool.store_dtype = torch.bfloat16
+        return pool
+
+    def use_hisparse(self):
+        self.pool = self.make_pool(HiSparseDSATokenToKVPool)
+        mapping = torch.zeros(65, dtype=torch.int64)
+        mapping[17], mapping[18], mapping[-1] = 3, 5, -1
+        self.pool.register_mapping(mapping)
+        return mapping
+
+    def invoke(self, locations):
+        with (
+            patch.object(forward_mla_rocm, "get_token_to_kv_pool", lambda: self.pool),
+            patch.object(
+                forward_mla_rocm,
+                "fused_qk_rope_cat_and_cache_mla",
+                lambda *args, **kwargs: args,
+                create=True,
+            ),
+        ):
+            args = forward_mla_rocm._fused_rope_cat_and_cache(
+                self.attn, torch.empty(0), None, None, None, None, locations
+            )
+        self.assertIs(args[4], self.cache)
+        return args[5]
+
+    def test_resident_locations_are_unchanged(self):
+        locations = torch.tensor([17, 0, -1], dtype=torch.int64)
+        self.assertIs(self.invoke(locations), locations)
+
+    def test_resident_strided_locations(self):
+        """Strided locations must not send interleaved storage values to AITER."""
+        storage = torch.tensor([3, 21, 5, 22, -1, 23])
+        locations = storage[::2]
+        actual = self.invoke(locations)
+        with self.subTest(layout="strided"):
+            self.assertTrue(actual.is_contiguous())
+            torch.testing.assert_close(actual, torch.tensor([3, 5, -1]))
+            torch.testing.assert_close(storage, torch.tensor([3, 21, 5, 22, -1, 23]))
+
+    def test_hisparse_maps_logical_slots(self):
+        """Logical slots beyond device capacity must write their physical rows."""
+        self.use_hisparse()
+        locations = torch.tensor([17, 18], dtype=torch.int64)
+        with self.subTest(logical_slots=[17, 18]):
+            torch.testing.assert_close(self.invoke(locations), torch.tensor([3, 5]))
+            torch.testing.assert_close(locations, torch.tensor([17, 18]))
+
+    def test_hisparse_padding_and_unmapped_slots(self):
+        self.use_hisparse()
+        locations = torch.tensor([17, -1, 19, 0])
+        with self.subTest(padding=-1, unmapped=19):
+            torch.testing.assert_close(
+                self.invoke(locations), torch.tensor([3, -1, 0, 0])
+            )
+            torch.testing.assert_close(locations, torch.tensor([17, -1, 19, 0]))
+
+
+if __name__ == "__main__":
+    unittest.main()