diff --git a/python/sglang/kernels/ops/kvcache/hisparse_slot_mapping.py b/python/sglang/kernels/ops/kvcache/hisparse_slot_mapping.py new file mode 100644 index 000000000..f1609d552 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/mem_cache/hisparse_memory_pool.py b/python/sglang/srt/mem_cache/hisparse_memory_pool.py index 4808dadd3..67b5e2d74 100644 --- a/python/sglang/srt/mem_cache/hisparse_memory_pool.py +++ b/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): diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 91b3bc929..28ee0c1f0 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/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, diff --git a/test/registered/amd/test_rocm_hisparse_fused_kv_gpu.py b/test/registered/amd/test_rocm_hisparse_fused_kv_gpu.py new file mode 100644 index 000000000..e3ac6c9c2 --- /dev/null +++ b/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() diff --git a/test/registered/kernels/ops/kvcache/test_hisparse_slot_mapping.py b/test/registered/kernels/ops/kvcache/test_hisparse_slot_mapping.py new file mode 100644 index 000000000..93ca37657 --- /dev/null +++ b/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() diff --git a/test/registered/unit/mem_cache/test_hisparse_slot_translation.py b/test/registered/unit/mem_cache/test_hisparse_slot_translation.py new file mode 100644 index 000000000..b0833784d --- /dev/null +++ b/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() diff --git a/test/registered/unit/models/test_rocm_hisparse_fused_kv.py b/test/registered/unit/models/test_rocm_hisparse_fused_kv.py new file mode 100644 index 000000000..e5c49e652 --- /dev/null +++ b/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()