minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Enable HiSparse on ROCm (#26639)

Browse Source 
Co-authored-by: clintg6 <7388379+clintg6@users.noreply.github.com>
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Clint
2026-06-19 11:59:45 -07:00
committed by GitHub
co-authored by clintg6 HAI
parent ca88b7f1d2
commit c436a8161a
12 changed files with 971 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/amd/accuracy/mi30x/test_glm51_hisparse_eval_mi30x.py
+90
View File
@@ -0,0 +1,90 @@
"""AMD GLM-5.1 HiSparse GSM8K evaluation test (8-GPU MI30x)."""
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_amd_ci(
est_time=3600,
suite="nightly-amd-accuracy-8-gpu-glm51-hisparse",
nightly=True,
)
class TestGLM51HiSparseEvalAMD(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = "/models/GLM-5.1-FP8/"
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=7200,
other_args=[
"--tp",
"8",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 8}',
"--trust-remote-code",
"--tool-call-parser",
"glm47",
"--reasoning-parser",
"glm45",
"--mem-fraction-static",
"0.65",
"--dsa-prefill-backend",
"aiter",
"--dsa-decode-backend",
"aiter",
"--kv-cache-dtype",
"fp8_e4m3",
"--max-running-requests",
"2",
"--watchdog-timeout",
"1200",
"--skip-server-warmup",
"--enable-hisparse",
"--hisparse-config",
'{"top_k": 2048, "device_buffer_size": 2048, "host_to_device_ratio": 1}',
"--disable-radix-cache",
],
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process"):
kill_process_tree(cls.process.pid)
def test_gsm8k_accuracy(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=4000,
num_examples=500,
num_threads=100,
num_shots=24,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (glm-5.1 hisparse mi30x)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.93)
if __name__ == "__main__":
unittest.main()
test/registered/amd/accuracy/mi35x/test_glm51_hisparse_eval_mi35x.py
+90
View File
@@ -0,0 +1,90 @@
"""MI35x GLM-5.1 HiSparse GSM8K evaluation test (8-GPU)."""
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_amd_ci(
est_time=5400,
suite="nightly-amd-8-gpu-mi35x-glm51-hisparse",
nightly=True,
)
class TestGLM51HiSparseEvalMI35x(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.model = "/models/GLM-5.1-FP8/"
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=7200,
other_args=[
"--tp",
"8",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 8}',
"--trust-remote-code",
"--tool-call-parser",
"glm47",
"--reasoning-parser",
"glm45",
"--mem-fraction-static",
"0.65",
"--dsa-prefill-backend",
"aiter",
"--dsa-decode-backend",
"aiter",
"--kv-cache-dtype",
"fp8_e4m3",
"--max-running-requests",
"2",
"--watchdog-timeout",
"1200",
"--skip-server-warmup",
"--enable-hisparse",
"--hisparse-config",
'{"top_k": 2048, "device_buffer_size": 2048, "host_to_device_ratio": 1}',
"--disable-radix-cache",
],
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process"):
kill_process_tree(cls.process.pid)
def test_gsm8k_accuracy(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=4000,
num_examples=500,
num_threads=100,
num_shots=24,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (glm-5.1 hisparse mi35x)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.93)
if __name__ == "__main__":
unittest.main()
test/registered/jit/test_hisparse.py
+152 -1
View File
@@ -9,8 +9,9 @@ from sglang.jit_kernel.hisparse import (
transfer_cache_dsv4_mla,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd")
register_cuda_ci(est_time=10, suite="base-b-kernel-unit-1-gpu-large")
register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
@@ -366,6 +367,33 @@ def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
def test_load_cache_to_device_buffer_multiple_misses_copy_all_slots() -> None:
state = _make_state(
[[9, 7, 3, 5, 11]],
[[0, 1, 2, 3, -1]],
[8],
)
out = _run_kernel(
top_k_tokens=torch.tensor([[4, 5, 6, 7]], dtype=torch.int32, device=DEVICE),
seq_len=9,
**state,
)
assert torch.equal(out.cpu(), torch.tensor([[9, 7, 3, 5]], dtype=torch.int32))
assert torch.equal(
state["device_buffer_tokens"].cpu(),
torch.tensor([[4, 5, 6, 7, -1]], dtype=torch.int32),
)
assert torch.equal(
state["lru_slots"].cpu(), torch.tensor([[0, 1, 2, 3]], dtype=torch.int16)
)
for token, loc in zip([4, 5, 6, 7], [9, 7, 3, 5]):
assert torch.equal(
state["device_buffer"][loc].cpu(), state["host_cache"][token]
)
def test_load_cache_to_device_buffer_batched_with_padding() -> None:
state = _make_state(
[
@@ -422,5 +450,128 @@ def test_load_cache_to_device_buffer_batched_with_padding() -> None:
assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
def test_load_cache_to_device_buffer_dsv4_mla_miss_copy_layout() -> None:
# Both the host cache and the device buffer use the page-padded C4 layout,
# matching DeepSeekV4PagedHostPool, the backup/write path, and the swap-in
# kernel on both CUDA and ROCm. The miss copy must read the host source with
# paged addressing (get_pointer_paged), not a linear per-item stride.
num_pages = (HOST_CACHE_SIZE + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE
state = _long_case()
host_cache = torch.zeros(
(num_pages, DSV4_PAGE_BYTES),
dtype=torch.uint8,
device="cpu",
pin_memory=True,
)
for token in range(HOST_CACHE_SIZE):
_write_dsv4_token(host_cache, token, seed=token + 1)
device_buffer = torch.full(
(num_pages, DSV4_PAGE_BYTES),
0xFF,
dtype=torch.uint8,
device=DEVICE,
)
out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE)
# Token 6 is a miss in _long_case(), so it should be copied into evict slot 0,
# whose physical device loc is 9.
load_cache_to_device_buffer_dsv4_mla(
top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE),
device_buffer_tokens=state["device_buffer_tokens"],
host_cache_locs=state["host_cache_locs"],
device_buffer_locs=state["device_buffer_locs"],
host_cache=host_cache,
device_buffer=device_buffer,
top_k_device_locs=out,
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE),
lru_slots=state["lru_slots"],
item_size_bytes=DSV4_ITEM_BYTES,
num_top_k=1,
hot_buffer_size=HOT_BUFFER_SIZE,
page_size=DSV4_PAGE_SIZE,
block_size=256,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32))
# host_cache_locs[token=6] == 6 in _long_case(); evict slot 0 -> device loc 9.
assert torch.equal(
_read_dsv4_token(device_buffer, 9).cpu(),
_read_dsv4_token(host_cache, 6),
)
@pytest.mark.skipif(
not is_hip(), reason="Covers a ROCm wavefront64 LRU writeback regression."
)
def test_load_cache_to_device_buffer_rocm_large_lru_writeback() -> None:
top_k = 2048
hot_buffer_size = 4096
seq_len = 7299
kv_dim = 4
item_size_bytes = kv_dim * torch.empty((), dtype=DTYPE).element_size()
top_k_tokens = torch.cat(
[
torch.arange(1000, 2000, dtype=torch.int32),
torch.arange(5000, 6048, dtype=torch.int32),
]
).view(1, -1)
device_buffer_tokens = torch.arange(hot_buffer_size, dtype=torch.int32).view(1, -1)
device_buffer_locs = torch.arange(hot_buffer_size + 1, dtype=torch.int32).view(
1, -1
)
lru_slots = torch.arange(hot_buffer_size, dtype=torch.int16).view(1, -1)
host_cache_locs = torch.arange(seq_len, dtype=torch.int64).view(1, -1)
top_k_tokens = top_k_tokens.to(DEVICE)
device_buffer_tokens = device_buffer_tokens.to(DEVICE)
device_buffer_locs = device_buffer_locs.to(DEVICE)
lru_slots = lru_slots.to(DEVICE)
host_cache_locs = host_cache_locs.to(DEVICE)
host_cache = torch.empty((seq_len, 1, kv_dim), dtype=DTYPE, pin_memory=True)
host_cache.zero_()
device_buffer = torch.empty(
(hot_buffer_size + 1, 1, kv_dim), dtype=DTYPE, device=DEVICE
)
out = torch.full_like(top_k_tokens, -1)
load_cache_to_device_buffer_mla(
top_k_tokens=top_k_tokens,
device_buffer_tokens=device_buffer_tokens,
host_cache_locs=host_cache_locs,
device_buffer_locs=device_buffer_locs,
host_cache=host_cache,
device_buffer=device_buffer,
top_k_device_locs=out,
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE),
lru_slots=lru_slots,
item_size_bytes=item_size_bytes,
num_top_k=top_k,
hot_buffer_size=hot_buffer_size,
page_size=1,
block_size=1024,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
torch.cuda.synchronize()
expected_lru = torch.cat(
[
torch.arange(2048, 4096, dtype=torch.int16),
torch.arange(0, 1000, dtype=torch.int16),
torch.arange(2000, 2048, dtype=torch.int16),
torch.arange(1000, 2000, dtype=torch.int16),
]
)
assert torch.equal(lru_slots.cpu().view(-1), expected_lru)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v", "-s"]))
test/registered/unit/managers/test_hisparse_unit.py
+92 -4
View File
@@ -15,9 +15,10 @@ from types import SimpleNamespace
import torch
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
# ---------------------------------------------------------------------------
# Test configuration (small-scale for fast CI runs)
@@ -92,7 +93,14 @@ class TestHiSparseUnit(unittest.TestCase):
cls._original_alloc = ALLOC_MEMORY_FUNCS["cuda"]
ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory
global_page_size = 1 if is_hip() else PAGE_SIZE
if is_hip():
from sglang.srt.layers.attention.dsa.utils import (
aiter_can_use_preshuffle_paged_mqa,
)
global_page_size = 64 if aiter_can_use_preshuffle_paged_mqa() else 1
else:
global_page_size = PAGE_SIZE
from sglang.srt.mem_cache.allocator.hisparse import (
HiSparseTokenToKVPoolAllocator,
@@ -524,6 +532,80 @@ class TestHiSparseUnit(unittest.TestCase):
self.allocator.logical_attn_allocator.free(kv_loc)
self._assert_sizes_restored(initial, "alloc_free_cycle")
def test_allocator_page_size_one_alloc_free_cycle(self):
"""alloc() maps logical to hisparse indices for ROCm page_size=1."""
if self.page_size != 1:
self.skipTest("page_size=1 alloc path is ROCm-specific")
initial = self._get_initial_sizes()
need_size = 16
kv_loc = self.allocator.alloc(need_size)
self.assertIsNotNone(kv_loc)
self.assertEqual(len(kv_loc), need_size)
mapping = self.allocator.full_to_hisparse_device_index_mapping[kv_loc]
self.assertTrue(torch.all(mapping > 0), "Mapping should be non-zero")
self.assertLess(self.allocator.available_size(), initial[0])
self.allocator.free(kv_loc)
mapping_after = self.allocator.full_to_hisparse_device_index_mapping[kv_loc]
self.assertTrue(torch.all(mapping_after == 0), "Mapping should be cleared")
self._assert_sizes_restored(initial, "page_size_one_alloc_free_cycle")
def test_decode_remap_frees_stale_page_size_one_mapping(self):
"""map_last_loc_to_buffer frees the temporary alloc() hisparse slot."""
if self.page_size != 1:
self.skipTest("page_size=1 decode remap path is ROCm-specific")
initial = self._get_initial_sizes()
device = self.allocator.device
fill_len = 2
req = _make_req("decode-remap", list(range(fill_len)))
self._alloc_req_slot(req)
kv_loc = self._alloc_kv(req, fill_len)
self.coordinator.alloc_device_buffer(req)
self.coordinator._skip_first_backup[req.req_pool_idx] = True
out_loc = self.allocator.alloc(1)
self.assertIsNotNone(out_loc)
stale_loc = self.allocator.full_to_hisparse_device_index_mapping[
out_loc
].clone()
self.assertTrue(torch.all(stale_loc > 0), "Temporary mapping should exist")
seq_len = fill_len + 1
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc)
req.kv_allocated_len = seq_len
req.kv_committed_len = seq_len
self.coordinator.map_last_loc_to_buffer(
seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device),
out_cache_loc=out_loc,
req_pool_indices=torch.tensor(
[req.req_pool_idx], dtype=torch.int64, device=device
),
seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64),
req_pool_indices_cpu=torch.tensor([req.req_pool_idx], dtype=torch.int64),
)
remapped_loc = self.allocator.full_to_hisparse_device_index_mapping[out_loc]
self.assertTrue(torch.all(remapped_loc > 0), "Remapped loc should exist")
self.assertFalse(
torch.equal(stale_loc, remapped_loc),
"Decode loc should move from temporary mapping to device buffer",
)
self.assertEqual(
self.allocator.hisparse_attn_allocator.available_size(),
initial[1] - seq_len,
)
self.coordinator.request_finished(req)
self.allocator.logical_attn_allocator.free(torch.cat([kv_loc, out_loc]))
self._free_req_slot(req)
self._assert_sizes_restored(initial, "decode_remap")
# ==================================================================
# Test: Staging (PD Colocate) path
# ==================================================================
@@ -596,17 +678,23 @@ class TestHiSparseUnit(unittest.TestCase):
fill_len,
1,
)
# With page_size>1 the rounded-up staging allocation provides headroom,
# so no new pages are needed. With page_size=1 there is no headroom and
# exactly one new page is allocated for the next token.
expected_new_pages = 0 if fill_len < rounded_len else 1
self.assertEqual(
self.coordinator.mem_pool_host.available_size(), available_size
self.coordinator.mem_pool_host.available_size(),
available_size - expected_new_pages,
)
self.assertTrue(torch.all(next_host_index >= 0))
expected_total = rounded_len + expected_new_pages * self.page_size
allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices(
self.coordinator.req_to_host_pool,
req.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]),
)
self.assertEqual(allocated_host_indices.numel(), rounded_len)
self.assertEqual(allocated_host_indices.numel(), expected_total)
self._cleanup_req(req, kv_loc)
self._assert_sizes_restored(initial, "single_node_staging_pages")
test/registered/unit/server_args/test_server_args.py
+109
View File
@@ -147,6 +147,115 @@ class TestLoadBalanceMethod(unittest.TestCase):
self.assertEqual(server_args.disaggregation_transfer_backend, "mooncake")
class TestHiSparseDsaBackendPolicy(unittest.TestCase):
@patch("sglang.srt.server_args.is_hip", return_value=False)
def test_hisparse_defaults_to_flashmla_sparse_on_cuda_bfloat16(self, _mock_is_hip):
server_args = ServerArgs(model_path="dummy", enable_hisparse=True)
server_args._set_default_dsa_backends(kv_cache_dtype="bfloat16", major=9)
self.assertEqual(server_args.dsa_prefill_backend, "flashmla_sparse")
self.assertEqual(server_args.dsa_decode_backend, "flashmla_sparse")
@patch("sglang.srt.server_args.is_hip", return_value=False)
def test_hisparse_defaults_to_flashmla_kv_on_cuda_fp8(self, _mock_is_hip):
server_args = ServerArgs(model_path="dummy", enable_hisparse=True)
server_args._set_default_dsa_backends(kv_cache_dtype="fp8_e4m3", major=9)
self.assertEqual(server_args.dsa_prefill_backend, "flashmla_kv")
self.assertEqual(server_args.dsa_decode_backend, "flashmla_kv")
@patch("sglang.srt.server_args.is_hip", return_value=True)
def test_hisparse_defaults_to_tilelang_on_rocm(self, _mock_is_hip):
server_args = ServerArgs(model_path="dummy", enable_hisparse=True)
server_args._set_default_dsa_backends(kv_cache_dtype="bfloat16", major=9)
self.assertEqual(server_args.dsa_prefill_backend, "tilelang")
self.assertEqual(server_args.dsa_decode_backend, "tilelang")
@patch("sglang.srt.server_args.is_hip", return_value=True)
def test_hisparse_preserves_rocm_user_backend_and_defaults_missing_side(
self, _mock_is_hip
):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
dsa_prefill_backend="tilelang",
)
server_args._set_default_dsa_backends(kv_cache_dtype="bfloat16", major=9)
self.assertEqual(server_args.dsa_prefill_backend, "tilelang")
self.assertEqual(server_args.dsa_decode_backend, "tilelang")
@patch("sglang.srt.server_args.is_hip", return_value=True)
def test_hisparse_accepts_aiter_backend_on_rocm(self, _mock_is_hip):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
kv_cache_dtype="bfloat16",
dsa_prefill_backend="aiter",
dsa_decode_backend="aiter",
)
server_args._validate_hisparse_dsa_backend("dsa_prefill_backend", "prefill")
server_args._validate_hisparse_dsa_backend("dsa_decode_backend", "decode")
@patch("sglang.srt.server_args.is_hip", return_value=True)
def test_hisparse_rejects_cuda_backend_on_rocm(self, _mock_is_hip):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
kv_cache_dtype="bfloat16",
dsa_prefill_backend="flashmla_sparse",
)
with self.assertRaisesRegex(ValueError, "tilelang"):
server_args._validate_hisparse_dsa_backend("dsa_prefill_backend", "prefill")
@patch("sglang.srt.server_args.is_hip", return_value=False)
def test_hisparse_rejects_rocm_backend_on_cuda(self, _mock_is_hip):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
kv_cache_dtype="bfloat16",
dsa_decode_backend="tilelang",
)
with self.assertRaisesRegex(ValueError, "flashmla_sparse"):
server_args._validate_hisparse_dsa_backend("dsa_decode_backend", "decode")
def test_hisparse_accepts_bfloat16_kv_cache_dtype(self):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
kv_cache_dtype="bfloat16",
)
server_args._validate_hisparse_kv_cache_dtype()
def test_hisparse_accepts_fp8_e4m3_kv_cache_dtype(self):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
kv_cache_dtype="fp8_e4m3",
)
server_args._validate_hisparse_kv_cache_dtype()
def test_hisparse_rejects_unsupported_kv_cache_dtype(self):
server_args = ServerArgs(
model_path="dummy",
enable_hisparse=True,
kv_cache_dtype="float16",
)
with self.assertRaisesRegex(ValueError, r"fp8_e4m3"):
server_args._validate_hisparse_kv_cache_dtype()
class TestContextParallelServerArgs(CustomTestCase):
def setUp(self):
self.parser = server_args_module.argparse.ArgumentParser()