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

support GLM-5.2 MTP index sharing with prefill CP (#30992)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Yuxuan Zhang
2026-07-13 21:11:27 -07:00
committed by GitHub
co-authored by Baizhou Zhang
parent 1b4176cc46
commit 7e229e2a81
7 changed files with 529 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/transform_index.py
+142 -20
View File
@@ -1,3 +1,4 @@
from itertools import accumulate
from typing import List, Optional
import torch
@@ -6,11 +7,39 @@ import triton.language as tl
def transform_index_page_table_prefill(**kwargs):
return transform_index_page_table_prefill_ref(**kwargs)
return transform_index_page_table_prefill_fast(**kwargs)
def transform_index_page_table_decode(**kwargs):
return transform_index_page_table_decode_ref(**kwargs)
return transform_index_page_table_decode_fast(**kwargs)
def _allocate_prefill_result(
topk_indices: torch.Tensor,
real_num_tokens: int,
output_num_tokens: Optional[int],
) -> torch.Tensor:
topk_num_tokens = topk_indices.shape[0]
if output_num_tokens is None:
output_num_tokens = topk_num_tokens
assert real_num_tokens <= topk_num_tokens, (
f"sum(extend_lens_cpu) ({real_num_tokens}) exceeds "
f"topk_indices rows ({topk_num_tokens})"
)
assert topk_num_tokens <= output_num_tokens, (
f"topk_indices rows ({topk_num_tokens}) exceeds "
f"output_num_tokens ({output_num_tokens})"
)
result = torch.empty(
(output_num_tokens, topk_indices.shape[1]),
dtype=torch.int32,
device=topk_indices.device,
)
if real_num_tokens < output_num_tokens:
result[real_num_tokens:].fill_(-1)
return result
@triton.jit
@@ -19,11 +48,11 @@ def transform_index_page_table_decode_kernel(
topk_indices_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_size: tl.constexpr,
max_seqlen_k: tl.constexpr,
page_table_row_stride: tl.constexpr,
):
TOPK: tl.constexpr = 2048
req_id = tl.program_id(0)
page_table_ptr = page_table_ptr + req_id * max_seqlen_k
page_table_ptr = page_table_ptr + req_id * page_table_row_stride
topk_indices_ptr = topk_indices_ptr + req_id * TOPK
result_ptr = result_ptr + req_id * TOPK
@@ -35,6 +64,61 @@ def transform_index_page_table_decode_kernel(
tl.store(result_ptr + offset, -1, mask=~mask)
@triton.jit
def transform_index_page_table_prefill_kernel(
page_table_ptr: torch.Tensor,
topk_indices_ptr: torch.Tensor,
cu_seqlens_q_ptr: torch.Tensor,
result_ptr: torch.Tensor,
page_table_stride_0: tl.constexpr,
page_table_stride_1: tl.constexpr,
topk_indices_stride_0: tl.constexpr,
topk_indices_stride_1: tl.constexpr,
result_stride_0: tl.constexpr,
result_stride_1: tl.constexpr,
PAGE_TABLE_IS_EXPANDED: tl.constexpr,
TOPK: tl.constexpr,
BLOCK_Q: tl.constexpr,
BLOCK_TOPK: tl.constexpr,
):
request_id = tl.program_id(0)
query_offsets = tl.program_id(1) * BLOCK_Q + tl.arange(0, BLOCK_Q)
topk_offsets = tl.program_id(2) * BLOCK_TOPK + tl.arange(0, BLOCK_TOPK)
query_start = tl.load(cu_seqlens_q_ptr + request_id)
query_end = tl.load(cu_seqlens_q_ptr + request_id + 1)
token_indices = query_start + query_offsets
mask = (token_indices[:, None] < query_end) & (topk_offsets[None, :] < TOPK)
loaded_topk_indices = tl.load(
topk_indices_ptr
+ token_indices[:, None] * topk_indices_stride_0
+ topk_offsets[None, :] * topk_indices_stride_1,
mask=mask,
other=-1,
)
valid_topk_mask = mask & (loaded_topk_indices >= 0)
if PAGE_TABLE_IS_EXPANDED:
page_table_rows = token_indices
else:
page_table_rows = token_indices * 0 + request_id
loaded_kv_indices = tl.load(
page_table_ptr
+ page_table_rows[:, None] * page_table_stride_0
+ loaded_topk_indices * page_table_stride_1,
mask=valid_topk_mask,
other=-1,
)
tl.store(
result_ptr
+ token_indices[:, None] * result_stride_0
+ topk_offsets[None, :] * result_stride_1,
loaded_kv_indices,
mask=mask,
)
def transform_index_page_table_decode_fast(
page_table: torch.Tensor,
topk_indices: torch.Tensor,
@@ -54,7 +138,6 @@ def transform_index_page_table_decode_fast(
assert page_table.shape[0] == topk_indices.shape[0]
assert topk_indices.shape[1] == 2048
qo_len = topk_indices.shape[0]
max_seqlen_k = page_table.shape[1]
if result is None:
result = torch.empty_like(topk_indices, dtype=torch.int32)
# Launch triton kernel
@@ -64,7 +147,7 @@ def transform_index_page_table_decode_fast(
topk_indices,
result,
page_size,
max_seqlen_k=max_seqlen_k,
page_table_row_stride=page_table.stride(0),
)
return result
@@ -74,20 +157,48 @@ def transform_index_page_table_prefill_fast(
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
output_num_tokens: Optional[int] = None,
page_table_is_expanded: bool = False,
cu_seqlens_q: Optional[torch.Tensor] = None,
) -> torch.Tensor:
# TODO(baizhou): can be implemented with another triton kernel
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_fast(
page_table[i].unsqueeze(0).expand(l, -1),
topk_indices[offset : offset + l],
result=result[offset : offset + l],
assert topk_indices.shape[1] == 2048
real_num_tokens = sum(extend_lens_cpu)
result = _allocate_prefill_result(topk_indices, real_num_tokens, output_num_tokens)
if real_num_tokens == 0:
return result
max_extend_len = max(extend_lens_cpu)
block_q = 1 if max_extend_len == 1 else 2 if max_extend_len == 2 else 4
block_topk = 256
if cu_seqlens_q is None:
cu_seqlens_q = torch.tensor(
[0, *accumulate(extend_lens_cpu)],
dtype=torch.int32,
device=topk_indices.device,
)
offset += l
assert offset == topk_indices.shape[0]
grid = (
cu_seqlens_q.shape[0] - 1,
triton.cdiv(max_extend_len, block_q),
triton.cdiv(topk_indices.shape[1], block_topk),
)
transform_index_page_table_prefill_kernel[grid](
page_table,
topk_indices,
cu_seqlens_q,
result,
page_table.stride(0),
page_table.stride(1),
topk_indices.stride(0),
topk_indices.stride(1),
result.stride(0),
result.stride(1),
PAGE_TABLE_IS_EXPANDED=page_table_is_expanded,
TOPK=topk_indices.shape[1],
BLOCK_Q=block_q,
BLOCK_TOPK=block_topk,
num_warps=4,
)
return result
@@ -117,10 +228,22 @@ def transform_index_page_table_prefill_ref(
topk_indices: torch.Tensor,
extend_lens_cpu: List[int],
page_size: int = 1,
output_num_tokens: Optional[int] = None,
page_table_is_expanded: bool = False,
) -> torch.Tensor:
assert page_size == 1
result = torch.empty_like(topk_indices, dtype=torch.int32)
assert len(extend_lens_cpu) == page_table.shape[0]
real_num_tokens = sum(extend_lens_cpu)
result = _allocate_prefill_result(topk_indices, real_num_tokens, output_num_tokens)
if page_table_is_expanded:
if real_num_tokens > 0:
transform_index_page_table_decode_ref(
page_table[:real_num_tokens],
topk_indices[:real_num_tokens],
result=result[:real_num_tokens],
)
return result
offset = 0
for i, l in enumerate(extend_lens_cpu):
transform_index_page_table_decode_ref(
@@ -129,7 +252,6 @@ def transform_index_page_table_prefill_ref(
result=result[offset : offset + l],
)
offset += l
assert offset == topk_indices.shape[0]
return result
python/sglang/srt/layers/attention/dsa_backend.py
+22 -9
View File
@@ -1898,19 +1898,20 @@ class DeepseekSparseAttnBackend(
q_nope = q_all[:, :, : layer.v_head_dim]
q_rope = q_all[:, :, layer.v_head_dim :]
# Align topk_indices with q dimensions
# This handles cases where q is padded (TP + partial DP attention)
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0])
# NOTE(dark): here, we use page size = 1
topk_transform_method = self.get_topk_transform_method(
forward_batch.forward_mode
)
if self.use_fused_topk:
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0])
page_table_1 = self._get_fused_topk_page_table(topk_indices)
else:
if topk_transform_method == TopkTransformMethod.RAGGED:
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0])
topk_indices_offset = metadata.topk_indices_offset
assert topk_indices_offset is not None
mask = topk_indices != -1
@@ -1929,6 +1930,12 @@ class DeepseekSparseAttnBackend(
topk_indices=topk_indices,
extend_lens_cpu=metadata.dsa_extend_seq_lens_list,
page_size=1,
output_num_tokens=q_nope.shape[0],
page_table_is_expanded=(
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend_v2()
),
cu_seqlens_q=metadata.cu_seqlens_q,
)
# todo hisparse: to cover more backends
@@ -2675,11 +2682,9 @@ class DeepseekSparseAttnBackend(
else:
q_all = q.view(-1, layer.tp_q_head_num, layer.head_dim)
# Align topk_indices with q dimensions
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q.shape[0])
if self.use_fused_topk:
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q.shape[0])
page_table_1 = self._get_fused_topk_page_table(topk_indices)
elif is_prefill:
page_table_1 = transform_index_page_table_prefill(
@@ -2687,8 +2692,16 @@ class DeepseekSparseAttnBackend(
topk_indices=topk_indices,
extend_lens_cpu=metadata.dsa_extend_seq_lens_list,
page_size=1,
output_num_tokens=q.shape[0],
page_table_is_expanded=(
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend_v2()
),
cu_seqlens_q=metadata.cu_seqlens_q,
)
else:
if topk_indices is not None:
topk_indices = self._pad_topk_indices(topk_indices, q.shape[0])
page_table_1 = transform_index_page_table_decode(
page_table=metadata.page_table_1,
topk_indices=topk_indices,
python/sglang/srt/models/deepseek_nextn.py
+61 -18
View File
@@ -33,6 +33,7 @@ from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
@@ -60,6 +61,34 @@ from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda, is_npu
def _gather_dsa_topk_indices_for_cp(
topk_indices: torch.Tensor,
local_num_tokens: int,
cp_size: int,
forward_batch: ForwardBatch,
stream,
) -> torch.Tensor:
if (
is_dsa_prefill_cp_round_robin_split()
and topk_indices.shape[0] < local_num_tokens
):
pad_rows = local_num_tokens - topk_indices.shape[0]
topk_indices = torch.cat(
[
topk_indices,
topk_indices.new_full((pad_rows, topk_indices.shape[1]), -1),
],
dim=0,
)
return cp_all_gather_rerange_output(
topk_indices,
cp_size,
forward_batch,
stream,
)
logger = logging.getLogger(__name__)
@@ -222,12 +251,21 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.eh_proj(eh_input)
if dsa_use_prefill_cp(
use_cp = dsa_use_prefill_cp(
forward_batch, self.dsa_enable_prefill_cp
) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp):
) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp)
if use_cp:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
residual = None
seed_buf = (
forward_batch.spec_info.dsa_seed_topk_capture
if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True)
else None
)
should_update_dsa_topk_indices = (
forward_batch.reuse_dsa_topk_indices or seed_buf is not None
)
with get_global_expert_distribution_recorder().disable_this_region():
hidden_states, residual, topk_indices = self.decoder(
positions,
@@ -241,34 +279,39 @@ class DeepseekModelNextN(nn.Module):
else None
),
)
if forward_batch.reuse_dsa_topk_indices:
forward_batch.spec_info.dsa_topk_indices = topk_indices
# MTP IndexShare: on draft-extend, publish the last-token DSA
# indexer top-k to seed (avoid recomputing in) the draft-decode loop.
if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
seed_buf = forward_batch.spec_info.dsa_seed_topk_capture
if seed_buf is not None and topk_indices is not None:
sel = forward_batch.spec_info.dsa_seed_topk_select
src = topk_indices if sel is None else topk_indices[sel]
seed_buf[: src.shape[0]].copy_(src)
if not forward_batch.forward_mode.is_idle():
if residual is not None:
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
else:
hidden_states = self.shared_head.norm(hidden_states)
if dsa_use_prefill_cp(
forward_batch, self.dsa_enable_prefill_cp
) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp):
# allgather + rerrange
if use_cp:
local_num_tokens = hidden_states.shape[0]
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if should_update_dsa_topk_indices and topk_indices is not None:
topk_indices = _gather_dsa_topk_indices_for_cp(
topk_indices,
local_num_tokens,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if should_update_dsa_topk_indices and topk_indices is not None:
if forward_batch.reuse_dsa_topk_indices:
forward_batch.spec_info.dsa_topk_indices = topk_indices
if seed_buf is not None:
sel = forward_batch.spec_info.dsa_seed_topk_select
src = (
topk_indices[: seed_buf.shape[0]]
if sel is None
else topk_indices[sel]
)
seed_buf[: src.shape[0]].copy_(src)
finally:
exit_stack.close()
python/sglang/srt/server_args.py
+1 -1
View File
@@ -5204,7 +5204,7 @@ class ServerArgs:
mode = strategy_to_legacy_mode[self.cp_strategy]
use_dsa_legacy_aliases = self.enable_dsa_prefill_context_parallel or getattr(
self, "attention_backend", None
self._resolved(), "attention_backend", None
) in ("dsa", "dsv4")
if use_dsa_legacy_aliases:
self.enable_dsa_prefill_context_parallel = True
python/sglang/srt/speculative/eagle_worker_v2.py
-3
View File
@@ -15,7 +15,6 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner i
)
from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
from sglang.srt.kv_canary.runner.canary_manager import context_tuple
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
from sglang.srt.layers.attention.tokenspeed_mla_backend import TokenspeedMLABackend
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
@@ -825,11 +824,9 @@ class EagleDraftWorker(EagleDraftWorkerBase):
# Seed the first draft-decode loop from each request's last prefill
# position. Gather last-per-req before the copy (prefill can be long).
# Skipped under context-parallel prefill (token layout wouldn't match).
seed_from_extend = (
self.seed_dsa_topk_from_draft_extend
and not forward_batch.forward_mode.is_idle()
and not dsa_use_prefill_cp(forward_batch)
)
if seed_from_extend:
bs = forward_batch.batch_size