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

[Feature] Unified memory: support decode context parallelism for the trtllm_mla family (#37693)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-03 03:37:12 -07:00
committed by GitHub
parent f59a4840c5
commit a11dba1a01
8 changed files with 378 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dcp_kernels.py
+18 -2
View File
@@ -82,13 +82,23 @@ def create_mla_kv_page_table_for_dcp(
req_pool_indices_ptr,
local_seq_lens_ptr,
block_kv_indices_ptr,
v2p_ptr, # in: [num_pages + 1] int64 -- virtual->physical page table
req_to_token_stride: tl.constexpr,
block_table_stride: tl.constexpr,
mult, # runtime: kernel_page_multiplier of the target sub-pool
PHYSICAL_PAGE_SIZE: tl.constexpr,
DCP_SIZE: tl.constexpr,
DCP_RANK: tl.constexpr,
PAGES_PER_BLOCK: tl.constexpr,
HAS_V2P: tl.constexpr,
):
"""This rank's cyclic slice of each request, as a page table.
``HAS_V2P`` picks the id space the emitted page number is in: the
DCP-collapsed page IS physical on a static pool, and still VIRTUAL under
the unified memory pool, where it takes one more gather through ``v2p_ptr``
and a ``mult`` scale to reach the per-layer views.
"""
req = tl.program_id(0)
page_block = tl.program_id(1)
page_offsets = page_block * PAGES_PER_BLOCK + tl.arange(0, PAGES_PER_BLOCK)
@@ -102,10 +112,16 @@ def create_mla_kv_page_table_for_dcp(
mask=mask,
other=0,
)
physical_pages = virtual_locs // DCP_SIZE // PHYSICAL_PAGE_SIZE
pages = virtual_locs // DCP_SIZE // PHYSICAL_PAGE_SIZE
if HAS_V2P:
# A `-1` in req_to_token and a freed (`-1`) v2p row both clamp to entry
# 0, the reserved padding page.
pages = tl.where(virtual_locs < 0, 0, pages)
physical = tl.load(v2p_ptr + pages, mask=mask, other=0)
pages = tl.maximum(physical * mult, 0)
tl.store(
block_kv_indices_ptr + req * block_table_stride + page_offsets,
physical_pages,
pages.to(tl.int32),
mask=mask,
)
python/sglang/srt/arg_groups/kv_cache_hook.py
+4 -6
View File
@@ -305,12 +305,10 @@ def _validate_unified_memory_dcp(server_args: Any) -> None:
"transfer, where translate_kv_indices_for_transfer would abort a "
"server that had already booted."
)
# trtllm_mla (and its cutedsl_mla / tokenspeed_mla subclasses) build the
# MLA block table straight from req_to_token with
# create_flashmla_kv_indices_triton, whose v2p gather assumes UNWIDENED
# page ids; the DCP variant (create_mla_kv_page_table_for_dcp) has no v2p
# gather at all. Wire one of them through the other to add those here.
dcp_allowed = {"flashinfer"}
# The trtllm_mla family builds its DCP block table through the pool's v2p
# gather (create_mla_kv_page_table_for_dcp), so it speaks the same
# two-stage contract as flashinfer.
dcp_allowed = {"flashinfer", "trtllm_mla", "cutedsl_mla", "tokenspeed_mla"}
backends = set(attention_backends_of(resolved_view(server_args)))
backends.discard(None)
assert backends <= dcp_allowed, (
python/sglang/srt/layers/attention/cutedsl_mla_backend.py
+12 -17
View File
@@ -174,22 +174,17 @@ class CuteDslMLABackend(TRTLLMMLABackend):
if self.data_type == torch.float8_e4m3fn:
assert q_rope is not None and k_rope is not None
if cos_sin_cache is None:
if (
save_kv_cache
and self._fused_set_kv_concat_q_fp8
and not self.kv_index_translator.is_translating
):
# Static pool: out_cache_loc is already the physical loc.
# Fused: bf16->fp8 quantize + KV scatter + q concat in one
# launch; None when not covered.
query = self._set_kv_and_concat_q_fp8_fused(
layer=layer,
loc=forward_batch.out_cache_loc,
q=q,
q_rope=q_rope,
k=k,
k_rope=k_rope,
)
if save_kv_cache and self._fused_set_kv_concat_q_fp8:
loc = self._resolve_fused_write_loc(forward_batch)
if loc is not None:
query = self._set_kv_and_concat_q_fp8_fused(
layer=layer,
loc=loc,
q=q,
q_rope=q_rope,
k=k,
k_rope=k_rope,
)
if query is None:
q, k, k_rope = mla_quantize_without_rope_for_fp8(
q, q_rope, k.squeeze(1), k_rope.squeeze(1)
@@ -211,7 +206,7 @@ class CuteDslMLABackend(TRTLLMMLABackend):
if query is None and save_kv_cache:
assert k is not None and k_rope is not None
self.token_to_kv_pool.set_mla_kv_buffer(
layer, forward_batch.out_cache_loc, k, k_rope
layer, self._kv_write_loc(forward_batch), k, k_rope
)
if query is not None:
python/sglang/srt/layers/attention/tokenspeed_mla_backend.py
+1 -1
View File
@@ -389,7 +389,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend):
if save_kv_cache:
self.token_to_kv_pool.set_mla_kv_buffer(
layer, forward_batch.out_cache_loc, k, k_rope
layer, self._kv_write_loc(forward_batch), k, k_rope
)
query = q.view(-1, layer.tp_q_head_num, layer.head_dim)
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+21 -2
View File
@@ -377,6 +377,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
) -> None:
parallel = get_parallel()
pages_per_block = get_num_page_per_block_flashmla(self.page_size)
# None on a static pool, whose collapsed page is already physical.
v2p = self.kv_index_translator.full_v2p_table
create_mla_kv_page_table_for_dcp[
(
block_kv_indices.shape[0],
@@ -389,12 +391,15 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
req_pool_indices,
local_seq_lens,
block_kv_indices,
v2p,
self.req_to_token.stride(0),
block_kv_indices.stride(0),
self.kv_index_translator.full_page_multiplier,
PHYSICAL_PAGE_SIZE=self.page_size,
DCP_SIZE=parallel.dcp_size,
DCP_RANK=parallel.dcp_rank,
PAGES_PER_BLOCK=pages_per_block,
HAS_V2P=v2p is not None,
)
def _create_block_kv_indices(
@@ -777,6 +782,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
else:
self._decode_kernel_loc = None
def _kv_write_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
"""The loc an unfused KV scatter must write at: the capture-stable
buffer under a captured unified-pool decode, since the translate
rebinds `out_cache_loc` to a fresh tensor the graph never recorded;
the batch's own loc everywhere else.
"""
if self._decode_kernel_loc is not None:
return self._decode_kernel_loc
return forward_batch.out_cache_loc
def _resolve_fused_write_loc(
self, forward_batch: ForwardBatch
) -> Optional[torch.Tensor]:
@@ -1198,6 +1213,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
):
return None
parallel = get_parallel()
# `loc` is WIDENED: the kernel resolves the owner rule itself, and that
# is also its only skip. A DCP-resolved loc never reaches here -- see
# the `_fused_set_kv_concat_q_fp8` gate.
assert not (parallel.dcp_enabled and self.kv_index_translator.is_translating), (
"fused fp8 KV write reached with a DCP-resolved loc"
)
return set_mla_kv_concat_q_fp8(
kv_buffer=kv_2d,
loc=loc,
@@ -1205,8 +1226,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
cache_k_rope=k_rope_2d,
q_nope=q_nope,
q_rope=q_rope_3d,
# DCP cyclic KV sharding: virtual loc -> owner mask + loc//world
# (identity when attn_dcp_size == 1).
dcp_world_size=parallel.attn_dcp_size,
dcp_rank=parallel.attn_dcp_rank,
)
python/sglang/srt/mem_cache/kv_index_translator.py
+16
View File
@@ -392,6 +392,22 @@ class KVIndexTranslator:
self._index_table_memo = (weakref.ref(forward_batch), view)
return view
@property
def full_v2p_table(self) -> Optional[torch.Tensor]:
"""The full-attention virtual->physical PAGE table, or None when this
pool needs no translation.
For the DCP page-table builders, whose gather is over a rank's cyclic
slice rather than a row prefix, so `build_index_table` cannot serve
them.
"""
return self._full_v2p_table
@property
def full_page_multiplier(self) -> int:
"""Scales a physical page into the id space the per-layer views use."""
return self._full_page_multiplier
def bind_and_verify_backends(self, backends) -> None:
"""Boot: make every reachable backend carry THIS translator.
test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py
+219
View File
@@ -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."""
test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py
+87
View File
@@ -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()