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

[DSV4] Chunk the indexer MQA logits by query rows under a free-memory budget (#39095)

Browse Source 
Signed-off-by: Shiki Wu <shikiw@nvidia.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
This commit is contained in:
Junpan Wu
2026-09-19 11:53:20 -07:00
committed by GitHub
co-authored by Baizhou Zhang Yuwei An
parent 993d1fccba
commit 7b67a96640
6 changed files with 720 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/attention/test_dsa_mqa_logits_chunking.py
+13 -16
View File
@@ -1,4 +1,5 @@
"""Contract tests for the DSA indexer's MQA-logits chunk budget.
"""Contract tests for the MQA-logits chunk decision shared by the DSA and DSV4
indexers.
On ROCm the `[num_q x num_k]` fp32 logits tensor goes to aiter's
`fp8_mqa_logits`, which only compiles below 2 GiB, so the budget that decides
@@ -8,34 +9,30 @@ The measured memory budget is stubbed: it is the only input the limit has to
beat, and stubbing it keeps these tests on CPU.
"""
from unittest import mock
import pytest
torch = pytest.importorskip("torch")
from sglang.srt.layers.attention.dsa import dsa_indexer # noqa: E402
from sglang.srt.layers.attention.mqa_logits_utils import ( # noqa: E402
MQA_LOGITS_MAX_BYTES_ROCM,
mqa_logits_should_chunk,
)
from sglang.test.ci.ci_register import register_cpu_ci # noqa: E402
register_cpu_ci(est_time=9, suite="base-a-test-cpu")
CEILING = dsa_indexer.Indexer._MQA_LOGITS_MAX_BYTES_ROCM
CEILING = MQA_LOGITS_MAX_BYTES_ROCM
# More than any single logits tensor here needs, so it never decides a case.
HUGE_MEM_BUDGET = 64 * 2**30
def _decide(num_q, num_k, mem_budget=HUGE_MEM_BUDGET, is_hip=True):
# __new__ skips an __init__ that needs a model config and a device.
indexer = dsa_indexer.Indexer.__new__(dsa_indexer.Indexer)
with (
mock.patch.object(dsa_indexer, "_is_hip", is_hip),
mock.patch.object(
dsa_indexer.Indexer,
"_get_mqa_logits_budget_bytes",
return_value=mem_budget,
),
):
return indexer._should_chunk_mqa_logits(num_q, num_k, 0)
return mqa_logits_should_chunk(
num_rows=num_q,
num_cols=num_k,
get_budget_bytes=lambda: mem_budget,
rocm=is_hip,
)
def test_the_ceiling_is_the_largest_logits_aiter_still_takes():
test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+353 -9
View File
@@ -11,19 +11,35 @@ from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
from sglang.srt.layers.attention.dsv4.indexer import (
FP8_DTYPE,
C4IndexerBackendMixin,
topk_transform_pytorch_vectorized,
)
from sglang.srt.layers.attention.dsv4.metadata import (
NonPagedIndexerPlan,
PagedIndexerMetadata,
iter_row_chunks,
plan_indexer_row_chunks,
)
from sglang.srt.layers.attention.mqa_logits_utils import (
MQA_LOGITS_MAX_BYTES_ROCM,
mqa_logits_budget_bytes,
mqa_logits_row_bytes,
mqa_logits_rows_per_chunk,
mqa_logits_should_chunk,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.runtime_context import get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
register_cpu_ci(est_time=17, suite="base-a-test-cpu")
_INDEXER = "sglang.srt.layers.attention.dsv4.indexer"
_METADATA = "sglang.srt.layers.attention.dsv4.metadata"
_MQA_UTILS = "sglang.srt.layers.attention.mqa_logits_utils"
# issue #35201: 372K raw tokens -> 92992 c4 columns, padded to 93184 by DeepGEMM.
_ISSUE_C4_COLS = 92992
_ISSUE_ALIGNED_COLS = 93184
class TestDSV4PagedIndexerMetadata(CustomTestCase):
@@ -247,8 +263,11 @@ class TestDSV4TopKDispatch(CustomTestCase):
torch.empty((1, 1, 1)),
)
)
backend._get_nonpaged_indexer_plan = MagicMock(return_value=object())
backend._forward_nonpaged_indexer = MagicMock(return_value=logits)
backend._get_nonpaged_indexer_plan = MagicMock(
return_value=SimpleNamespace(query_rows=1, rows_per_chunk=None)
)
backend._gather_nonpaged_index_k = MagicMock(return_value=(object(), object()))
backend._nonpaged_mqa_logits = MagicMock(return_value=logits)
indexer_capturer = MagicMock()
with (
@@ -353,7 +372,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
extend_start_loc=torch.tensor([0], dtype=torch.int32),
extend_num_tokens=query_rows,
)
metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64)
metadata = SimpleNamespace(
nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None
)
page_table = torch.tensor([[3, 1]], dtype=torch.int32).repeat(query_rows, 1)
c4_seq_lens = torch.tensor([62, 63, 64, 65], dtype=torch.int32)
@@ -388,6 +409,50 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
with threshold.override(0):
self.assertIsNone(build_plan())
def test_plan_row_chunking_follows_the_forward_budget(self):
backend = SimpleNamespace(_can_use_nonpaged_indexer=lambda **_: True)
backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True)
c4_indexer = SimpleNamespace(use_fp4_indexer=False, index_topk=512)
query_rows = 8192
seq_len = 372_000
batch = SimpleNamespace(
seq_lens=torch.tensor([seq_len], dtype=torch.int32),
seq_lens_cpu=[seq_len],
extend_seq_lens_cpu=[query_rows],
extend_seq_lens=torch.tensor([query_rows], dtype=torch.int32),
extend_start_loc=torch.tensor([0], dtype=torch.int32),
extend_num_tokens=query_rows,
)
c4_seq_lens = torch.full((query_rows,), seq_len // 4, dtype=torch.int32)
page_table = torch.zeros((query_rows, 1), dtype=torch.int32)
def build_plan(budget):
metadata = SimpleNamespace(
nonpaged_plan=None,
compressed_page_size=64,
mqa_logits_budget_bytes=budget,
)
return C4IndexerBackendMixin._get_nonpaged_indexer_plan(
backend,
c4_indexer=c4_indexer,
forward_batch=batch,
indexer_metadata=metadata,
page_table=page_table,
c4_seq_lens=c4_seq_lens,
query_rows=query_rows,
)
# No budget measured this forward (small batch or graph): one call.
self.assertIsNone(build_plan(None).rows_per_chunk)
budget = 512 << 20
plan = build_plan(budget)
# 8192 rows x align256(93056) cols x 4 B is far over 512 MiB.
self.assertIsNotNone(plan.rows_per_chunk)
self.assertLess(plan.rows_per_chunk, query_rows)
self.assertLessEqual(
plan.rows_per_chunk * mqa_logits_row_bytes(plan.max_seqlen_k), budget
)
def test_extreme_plan_metadata_is_bounded_and_fail_closed(self):
backend = SimpleNamespace(_can_use_nonpaged_indexer=lambda **_: True)
backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True)
@@ -401,7 +466,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
extend_start_loc=torch.tensor([0], dtype=torch.int32),
extend_num_tokens=query_rows,
)
metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64)
metadata = SimpleNamespace(
nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None
)
page_table = torch.zeros((query_rows, 1), dtype=torch.int32)
c4_seq_lens = torch.tensor(
[124_997, 124_998, 124_999, 125_000], dtype=torch.int32
@@ -439,7 +506,9 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
backend = SimpleNamespace(_can_use_nonpaged_indexer=can_use_nonpaged_indexer)
backend.dsa_topk_backend = SimpleNamespace(is_sgl_kernel=lambda: True)
c4_indexer = SimpleNamespace(use_fp4_indexer=False, index_topk=512)
metadata = SimpleNamespace(nonpaged_plan=None, compressed_page_size=64)
metadata = SimpleNamespace(
nonpaged_plan=None, compressed_page_size=64, mqa_logits_budget_bytes=None
)
def build_plan(query_rows):
batch = SimpleNamespace(
@@ -503,13 +572,18 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
deep_gemm = SimpleNamespace(fp8_mqa_logits=MagicMock(return_value=expected))
with patch.dict(sys.modules, {"deep_gemm": deep_gemm}):
actual = C4IndexerBackendMixin._forward_nonpaged_indexer(
q_indexer=q_indexer,
weights=weights,
kv = C4IndexerBackendMixin._gather_nonpaged_index_k(
c4_indexer=c4_indexer,
token_to_kv_pool=token_to_kv_pool,
plan=plan,
)
actual = C4IndexerBackendMixin._nonpaged_mqa_logits(
q_indexer=q_indexer,
weights=weights,
kv=kv,
plan=plan,
rows=slice(0, plan.query_rows),
)
self.assertIs(actual, expected)
token_to_kv_pool.get_index_k_scale_buffer.assert_called_once_with(
@@ -531,6 +605,276 @@ class TestDSV4NonPagedIndexer(CustomTestCase):
self.assertEqual(call.kwargs, {"clean_logits": False, "max_seqlen_k": 128})
class TestMqaLogitsBudgetArithmetic(CustomTestCase):
def test_row_bytes_follow_deepgemm_stride_alignment(self):
# DeepGEMM pads the fp32 logits row stride to 1024 B, i.e. 256 columns.
self.assertEqual(mqa_logits_row_bytes(1), 256 * 4)
self.assertEqual(mqa_logits_row_bytes(256), 256 * 4)
self.assertEqual(mqa_logits_row_bytes(257), 512 * 4)
self.assertEqual(mqa_logits_row_bytes(_ISSUE_C4_COLS), _ISSUE_ALIGNED_COLS * 4)
def test_rows_per_chunk_keeps_one_chunk_inside_budget(self):
row_bytes = mqa_logits_row_bytes(_ISSUE_C4_COLS)
budget = 512 << 20
# 4096 rows x 93184 cols x 4 B = 1.42 GiB > 512 MiB: must slice.
rows = mqa_logits_rows_per_chunk(
num_rows=4096, row_bytes=row_bytes, budget_bytes=budget
)
self.assertIsNotNone(rows)
self.assertLess(rows, 4096)
self.assertLessEqual(rows * row_bytes, budget)
# The whole matrix fits: single call.
self.assertIsNone(
mqa_logits_rows_per_chunk(
num_rows=64, row_bytes=row_bytes, budget_bytes=budget
)
)
# A tight budget never yields a chunk larger than the budget, even when
# that means fewer rows than a full page of queries.
tight = 64 << 20
rows = mqa_logits_rows_per_chunk(
num_rows=4096, row_bytes=row_bytes, budget_bytes=tight
)
self.assertEqual(rows, tight // row_bytes)
self.assertLessEqual(rows * row_bytes, tight)
# Few query rows still chunk when they do not fit.
self.assertEqual(
mqa_logits_rows_per_chunk(
num_rows=100, row_bytes=row_bytes, budget_bytes=40 * row_bytes
),
40,
)
# A budget below one row degrades to single-row chunks, never None.
self.assertEqual(
mqa_logits_rows_per_chunk(
num_rows=4096, row_bytes=row_bytes, budget_bytes=1
),
1,
)
def test_should_chunk_caps_the_budget_only_on_rocm(self):
huge = 64 << 30
# 16384 x 32768 x 4 B is exactly 2 GiB, aiter's compile-time ceiling.
self.assertEqual(
mqa_logits_should_chunk(
num_rows=16384, num_cols=32768, get_budget_bytes=lambda: huge, rocm=True
),
(True, MQA_LOGITS_MAX_BYTES_ROCM),
)
self.assertEqual(
mqa_logits_should_chunk(
num_rows=16384,
num_cols=32768,
get_budget_bytes=lambda: huge,
rocm=False,
),
(False, huge),
)
def test_should_chunk_skips_small_matrices_without_querying_the_budget(self):
get_budget = MagicMock(return_value=1)
# 64 decode rows x 100K columns is far below the 8M-element threshold.
self.assertEqual(
mqa_logits_should_chunk(
num_rows=64, num_cols=100_000, get_budget_bytes=get_budget, rocm=False
),
(False, 0),
)
get_budget.assert_not_called()
def test_plan_combines_sm120_cap_with_budget(self):
budget = 512 << 20
by_budget = mqa_logits_rows_per_chunk(
num_rows=8192,
row_bytes=mqa_logits_row_bytes(_ISSUE_C4_COLS),
budget_bytes=budget,
)
cases = (
# (num_rows, num_cols, budget_bytes, sm120_row_cap) -> rows_per_chunk
((8192, 1024, None, None), None),
((8192, 1024, None, 4096), 4096),
((4096, 1024, None, 4096), None),
((8192, _ISSUE_C4_COLS, budget, None), by_budget),
((8192, _ISSUE_C4_COLS, budget, 4096), min(4096, by_budget)),
)
for (num_rows, num_cols, budget_bytes, cap), expected in cases:
with self.subTest(num_rows=num_rows, num_cols=num_cols, cap=cap):
self.assertEqual(
plan_indexer_row_chunks(
num_rows=num_rows,
num_cols=num_cols,
budget_bytes=budget_bytes,
sm120_row_cap=cap,
),
expected,
)
def test_iter_row_chunks_covers_rows_exactly_once(self):
self.assertEqual(
list(iter_row_chunks(num_rows=10, rows_per_chunk=4)),
[slice(0, 4), slice(4, 8), slice(8, 10)],
)
self.assertEqual(
list(iter_row_chunks(num_rows=10, rows_per_chunk=None)), [slice(0, 10)]
)
self.assertEqual(
list(iter_row_chunks(num_rows=10, rows_per_chunk=64)), [slice(0, 10)]
)
def test_static_budget_never_queries_free_memory(self):
total = 80 << 30
props = SimpleNamespace(total_memory=total)
device_module = SimpleNamespace(
get_device_properties=MagicMock(return_value=props)
)
schedule = SimpleNamespace(mem_fraction_static=0.9)
with (
envs.SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION.override(0.2),
patch(f"{_MQA_UTILS}.get_device_module", return_value=device_module),
patch(f"{_MQA_UTILS}.get_schedule", return_value=schedule),
patch(f"{_MQA_UTILS}.is_hip", return_value=False),
patch(f"{_MQA_UTILS}.is_xpu", return_value=False),
patch(
"torch.cuda.mem_get_info", return_value=(6 << 30, total)
) as mem_get_info,
):
static = mqa_logits_budget_bytes(device_index=0, allow_sync=False)
mem_get_info.assert_not_called()
live = mqa_logits_budget_bytes(device_index=0, allow_sync=True)
mem_get_info.assert_called_once_with(0)
# static: 80 GiB x (1 - 0.9) x 0.2; live is further capped by 6 GiB free x 0.2.
self.assertEqual(static, int(int(total * 0.1) * 0.2))
self.assertEqual(live, int((6 << 30) * 0.2))
class TestPagedIndexerMetadataChunking(CustomTestCase):
"""The schedule list and the top-k plan list must be built over the exact
row chunks the indexer loops over; a mismatch would silently score rows
with another chunk's schedule."""
def _build(self, *, num_rows: int, budget, use_topk_v2: bool):
deep_gemm = SimpleNamespace(
get_num_sms=MagicMock(return_value=1),
get_paged_mqa_logits_metadata=MagicMock(
side_effect=lambda c4, *_: torch.zeros((2, 2), dtype=torch.int32)
),
)
c4_seq_lens = torch.arange(1, num_rows + 1, dtype=torch.int32)
page_table = torch.zeros(
(num_rows, _ISSUE_ALIGNED_COLS // 64), dtype=torch.int32
)
with (
patch.dict(sys.modules, {"deep_gemm": deep_gemm}),
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(False),
envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.override(False),
patch(f"{_METADATA}.is_hip", return_value=False),
patch(f"{_METADATA}.is_xpu", return_value=False),
patch(f"{_METADATA}._IS_SM120", False),
patch.object(
PagedIndexerMetadata, "_mqa_logits_budget", return_value=budget
),
patch(
"sglang.kernels.ops.attention.dsv4.plan_topk_v2",
side_effect=lambda seq_lens: seq_lens.new_zeros(
(seq_lens.shape[0] + 1, 2)
),
) as plan_topk_v2,
):
metadata = PagedIndexerMetadata(
page_size=256,
compressed_page_size=64,
page_table=page_table,
compressed_seq_lens=c4_seq_lens,
use_topk_v2=use_topk_v2,
)
return metadata, deep_gemm, plan_topk_v2
def test_budget_splits_schedules_and_topk_plans_over_the_same_rows(self):
num_rows, budget = 4096, 512 << 20
metadata, deep_gemm, plan_topk_v2 = self._build(
num_rows=num_rows, budget=budget, use_topk_v2=True
)
expected_rows = mqa_logits_rows_per_chunk(
num_rows=num_rows,
row_bytes=mqa_logits_row_bytes(_ISSUE_ALIGNED_COLS),
budget_bytes=budget,
)
self.assertEqual(metadata.rows_per_chunk, expected_rows)
self.assertEqual(metadata.mqa_logits_budget_bytes, budget)
chunks = list(iter_row_chunks(num_rows=num_rows, rows_per_chunk=expected_rows))
self.assertGreater(len(chunks), 1)
self.assertIsInstance(metadata.deep_gemm_metadata, list)
self.assertEqual(len(metadata.deep_gemm_metadata), len(chunks))
schedule_rows = [
call.args[0]
for call in deep_gemm.get_paged_mqa_logits_metadata.call_args_list
]
torch.testing.assert_close(
torch.cat(schedule_rows), metadata.compressed_seq_lens.unsqueeze(-1)
)
self.assertEqual(
[r.shape[0] for r in schedule_rows], [c.stop - c.start for c in chunks]
)
self.assertEqual(len(metadata.topk_metadata_chunks), len(chunks))
# First call is the full-batch plan; the rest are one per chunk.
plan_rows = [call.args[0] for call in plan_topk_v2.call_args_list]
torch.testing.assert_close(plan_rows[0], metadata.compressed_seq_lens)
torch.testing.assert_close(
torch.cat(plan_rows[1:]), metadata.compressed_seq_lens
)
self.assertEqual(
[r.shape[0] for r in plan_rows[1:]], [c.stop - c.start for c in chunks]
)
def test_no_budget_keeps_the_single_call_shape(self):
metadata, deep_gemm, plan_topk_v2 = self._build(
num_rows=4096, budget=None, use_topk_v2=True
)
self.assertIsNone(metadata.rows_per_chunk)
self.assertIsNone(metadata.topk_metadata_chunks)
self.assertIsInstance(metadata.deep_gemm_metadata, torch.Tensor)
deep_gemm.get_paged_mqa_logits_metadata.assert_called_once()
plan_topk_v2.assert_called_once()
class TestChunkedTopKMatchesUnchunked(CustomTestCase):
"""Each row's top-k depends only on its own logits row, sequence length and
page-table row, so scoring the batch in row chunks must select the same
pages as one pass. This is the property the chunk loop relies on."""
def test_row_chunks_select_the_same_pages(self):
torch.manual_seed(0)
rows, width, topk, page_size = 37, 2048, 64, 64
logits = torch.randn(rows, width, dtype=torch.float32)
seq_lens = torch.randint(1, width, (rows,), dtype=torch.int32)
page_table = torch.randint(
0, 4096, (rows, width // page_size), dtype=torch.int32
)
def run(rows_per_chunk):
out = torch.full((rows, topk), -1, dtype=torch.int32)
for rows_slice in iter_row_chunks(
num_rows=rows, rows_per_chunk=rows_per_chunk
):
topk_transform_pytorch_vectorized(
logits[rows_slice],
seq_lens[rows_slice],
page_table[rows_slice],
out[rows_slice],
page_size,
None,
)
# Unsorted top-k: compare the selected sets row by row.
return out.sort(dim=1).values
expected = run(None)
for rows_per_chunk in (1, 7, 16, rows - 1):
with self.subTest(rows_per_chunk=rows_per_chunk):
self.assertTrue(torch.equal(run(rows_per_chunk), expected))
class TestCandidateIndexerGating(CustomTestCase):
def test_candidate_indexer_gating(self):
from sglang.srt.layers.attention.dsv4 import candidate_indexer