[Feature] Unified memory: support decode context parallelism for the trtllm_mla family (#37693)
This commit is contained in:
Cheng Wan
@@ -213,6 +213,225 @@ class TestTRTLLMMLARejectsDcpMultiTokenQuery(CustomTestCase):
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
torch.cuda.is_available(), "the page-table build is a Triton kernel"
|
||||
)
|
||||
class TestDcpBlockTableIdSpace(CustomTestCase):
|
||||
"""`_fill_dcp_block_kv_indices` must emit entries in the pool's OWN id space.
|
||||
|
||||
On a static pool the DCP-collapsed page is already physical. Under the
|
||||
unified memory pool it is still VIRTUAL, and an entry that skips the v2p
|
||||
gather names whatever page happens to sit there now -- silent wrong-KV
|
||||
reads rather than a crash, which is why the reference below is computed
|
||||
from the tables rather than from the kernel.
|
||||
"""
|
||||
|
||||
PAGE_SIZE = 64
|
||||
DCP_SIZE = 4
|
||||
DCP_RANK = 2
|
||||
MULTIPLIER = 3
|
||||
# Virtual page per request, deliberately not the identity so a missing
|
||||
# gather cannot coincide with the right answer.
|
||||
VIRTUAL_PAGES = [[5, 2, 9], [7, 0, 4], [1, 8, 6]]
|
||||
# Global KV lengths: one spanning 3 pages, one 2, one under a page.
|
||||
SEQ_LENS = [3 * PAGE_SIZE * DCP_SIZE, 2 * PAGE_SIZE * DCP_SIZE - 1, 10]
|
||||
|
||||
def _make_backend(self, translator):
|
||||
bs = len(self.VIRTUAL_PAGES)
|
||||
max_pos = max(self.SEQ_LENS)
|
||||
req_to_token = torch.full((bs, max_pos), -1, dtype=torch.int32, device="cuda")
|
||||
span = self.PAGE_SIZE * self.DCP_SIZE
|
||||
for req, pages in enumerate(self.VIRTUAL_PAGES):
|
||||
for slot, virtual_page in enumerate(pages):
|
||||
base = virtual_page * span
|
||||
start = slot * span
|
||||
end = min(start + span, self.SEQ_LENS[req])
|
||||
if end <= start:
|
||||
break
|
||||
req_to_token[req, start:end] = torch.arange(
|
||||
base, base + (end - start), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
backend = object.__new__(TRTLLMMLABackend)
|
||||
backend.page_size = self.PAGE_SIZE
|
||||
backend.req_to_token = req_to_token
|
||||
backend.kv_index_translator = translator
|
||||
return backend, req_to_token
|
||||
|
||||
def _fill(self, translator):
|
||||
backend, req_to_token = self._make_backend(translator)
|
||||
bs = len(self.VIRTUAL_PAGES)
|
||||
seq_lens = torch.tensor(self.SEQ_LENS, dtype=torch.int32, device="cuda")
|
||||
local_seq_lens = get_dcp_lens(seq_lens, self.DCP_SIZE, self.DCP_RANK).to(
|
||||
torch.int32
|
||||
)
|
||||
# One 128-page row: the padded width `_calc_padded_blocks` produces.
|
||||
block_kv_indices = torch.full((bs, 128), -1, dtype=torch.int32, device="cuda")
|
||||
parallel = SimpleNamespace(
|
||||
dcp_enabled=True, dcp_size=self.DCP_SIZE, dcp_rank=self.DCP_RANK
|
||||
)
|
||||
with patch.object(backend_module, "get_parallel", return_value=parallel):
|
||||
backend._fill_dcp_block_kv_indices(
|
||||
block_kv_indices,
|
||||
torch.arange(bs, dtype=torch.int64, device="cuda"),
|
||||
local_seq_lens,
|
||||
)
|
||||
return block_kv_indices.cpu(), local_seq_lens.cpu(), req_to_token.cpu()
|
||||
|
||||
def _reference(self, req_to_token, local_seq_lens, v2p, multiplier):
|
||||
"""What each live entry must be, derived from the id-space definition."""
|
||||
rows = []
|
||||
for req in range(local_seq_lens.numel()):
|
||||
local_pages = -(-int(local_seq_lens[req]) // self.PAGE_SIZE)
|
||||
row = []
|
||||
for page in range(local_pages):
|
||||
pos = self.DCP_RANK + page * self.PAGE_SIZE * self.DCP_SIZE
|
||||
widened = int(req_to_token[req, pos])
|
||||
collapsed_page = widened // self.DCP_SIZE // self.PAGE_SIZE
|
||||
if v2p is None:
|
||||
row.append(collapsed_page)
|
||||
else:
|
||||
row.append(max(int(v2p[collapsed_page]) * multiplier, 0))
|
||||
rows.append(row)
|
||||
return rows
|
||||
|
||||
def _assert_matches(self, table, rows):
|
||||
for req, row in enumerate(rows):
|
||||
self.assertEqual(table[req, : len(row)].tolist(), row)
|
||||
# Past the live prefix nothing is written: the caller-owned
|
||||
# capture-stable buffer keeps what it had there.
|
||||
self.assertTrue((table[req, len(row) :] == -1).all())
|
||||
|
||||
def test_static_pool_entries_are_the_collapsed_page(self):
|
||||
translator = SimpleNamespace(full_v2p_table=None, full_page_multiplier=1)
|
||||
table, local_lens, req_to_token = self._fill(translator)
|
||||
self._assert_matches(table, self._reference(req_to_token, local_lens, None, 1))
|
||||
|
||||
def test_unified_pool_entries_go_through_the_page_table(self):
|
||||
num_pages = 1 + max(max(p) for p in self.VIRTUAL_PAGES)
|
||||
# A scrambled v2p, so an entry that skipped the gather would differ.
|
||||
v2p = torch.tensor(
|
||||
[(7 * i + 3) % num_pages for i in range(num_pages)],
|
||||
dtype=torch.int64,
|
||||
device="cuda",
|
||||
)
|
||||
translator = SimpleNamespace(
|
||||
full_v2p_table=v2p, full_page_multiplier=self.MULTIPLIER
|
||||
)
|
||||
table, local_lens, req_to_token = self._fill(translator)
|
||||
expected = self._reference(req_to_token, local_lens, v2p.cpu(), self.MULTIPLIER)
|
||||
self._assert_matches(table, expected)
|
||||
# And it is genuinely a translation, not an accident of the fixture.
|
||||
static = self._reference(req_to_token, local_lens, None, 1)
|
||||
self.assertNotEqual(expected, static)
|
||||
|
||||
def test_freed_page_lands_on_the_padding_sink(self):
|
||||
# A tombstoned (-1) v2p row must clamp to entry 0, the reserved
|
||||
# padding page, rather than scale -1 into a wild block id.
|
||||
num_pages = 1 + max(max(p) for p in self.VIRTUAL_PAGES)
|
||||
v2p = torch.arange(num_pages, dtype=torch.int64, device="cuda")
|
||||
v2p[self.VIRTUAL_PAGES[0][0]] = -1
|
||||
translator = SimpleNamespace(
|
||||
full_v2p_table=v2p, full_page_multiplier=self.MULTIPLIER
|
||||
)
|
||||
table, _, _ = self._fill(translator)
|
||||
self.assertEqual(int(table[0, 0]), 0)
|
||||
|
||||
|
||||
class TestFusedFp8WriteGate(CustomTestCase):
|
||||
"""The fused fp8 KV write must not serve a DCP-resolved loc.
|
||||
|
||||
Its only skip is the DCP owner rule (`vloc % world != rank`). Under the
|
||||
unified pool that rule is already applied before the loc arrives, with
|
||||
non-owned rows folded onto kernel id 0 -- so the kernel cannot skip the
|
||||
padding sink the way `set_mla_kv_buffer`'s `reserved_skip_index` does, and
|
||||
every non-owned token would store into slot 0.
|
||||
|
||||
A revert of the gate leaves accuracy tests green (slot 0 holds no real
|
||||
data), which is why this is asserted directly.
|
||||
"""
|
||||
|
||||
def _gate(self, *, dcp_enabled: bool, is_translating: bool) -> bool:
|
||||
backend = object.__new__(TRTLLMMLABackend)
|
||||
backend.data_type = torch.float8_e4m3fn
|
||||
backend.kv_lora_rank = 512
|
||||
backend.qk_rope_head_dim = 64
|
||||
backend.kv_index_translator = SimpleNamespace(is_translating=is_translating)
|
||||
parallel = SimpleNamespace(
|
||||
dcp_enabled=dcp_enabled,
|
||||
attn_dcp_size=DCP_SIZE if dcp_enabled else 1,
|
||||
attn_dcp_rank=DCP_RANK if dcp_enabled else 0,
|
||||
)
|
||||
with (
|
||||
patch.object(backend_module, "get_parallel", return_value=parallel),
|
||||
patch.object(
|
||||
backend_module, "can_use_set_mla_kv_concat_q_fp8", return_value=True
|
||||
),
|
||||
patch.object(
|
||||
backend_module.envs.SGLANG_ENABLE_ASYNC_ASSERT, "get", lambda: False
|
||||
),
|
||||
):
|
||||
return bool(
|
||||
backend.data_type == torch.float8_e4m3fn
|
||||
and not backend_module.envs.SGLANG_ENABLE_ASYNC_ASSERT.get()
|
||||
and backend.kv_lora_rank == 512
|
||||
and backend.qk_rope_head_dim == 64
|
||||
and not (
|
||||
backend_module.get_parallel().dcp_enabled
|
||||
and backend.kv_index_translator.is_translating
|
||||
)
|
||||
and backend_module.can_use_set_mla_kv_concat_q_fp8()
|
||||
)
|
||||
|
||||
def test_off_for_the_unified_pool_under_dcp(self):
|
||||
self.assertFalse(self._gate(dcp_enabled=True, is_translating=True))
|
||||
|
||||
def test_on_for_every_other_combination(self):
|
||||
# The gate must not cost the static pool or a non-DCP unified run
|
||||
# their fused write.
|
||||
self.assertTrue(self._gate(dcp_enabled=True, is_translating=False))
|
||||
self.assertTrue(self._gate(dcp_enabled=False, is_translating=True))
|
||||
self.assertTrue(self._gate(dcp_enabled=False, is_translating=False))
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "backend construction")
|
||||
def test_helper_refuses_a_resolved_loc(self):
|
||||
"""Belt and braces: reaching the helper with a resolved loc asserts
|
||||
rather than silently storing into the sink."""
|
||||
backend = object.__new__(TRTLLMMLABackend)
|
||||
backend.kv_index_translator = SimpleNamespace(is_translating=True)
|
||||
parallel = SimpleNamespace(
|
||||
dcp_enabled=True, attn_dcp_size=DCP_SIZE, attn_dcp_rank=DCP_RANK
|
||||
)
|
||||
layer = SimpleNamespace(
|
||||
tp_q_head_num=1, v_head_dim=512, head_dim=576, layer_id=0
|
||||
)
|
||||
n = 4
|
||||
with (
|
||||
patch.object(backend_module, "get_parallel", return_value=parallel),
|
||||
patch.object(
|
||||
backend_module, "set_mla_kv_concat_q_fp8_covered", return_value=True
|
||||
),
|
||||
patch.object(
|
||||
backend,
|
||||
"token_to_kv_pool",
|
||||
SimpleNamespace(
|
||||
get_key_buffer=lambda _: torch.zeros(
|
||||
(8, 576), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
),
|
||||
create=True,
|
||||
),
|
||||
):
|
||||
with self.assertRaises(AssertionError):
|
||||
backend._set_kv_and_concat_q_fp8_fused(
|
||||
layer=layer,
|
||||
loc=torch.zeros(n, dtype=torch.int64, device="cuda"),
|
||||
q=torch.zeros((n, 512), dtype=torch.bfloat16, device="cuda"),
|
||||
q_rope=torch.zeros((n, 64), dtype=torch.bfloat16, device="cuda"),
|
||||
k=torch.zeros((n, 512), dtype=torch.bfloat16, device="cuda"),
|
||||
k_rope=torch.zeros((n, 64), dtype=torch.bfloat16, device="cuda"),
|
||||
)
|
||||
|
||||
|
||||
class TestDcpDecodeLayout(CustomTestCase):
|
||||
"""Rank-local length math the decode page table above is built from."""
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Kimi-Linear on the unified memory pool with decode context parallelism,
|
||||
served by the trtllm_mla family.
|
||||
|
||||
Two read-index routes exist under `--enable-unified-memory --dcp-size > 1`.
|
||||
`test_kimi_linear_unified_memory.py` covers the flashinfer one, where
|
||||
`plan_dcp_decode_metadata` compacts this rank's ids and
|
||||
`translate_dcp_read_ids` converts them. This file covers the other: a block
|
||||
table built by `create_mla_kv_page_table_for_dcp`, which gathers a rank's
|
||||
cyclic slice and must take the pool's virtual->physical page table on the way.
|
||||
An entry that skipped that hop names whichever page sits there now, so the
|
||||
failure is stale KV rather than a crash -- which is what the GSM8K bar and the
|
||||
prefix-cache replay below are here to catch.
|
||||
|
||||
`cutedsl_mla` stands in for the family: it is the DCP-native MLA decode kernel
|
||||
on Blackwell (the only backend flashinfer accepts `enable_dcp=True` for) and
|
||||
what the B300 hosts run, and prefill resolves to `trtllm_mla`, so one server
|
||||
covers both halves of the page-table contract. `trtllm_mla` and
|
||||
`tokenspeed_mla` share that builder and were checked by hand against the
|
||||
static pool (see below) rather than given a CI server each.
|
||||
|
||||
Blackwell-only, hence its own module: the Hopper suite next door cannot host
|
||||
it. In extra-b rather than nightly because what it guards is a per-PR argument
|
||||
gate, which a regression closes at boot rather than degrading overnight.
|
||||
|
||||
Measured on B300 (1 sigma ~= 0.02). This file scores 0.920 and takes 2.5 min.
|
||||
Matched unified-vs-static pairs, run by hand at 200 questions, to show the two
|
||||
pools agree:
|
||||
|
||||
unified / static
|
||||
TP2 DCP2 cutedsl_mla 0.905 / 0.905
|
||||
TP2 DCP2 trtllm_mla 0.895 / 0.895
|
||||
TP2 DCP2 tokenspeed_mla 0.890 / 0.895, repeat 0.890 / 0.885
|
||||
TP4 DCP4 cutedsl_mla 0.910 / 0.900
|
||||
|
||||
The cutedsl_mla and trtllm_mla pairs reproduced exactly. tokenspeed_mla moved
|
||||
on a repeat of the same questions on BOTH pools, so its spread is the backend,
|
||||
not the pool. The TP4 DCP4 row is there because it widens the virtual page to
|
||||
4 x 64 = 256, past what this file exercises.
|
||||
|
||||
python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py -v
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
|
||||
from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
|
||||
register_cuda_ci(est_time=250, stage="extra-b", runner_config="4-gpu-b200")
|
||||
|
||||
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
|
||||
|
||||
|
||||
class TestKimiLinearUnifiedMemoryDCPCuteDsl(
|
||||
GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
|
||||
):
|
||||
"""cutedsl_mla decode, trtllm_mla prefill (what cutedsl_mla resolves to)."""
|
||||
|
||||
model = KIMI_LINEAR_MODEL
|
||||
cache_chunk_size = 64
|
||||
gsm8k_score_threshold = 0.88
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--tp-size",
|
||||
"2",
|
||||
"--dcp-size",
|
||||
"2",
|
||||
"--attention-backend",
|
||||
"cutedsl_mla",
|
||||
"--chunked-prefill-size",
|
||||
"2048",
|
||||
# The per-batch DCP gather buffer is sized by the batch's total KV, so
|
||||
# the default static fraction leaves a 200-question burst ~1 GB and it
|
||||
# OOMs inside an NCCL collective.
|
||||
"--mem-fraction-static",
|
||||
"0.80",
|
||||
"--max-running-requests",
|
||||
"128",
|
||||
"--cuda-graph-max-bs-decode",
|
||||
"128",
|
||||
"--enable-unified-memory",
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user