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.