[2/N] [Kernel] Fuse padding-preserving HiSparse slot translation (#39837)
This commit is contained in:
Sasha Sidorov
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user