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

[sp] Make attention-TP sequence sharding a per-forward batch property (#37546)

Browse Source 
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Ming Yang
2026-09-04 02:11:03 -07:00
committed by GitHub
co-authored by Claude Lianmin Zheng
parent 67248e04b4
commit 44c786679f
40 changed files with 498 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/batch_overlap/two_batch_overlap.py
+8 -4
View File
@@ -516,7 +516,7 @@ class TboForwardBatchPreparer:
batch,
tbo_children_num_token_non_padded=tbo_children_num_token_non_padded,
# Eager split: the children can carry a CPU count too, so the
# attention 0-token skip (which reads num_token_non_padded_cpu)
# attention 0-token skip (which reads global_num_token_non_padded_cpu)
# survives the split. The cuda-graph plugin path below leaves this
# None because its device buffer is refreshed per replay.
tbo_children_num_token_non_padded_cpu=cls._split_num_token_non_padded(
@@ -806,7 +806,11 @@ class TboForwardBatchPreparer:
extend_num_tokens=extend_num_tokens,
num_token_non_padded=out_num_token_non_padded,
# TODO: handle it when we need TBO + DeepSeek V3.2
num_token_non_padded_cpu=out_num_token_non_padded_cpu,
global_num_token_non_padded=None,
global_num_token_non_padded_cpu=out_num_token_non_padded_cpu,
# The child runs the same forward, so it keeps the parent's
# sharding verdict; its counts above are already per-child.
attn_tp_sequence_sharded=batch.attn_tp_sequence_sharded,
tbo_split_seq_index=None,
tbo_parent_token_range=(start_token_index, end_token_index),
tbo_children=None,
@@ -863,8 +867,8 @@ class TboForwardBatchPreparer:
@staticmethod
def _get_num_token_non_padded_cpu(batch: ForwardBatch) -> int:
num_token_non_padded = (
batch.num_token_non_padded_cpu
if batch.num_token_non_padded_cpu is not None
batch.global_num_token_non_padded_cpu
if batch.global_num_token_non_padded_cpu is not None
else len(batch.input_ids)
)
return num_token_non_padded
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+32 -17
View File
@@ -1274,9 +1274,12 @@ class AscendAttnBackend(AttentionBackend):
if self.use_fia:
if self._can_use_tnd(layer):
num_token_padding = q.shape[0]
if num_token_padding > forward_batch.num_token_non_padded_cpu:
if (
num_token_padding
> forward_batch.global_num_token_non_padded_cpu
):
q, k, v = [
data[: forward_batch.num_token_non_padded_cpu]
data[: forward_batch.global_num_token_non_padded_cpu]
for data in [q, k, v]
]
q = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
@@ -1318,7 +1321,10 @@ class AscendAttnBackend(AttentionBackend):
attn_out = attn_out.view(
-1, layer.tp_q_head_num * layer.v_head_dim
)
if num_token_padding != forward_batch.num_token_non_padded_cpu:
if (
num_token_padding
!= forward_batch.global_num_token_non_padded_cpu
):
attn_out = torch.cat(
[
attn_out,
@@ -1411,9 +1417,12 @@ class AscendAttnBackend(AttentionBackend):
"""FIA supports multi-bs in the current version of CANN"""
q = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
num_token_padding = q.shape[0]
if num_token_padding > forward_batch.num_token_non_padded_cpu:
if (
num_token_padding
> forward_batch.global_num_token_non_padded_cpu
):
q, k, v = [
data[: forward_batch.num_token_non_padded_cpu]
data[: forward_batch.global_num_token_non_padded_cpu]
for data in [q, k, v]
]
attn_output, _ = torch_npu.npu_fused_infer_attention_score(
@@ -1439,7 +1448,10 @@ class AscendAttnBackend(AttentionBackend):
-1, layer.tp_q_head_num * layer.v_head_dim
)
if num_token_padding != forward_batch.num_token_non_padded_cpu:
if (
num_token_padding
!= forward_batch.global_num_token_non_padded_cpu
):
attn_output = torch.cat(
[
attn_output,
@@ -1699,7 +1711,8 @@ class AscendAttnBackend(AttentionBackend):
else:
num_token_padding = q.shape[0]
q, k, v = [
data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]
data[: forward_batch.global_num_token_non_padded_cpu]
for data in [q, k, v]
]
q_nope, q_rope = q.split(
[layer.v_head_dim, self.qk_rope_head_dim], dim=-1
@@ -1789,7 +1802,7 @@ class AscendAttnBackend(AttentionBackend):
attn_output = attn_output.reshape(
[-1, layer.tp_q_head_num, layer.v_head_dim]
)
if num_token_padding != forward_batch.num_token_non_padded_cpu:
if num_token_padding != forward_batch.global_num_token_non_padded_cpu:
attn_output = torch.cat(
[
attn_output,
@@ -1862,7 +1875,8 @@ class AscendAttnBackend(AttentionBackend):
else:
num_token_padding = q.shape[0]
q, k, v = [
data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]
data[: forward_batch.global_num_token_non_padded_cpu]
for data in [q, k, v]
]
q_nope, q_rope = q.split(
@@ -1891,7 +1905,7 @@ class AscendAttnBackend(AttentionBackend):
attn_output = attn_output.reshape(
-1, layer.tp_q_head_num, layer.v_head_dim
)
if num_token_padding != forward_batch.num_token_non_padded_cpu:
if num_token_padding != forward_batch.global_num_token_non_padded_cpu:
attn_output = torch.cat(
[
attn_output,
@@ -2011,7 +2025,7 @@ class AscendAttnBackend(AttentionBackend):
if not self.graph_mode:
num_token_padding = query.shape[0]
query = query[: forward_batch.num_token_non_padded_cpu]
query = query[: forward_batch.global_num_token_non_padded_cpu]
if self.forward_metadata.seq_lens_cpu_int is None:
actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list
@@ -2097,14 +2111,15 @@ class AscendAttnBackend(AttentionBackend):
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if (
not self.graph_mode
and forward_batch.num_token_non_padded_cpu is not None
and forward_batch.num_token_non_padded_cpu != num_token_padding
and forward_batch.global_num_token_non_padded_cpu is not None
and forward_batch.global_num_token_non_padded_cpu != num_token_padding
):
attn_output = torch.cat(
[
attn_output,
attn_output.new_zeros(
num_token_padding - forward_batch.num_token_non_padded_cpu,
num_token_padding
- forward_batch.global_num_token_non_padded_cpu,
*attn_output.shape[1:],
),
],
@@ -2132,8 +2147,8 @@ class AscendAttnBackend(AttentionBackend):
q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim)
if not self.graph_mode:
num_token_padding = q.shape[0]
q_nope = q_nope[: forward_batch.num_token_non_padded_cpu]
q_rope = q_rope[: forward_batch.num_token_non_padded_cpu]
q_nope = q_nope[: forward_batch.global_num_token_non_padded_cpu]
q_rope = q_rope[: forward_batch.global_num_token_non_padded_cpu]
if self.forward_metadata.seq_lens_cpu_int is None:
actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list
else:
@@ -2235,7 +2250,7 @@ class AscendAttnBackend(AttentionBackend):
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if (
not self.graph_mode
and forward_batch.num_token_non_padded_cpu != num_token_padding
and forward_batch.global_num_token_non_padded_cpu != num_token_padding
):
attn_output = torch.cat(
[
python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py
+5 -5
View File
@@ -201,12 +201,12 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
num_token_padding = mixed_qkv.shape[0]
if (
not self.graph_mode
and forward_batch.num_token_non_padded_cpu != num_token_padding
and forward_batch.global_num_token_non_padded_cpu != num_token_padding
):
mixed_qkv = mixed_qkv[: forward_batch.num_token_non_padded_cpu]
a = a[: forward_batch.num_token_non_padded_cpu]
b = b[: forward_batch.num_token_non_padded_cpu]
seq_len = forward_batch.num_token_non_padded_cpu
mixed_qkv = mixed_qkv[: forward_batch.global_num_token_non_padded_cpu]
a = a[: forward_batch.global_num_token_non_padded_cpu]
b = b[: forward_batch.global_num_token_non_padded_cpu]
seq_len = forward_batch.global_num_token_non_padded_cpu
batch_size = cache_indices.shape[0]
draft_token_num = forward_batch.spec_info.draft_token_num
python/sglang/srt/kv_canary/perturb/slot_picker.py
+4 -2
View File
@@ -43,7 +43,9 @@ def collect_active_slots(
if exclude_out_cache_loc:
out_cache_loc = maybe_inaccurate_forward_batch.out_cache_loc
if out_cache_loc is not None:
valid_num_tokens = maybe_inaccurate_forward_batch.num_token_non_padded_cpu
valid_num_tokens = (
maybe_inaccurate_forward_batch.global_num_token_non_padded_cpu
)
if valid_num_tokens is None:
valid_num_tokens = int(out_cache_loc.shape[0])
excluded = set(
@@ -87,7 +89,7 @@ def pick_out_cache_loc_slot(
total = int(out_cache_loc.shape[0])
if total <= 0:
return None
valid_num_tokens = maybe_inaccurate_forward_batch.num_token_non_padded_cpu
valid_num_tokens = maybe_inaccurate_forward_batch.global_num_token_non_padded_cpu
if valid_num_tokens is None:
valid_num_tokens = total
valid_num_tokens = int(valid_num_tokens)
python/sglang/srt/layers/attention/hpc_ops_backend.py
+1 -1
View File
@@ -681,7 +681,7 @@ def hpc_ops_fp8_rope_store_kv(
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layer = context.attention_layers[layer_id]
real_num_tokens = forward_batch.num_token_non_padded_cpu
real_num_tokens = forward_batch.global_num_token_non_padded_cpu
backend = get_attn_backend()
backend._run_fp8_rope_store_kv(
python/sglang/srt/layers/logits_processor.py
+4
View File
@@ -271,6 +271,9 @@ class LogitsMetadata:
# Whether this batch is prefill-only (no token generation needed)
is_prefill_only: bool = False
# Carried from ForwardBatch so logits pruning can reconstruct the SP gather.
attn_tp_sequence_sharded: bool = False
mm_input_embeds: Optional[torch.Tensor] = None
# DRAFT_EXTEND_V2: when set, lm_head and LAST hidden capture use only these
@@ -324,6 +327,7 @@ class LogitsMetadata:
token_ids_logprobs=forward_batch.token_ids_logprobs,
extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
is_prefill_only=forward_batch.is_prefill_only,
attn_tp_sequence_sharded=forward_batch.attn_tp_sequence_sharded,
global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
dp_local_start_pos=forward_batch.dp_local_start_pos,
dp_local_num_tokens=forward_batch.dp_local_num_tokens,
python/sglang/srt/layers/radix_attention.py
+3 -3
View File
@@ -323,7 +323,7 @@ def _unified_attention_with_output_impl(
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
real_query_num_tokens = forward_batch.num_token_non_padded_cpu
real_query_num_tokens = forward_batch.global_num_token_non_padded_cpu
# Ordinary PCG attention pads Q/K/V to the same token bucket. Prefix MHA
# instead supplies a fixed-capacity K/V chunk whose extent is independent
# of the suffix queries, so its caller must preserve that separate extent.
@@ -531,7 +531,7 @@ def unified_sparse_attention_with_output(
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layer = context.attention_layers[layer_id]
real_num_tokens = forward_batch.num_token_non_padded_cpu
real_num_tokens = forward_batch.global_num_token_non_padded_cpu
if real_num_tokens == 0:
_zero_skipped_attn_outputs(attn_out, idx_out)
@@ -602,7 +602,7 @@ def attention_with_output_extra_kwargs(
context = get_tc_piecewise_forward_context()
forward_batch = context.forward_batch
attention_layer = context.attention_layers[layer_id]
real_num_tokens = forward_batch.num_token_non_padded_cpu
real_num_tokens = forward_batch.global_num_token_non_padded_cpu
if real_num_tokens == 0:
_zero_skipped_attn_outputs(output)
python/sglang/srt/layers/radix_linear_attention.py
+4 -2
View File
@@ -117,7 +117,7 @@ class RadixLinearAttention(nn.Module):
is_extend and not forward_batch.forward_mode.is_target_verify()
)
real_num_tokens = (
getattr(forward_batch, "num_token_non_padded_cpu", None)
getattr(forward_batch, "global_num_token_non_padded_cpu", None)
if should_trim_padded_extend
else None
)
@@ -157,7 +157,9 @@ def _linear_attention_with_output_impl(
forward_batch: ForwardBatch,
) -> None:
"""Run linear attention on the real prefix and initialize physical padding."""
real_num_tokens = min(forward_batch.num_token_non_padded_cpu, mixed_qkv.shape[0])
real_num_tokens = min(
forward_batch.global_num_token_non_padded_cpu, mixed_qkv.shape[0]
)
original_out_cache_loc = forward_batch.out_cache_loc
# Keep the original ForwardBatch object and only narrow cache locations for
python/sglang/srt/model_executor/cpu_graph_runner.py
+3 -1
View File
@@ -940,8 +940,10 @@ class CPUGraphRunner:
)
captured_forward_batch.encoder_out_cache_loc = None
if enable_num_token_non_padded():
# CPUGraphRunner asserts not require_gathered_buffer, so this path is
# never attn-TP sharded: LOCAL == GLOBAL.
captured_forward_batch.num_token_non_padded.copy_(
forward_batch.num_token_non_padded
forward_batch.global_num_token_non_padded
)
self.model_runner.attn_backend.init_forward_metadata(captured_forward_batch)
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+39 -22
View File
@@ -522,6 +522,8 @@ def build_decode_registry(
require_gathered_buffer: bool = False,
enable_prefill_cp: bool = False,
require_mlp_tp_gather: bool = False,
# Per-bucket attn-TP sharded (SP) predicate; defaults to replicated.
attn_tp_sharded_fn: Callable[[int], bool] = lambda num_tokens: False,
dp_size: int = 1,
register_global_num_tokens: bool = True,
share_pool: bool = True,
@@ -648,15 +650,24 @@ def build_decode_registry(
)
def _num_token_non_padded_post_fill(buf, fb, ctx):
# Gathered (DP) path overwrites the plain FB copy with this rank's
# local count; the non-gathered path keeps the copied value.
if require_gathered_buffer and not enable_prefill_cp:
buf.copy_(
compute_local_num_token_non_padded(
global_num_token_non_padded=fb.num_token_non_padded,
num_tokens_per_dp=ctx.padded_num_tokens,
)
# init_new batches localize from the invariant GLOBAL scalar (a
# replicated / CP forward keeps the full count; sharded=False is a
# passthrough). The dense SBD draft and TBO sub-batches bypass
# init_new -- they leave the GLOBAL None and set the replicated LOCAL
# count directly, so carry that through.
if fb.global_num_token_non_padded is None:
buf.copy_(fb.num_token_non_padded)
return
sharded = not enable_prefill_cp and attn_tp_sharded_fn(
ctx.padded_num_tokens
)
buf.copy_(
compute_local_num_token_non_padded(
global_num_token_non_padded=fb.global_num_token_non_padded,
num_tokens_per_dp=ctx.padded_num_tokens,
sharded=sharded,
)
)
slots.append(
GraphSlot(
@@ -664,6 +675,7 @@ def build_decode_registry(
lambda _bs, _mt: (1,),
torch.int32,
axis="none",
copy_from_fb=False,
post_fill=_num_token_non_padded_post_fill,
)
)
@@ -810,6 +822,8 @@ def build_prefill_registry(
enable_num_token_non_padded: bool = False,
require_gathered_buffer: bool = False,
enable_prefill_cp: bool = False,
# Per-bucket attn-TP sharded (SP) predicate; defaults to replicated.
attn_tp_sharded_fn: Callable[[int], bool] = lambda num_tokens: False,
register_input_embeds: bool = True,
share_pool: bool = True,
source: Optional[Any] = None,
@@ -904,23 +918,25 @@ def build_prefill_registry(
)
def _prefill_num_token_non_padded_post_fill(buf, fb, ctx):
# The FB tensor was attn-TP-localized against the RAW length, but
# replay pads rows up to the capture bucket, moving the shard
# boundary — copying it verbatim would make the in-graph pad mask
# blank real tokens whenever raw < bucket. Recompute the local
# count against the padded bucket from the batch's un-adjusted
# global count, mirroring the decode registry's post_fill.
if require_gathered_buffer:
if not enable_prefill_cp:
buf.fill_(
compute_local_num_token_non_padded_cpu(
global_num_token_non_padded=fb.num_token_non_padded_cpu,
num_tokens_per_dp=ctx.padded_num_tokens,
)
# LOCAL count for the PADDED bucket, derived from the invariant
# GLOBAL host int: replay pads rows up to the capture bucket, which
# moves the shard boundary, so the count must be recomputed against
# the bucket rather than copied. A replicated / CP forward keeps the
# full count (sharded=False is a passthrough).
if fb.global_num_token_non_padded_cpu is not None:
sharded = not enable_prefill_cp and attn_tp_sharded_fn(
ctx.padded_num_tokens
)
buf.fill_(
compute_local_num_token_non_padded_cpu(
global_num_token_non_padded=fb.global_num_token_non_padded_cpu,
num_tokens_per_dp=ctx.padded_num_tokens,
sharded=sharded,
)
)
else:
# Non-gathered FullCG still needs the live boundary rather
# than a stale/absent ForwardBatch tensor.
# than a stale/absent ForwardBatch value.
buf.fill_(ctx.raw_num_tokens)
slots.append(
@@ -929,6 +945,7 @@ def build_prefill_registry(
lambda _bs2, _mt: (1,),
torch.int32,
axis="none",
copy_from_fb=False,
post_fill=_prefill_num_token_non_padded_post_fill,
)
)
python/sglang/srt/model_executor/forward_batch_info.py
+99 -49
View File
@@ -32,7 +32,7 @@ import warnings
from dataclasses import dataclass
from enum import IntEnum, auto
from functools import total_ordering
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Set, Tuple, Union
from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union
import torch
@@ -224,19 +224,6 @@ class CaptureHiddenMode(IntEnum):
return self.value < other.value
# Predicate for whether a forward's sequence is sharded across the attn-TP group
# (vs. replicated on every rank). Injected at init; unset defaults to sharded.
_attn_tp_sequence_sharded_predicate: Optional[Callable[[int], bool]] = None
def register_attn_tp_sequence_sharded_predicate(
predicate: Callable[[int], bool],
) -> None:
"""Register the predicate for whether a forward is sharded across attn-TP."""
global _attn_tp_sequence_sharded_predicate
_attn_tp_sequence_sharded_predicate = predicate
def get_server_return_hidden_states_mode() -> CaptureHiddenMode:
features = get_exec().features
mode = features.return_hidden_states_mode
@@ -257,15 +244,14 @@ def get_required_capture_hidden_mode(
return max(capture_hidden_mode, spec_capture_hidden_mode)
def _attn_tp_local_shard_bounds(num_tokens_per_dp: int) -> Tuple[int, int]:
def _attn_tp_local_shard_bounds(
num_tokens_per_dp: int, *, sharded: bool
) -> Tuple[int, int]:
"""(tokens_per_rank, rank_offset) of this attn-TP rank's slice of the sequence.
A replicated (non-sharded) forward puts the whole sequence on every rank, so
the slice is the full range with no offset; localizing it as a shard would
drop real tokens on non-zero ranks.
A replicated (non-sharded) forward keeps the full range on every rank.
"""
predicate = _attn_tp_sequence_sharded_predicate
if predicate is not None and not predicate(num_tokens_per_dp):
if not sharded:
return num_tokens_per_dp, 0
parallel = get_parallel()
tokens_per_rank = num_tokens_per_dp // parallel.attn_tp_size
@@ -275,13 +261,24 @@ def _attn_tp_local_shard_bounds(num_tokens_per_dp: int) -> Tuple[int, int]:
def compute_local_num_token_non_padded(
global_num_token_non_padded: torch.Tensor,
num_tokens_per_dp: int,
*,
sharded: bool,
) -> torch.Tensor:
"""Compute local non-padded token count for this attention-TP rank.
Converts a global count (across all TP ranks) to a local count for this rank.
The "global" scope is within the current DP rank; DP is handled via num_tokens_per_dp.
``num_tokens_per_dp`` is the padded bucket width for the DP group, so each rank
owns a contiguous ``chunk = num_tokens_per_dp // attn_tp_size`` slice: the local
count is ``clamp(global - chunk * attn_tp_rank, 0, chunk)``. The padded bucket
(not ``ceil(real / attn_tp_size)``) sets the chunk, so a trailing rank can own
zero real tokens. ``sharded`` False returns the global count unchanged
(replicated).
"""
tokens_per_rank, rank_offset = _attn_tp_local_shard_bounds(num_tokens_per_dp)
tokens_per_rank, rank_offset = _attn_tp_local_shard_bounds(
num_tokens_per_dp, sharded=sharded
)
return torch.clamp(
global_num_token_non_padded - rank_offset,
0,
@@ -292,15 +289,14 @@ def compute_local_num_token_non_padded(
def compute_local_num_token_non_padded_cpu(
global_num_token_non_padded: int,
num_tokens_per_dp: int,
*,
sharded: bool,
) -> int:
"""Int-scalar twin of ``compute_local_num_token_non_padded``.
Replay-time hooks hold the global count as a host int
(``num_token_non_padded_cpu``) and write the localized result into a
device buffer; keeping the math on ints lets them use ``Tensor.fill_``
instead of staging a CPU tensor through a host-to-device copy per replay.
"""
tokens_per_rank, rank_offset = _attn_tp_local_shard_bounds(num_tokens_per_dp)
"""Int-scalar twin of ``compute_local_num_token_non_padded`` for replay-time
hooks that hold the global count as a host int."""
tokens_per_rank, rank_offset = _attn_tp_local_shard_bounds(
num_tokens_per_dp, sharded=sharded
)
return min(max(global_num_token_non_padded - rank_offset, 0), tokens_per_rank)
@@ -529,9 +525,36 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# Has to be None when cuda graph is captured.
global_num_tokens_for_logprob_cpu: Optional[List[int]] = None
global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
# For padding
num_token_non_padded: Optional[torch.Tensor] = None # scalar tensor
num_token_non_padded_cpu: int = None
# Real (non-padding) token count, held at two scopes whose meaning never
# changes once set:
# GLOBAL — the real-token count across the attn-TP group before sharding.
# global_num_token_non_padded GPU int32 scalar. The invariant source
# the eager forward and the cuda-graph
# registry localize from on each forward /
# replay. Present only when
# enable_num_token_non_padded()
# (moe_ep_size > 1).
# global_num_token_non_padded_cpu host int. Host-side attention/backend
# slices read it directly; the prefill
# graph registry derives its per-rank GPU
# scalar from it.
# LOCAL — this attn-TP rank's owned count after sharding.
# num_token_non_padded GPU int32 scalar, derived from
# global_num_token_non_padded (see
# compute_local_num_token_non_padded). The
# MoE topk kernel masks padded rows with
# it; replicated forwards keep the full
# count. Left None until localized (eager
# prep / graph replay). Present only when
# enable_num_token_non_padded().
global_num_token_non_padded: Optional[torch.Tensor] = None # scalar, GLOBAL
global_num_token_non_padded_cpu: int = None # host int, GLOBAL
num_token_non_padded: Optional[torch.Tensor] = None # scalar, LOCAL (derived)
# Whether this forward's sequence is sharded across the attn-TP group (SP on)
# vs. replicated; stamped per forward, defaults to replicated.
attn_tp_sequence_sharded: bool = False
# === Runtime-filled (set during the forward pass / cuda graph / managers; not at construction) ===
# Preallocated piecewise-graph attention output, set by RadixAttention.
@@ -855,12 +878,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
num_tokens = len(batch.input_ids) if batch.input_ids is not None else 0
if enable_num_token_non_padded():
ret.num_token_non_padded = torch.tensor(
ret.global_num_token_non_padded = torch.tensor(
num_tokens,
dtype=torch.int32,
pin_memory=is_pin_memory_available(device),
).to(device, non_blocking=True)
ret.num_token_non_padded_cpu = num_tokens
ret.global_num_token_non_padded_cpu = num_tokens
ret.init_mlp_sync_metadata(batch, device)
@@ -1005,21 +1028,29 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
pin_memory=is_pin_memory_available(batch.device),
).to(batch.device, non_blocking=True)
def adjust_num_token_non_padded_for_attn_tp(self) -> None:
"""Make num_token_non_padded local to this attention-TP rank."""
def set_local_num_token_non_padded(self, *, sharded: bool) -> None:
"""Derive the LOCAL num_token_non_padded from the invariant GLOBAL scalar.
A replicated (``sharded=False``) forward keeps the full DP-group count.
"""
from sglang.srt.utils.common import require_mlp_tp_gather
dp_rank = get_parallel().attn_dp_rank
assert self.global_num_tokens_cpu is not None
if require_mlp_tp_gather():
num_tokens_per_dp = self.global_num_tokens_cpu[dp_rank]
if self.global_num_tokens_cpu is not None:
# DP / MLP-sync path: per-DP padded width.
if require_mlp_tp_gather():
num_tokens_per_dp = self.global_num_tokens_cpu[
get_parallel().attn_dp_rank
]
else:
num_tokens_per_dp = self.global_num_tokens_cpu[0]
else:
num_tokens_per_dp = self.global_num_tokens_cpu[0]
# Pure TP+SP: local input width.
num_tokens_per_dp = self._forward_num_tokens()
self.num_token_non_padded = compute_local_num_token_non_padded(
global_num_token_non_padded=self.num_token_non_padded,
global_num_token_non_padded=self.global_num_token_non_padded,
num_tokens_per_dp=num_tokens_per_dp,
sharded=sharded,
)
def merge_mm_inputs(self) -> Optional[MultimodalInputs]:
@@ -1369,6 +1400,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
else:
num_tokens = global_num_tokens[0]
self.attn_tp_sequence_sharded = model_runner.attn_tp_sequence_sharded(
num_tokens
)
self.global_dp_buffer_len = buffer_len
set_dp_buffer_len(
buffer_len,
@@ -1456,14 +1491,17 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
mask_dummy_tokens = (
not hybrid_ssm and self._original_forward_mode.is_idle()
)
# Bump the GLOBAL scalar; the LOCAL count is derived from it
# downstream. (global_num_token_non_padded is None unless
# moe_ep_size > 1.)
if mask_dummy_tokens:
if self.num_token_non_padded is not None:
self.num_token_non_padded.fill_(0)
self.num_token_non_padded_cpu = 0
if self.global_num_token_non_padded is not None:
self.global_num_token_non_padded.fill_(0)
self.global_num_token_non_padded_cpu = 0
else:
if self.num_token_non_padded is not None:
self.num_token_non_padded.fill_(num_tokens)
self.num_token_non_padded_cpu = num_tokens
if self.global_num_token_non_padded is not None:
self.global_num_token_non_padded.fill_(num_tokens)
self.global_num_token_non_padded_cpu = num_tokens
else:
self.extend_num_tokens = bs
self.extend_seq_lens = torch.full_like(self.seq_lens, 1)
@@ -1633,9 +1671,21 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
spec_info.hidden_states, num_tokens
)
def _forward_num_tokens(self) -> int:
"""Token width of this forward (``input_embeds`` over ``input_ids``), as
the model counts it for its SP gate."""
if self.input_embeds is not None:
return self.input_embeds.shape[0]
return self.input_ids.shape[0]
def prepare_attn_tp_scatter_input(self, model_runner: ModelRunner):
from sglang.srt.layers.communicator import get_attn_tp_context
# Pure TP+SP has no MLP-sync pass, so stamp the decision here.
self.attn_tp_sequence_sharded = model_runner.attn_tp_sequence_sharded(
self._forward_num_tokens()
)
attn_tp_context = get_attn_tp_context()
input_scattered = attn_tp_context.use_input_scattered(self)
if not input_scattered:
python/sglang/srt/model_executor/model_runner.py
+22 -14
View File
@@ -1510,6 +1510,13 @@ class ModelRunner:
"""Customize a runner-created dummy batch before attention metadata initialization."""
return forward_batch
def attn_tp_sequence_sharded(self, num_tokens: int) -> bool:
"""Whether this forward (``num_tokens`` tokens) is attn-TP sharded (SP).
Extension point for per-forward SP gating; the default behavior shards iff
it holds a gathered buffer.
"""
return require_gathered_buffer()
def _prepare_eager_forward_batch(self, forward_batch: ForwardBatch) -> None:
"""Pad / normalize a batch for the eager (non-cuda-graph) forward.
@@ -1524,20 +1531,21 @@ class ModelRunner:
else:
forward_batch.prepare_attn_tp_scatter_input(self)
# Normalize num_token_non_padded to be local to this attention TP rank if needed.
# The skip is scoped to DSACPLayerCommunicator-style CP (DSA, MLA): those
# flavors already feed a zigzag-split rank-local layout whose token count
# should not be further divided by attn_tp_size. MHA-arch prefill CP
# (Qwen3/Qwen2 MoE) keeps the attn_tp-replicated layout and wants the
# adjustment to run — see docs/design/prefill-cp-mla.md §Phase 5.
if (
forward_batch.num_token_non_padded is not None
and forward_batch.global_num_tokens_gpu is not None
and require_gathered_buffer()
and not is_dsa_enable_prefill_cp()
and not is_mla_prefill_cp_enabled()
):
forward_batch.adjust_num_token_non_padded_for_attn_tp()
# Derive the LOCAL num_token_non_padded from the GLOBAL scalar. sharded is
# cleared for DSACPLayerCommunicator-style CP (DSA, MLA): those flavors
# already feed a zigzag-split rank-local layout whose token count should
# not be further divided by attn_tp_size, so they keep the full count.
# MHA-arch prefill CP (Qwen3/Qwen2 MoE) keeps the attn_tp-replicated
# layout and wants sharding to apply — see docs/design/prefill-cp-mla.md
# §Phase 5.
if forward_batch.global_num_token_non_padded is not None:
forward_batch.set_local_num_token_non_padded(
sharded=(
forward_batch.attn_tp_sequence_sharded
and not is_dsa_enable_prefill_cp()
and not is_mla_prefill_cp_enabled()
),
)
# Hisparse coordinator — backends now read it from self.model_runner.
if self.hisparse_coordinator is not None:
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+6 -3
View File
@@ -459,6 +459,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
require_gathered_buffer=self.require_gathered_buffer,
enable_prefill_cp=self.enable_prefill_cp,
require_mlp_tp_gather=self.require_mlp_tp_gather,
attn_tp_sharded_fn=self.model_runner.attn_tp_sequence_sharded,
dp_size=self.dp_size,
source=self.buffers,
)
@@ -915,17 +916,18 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
else None
)
# Adjust for attention TP if needed (matching replay path in
# populate_from_forward_batch).
# Localize the count when this bucket is attn-TP sharded (SP on).
attn_tp_sharded = self.model_runner.attn_tp_sequence_sharded(num_tokens)
buffers.num_token_non_padded[...] = num_tokens
if (
enable_num_token_non_padded()
and self.require_gathered_buffer
and not self.enable_prefill_cp
and attn_tp_sharded
):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=buffers.num_token_non_padded,
num_tokens_per_dp=num_tokens,
sharded=True,
)
buffers.num_token_non_padded.copy_(local)
@@ -1009,6 +1011,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
spec_info=spec_info,
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=buffers.num_token_non_padded,
attn_tp_sequence_sharded=attn_tp_sharded,
global_forward_mode=self.capture_forward_mode,
lora_ids=lora_ids,
rids_int=rids_int,
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+26 -7
View File
@@ -382,6 +382,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
enable_prefill_cp=(
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
),
attn_tp_sharded_fn=self.model_runner.attn_tp_sequence_sharded,
source=self.buffers,
)
@@ -646,10 +647,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
buf = self.buffer_registry.get_slot("num_token_non_padded").buffer
buf.fill_(num_tokens)
if require_gathered_buffer():
# Localize the count when this bucket is attn-TP sharded (SP on).
if self.model_runner.attn_tp_sequence_sharded(num_tokens):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=buf,
num_tokens_per_dp=num_tokens,
sharded=True,
)
buf.copy_(local)
return buf
@@ -1390,7 +1393,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# Ported from main #27468.
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=self._capture_num_token_non_padded(num_tokens),
num_token_non_padded_cpu=num_tokens,
global_num_token_non_padded_cpu=num_tokens,
attn_tp_sequence_sharded=self.model_runner.attn_tp_sequence_sharded(
num_tokens
),
global_forward_mode=ForwardMode.EXTEND,
# All-None ids are safe: kernels no-op at rank 0 and replay
# refreshes the static batch info with live values.
@@ -1660,7 +1666,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
spec_info=padded_spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=num_token_non_padded,
num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
global_num_token_non_padded_cpu=forward_batch.global_num_token_non_padded_cpu,
attn_tp_sequence_sharded=self.model_runner.attn_tp_sequence_sharded(
static_num_tokens
),
global_forward_mode=pcg_global_forward_mode,
lora_ids=forward_batch.lora_ids,
sampling_info=forward_batch.sampling_info,
@@ -1785,10 +1794,20 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
original_layer_forward = self.layer_model.forward
self.layer_model.forward = replay_layer_forward
# For Full, run the eager tail against the raw user-facing batch so it
# uses real request metadata instead of padded slots. BCG has no
# request-slot padding, so static_forward_batch is already the serving batch.
tail_batch = forward_batch if full_path else static_forward_batch
if full_path:
# Run the eager tail against the raw user-facing batch so it uses
# real request metadata instead of padded slots. Its SP verdict came
# from the unpadded count, though, while the captured body froze one
# at the padded bucket count; the two can differ, so carry the body's
# on a private view rather than mutating the caller's batch.
tail_batch = copy.copy(forward_batch)
tail_batch.attn_tp_sequence_sharded = (
static_forward_batch.attn_tp_sequence_sharded
)
else:
# BCG has no request-slot padding, so static_forward_batch is
# already the serving batch and already holds the body's verdict.
tail_batch = static_forward_batch
if not full_path:
# MTP consumes the target model's live multimodal embeddings in its
# eager wrapper before the captured transformer body is replayed.
python/sglang/srt/model_executor/runner_utils/buffers.py
+1
View File
@@ -329,6 +329,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=forward_batch.num_token_non_padded,
num_tokens_per_dp=num_tokens_per_dp,
sharded=forward_batch.attn_tp_sequence_sharded,
)
dsts.append(self.num_token_non_padded)
srcs.append(local)
python/sglang/srt/models/deepseek_v4.py
+1 -1
View File
@@ -596,7 +596,7 @@ def deepseek_v4_attention_with_output(
forward_batch = context.forward_batch
attention_layers = context.attention_layers
attention_layer = attention_layers[layer_id]
real_num_tokens = forward_batch.num_token_non_padded_cpu
real_num_tokens = forward_batch.global_num_token_non_padded_cpu
if real_num_tokens == 0:
output.zero_()
python/sglang/srt/models/inkling.py
+2 -2
View File
@@ -405,7 +405,7 @@ class InklingDecoderLayer(nn.Module):
attn_out / residual_out and returns None (the eager_on_graph copy-back is
per-tensor, not per-tuple, so outputs must be pre-allocated buffers)."""
forward_batch = get_tc_piecewise_forward_context().forward_batch
n = forward_batch.num_token_non_padded_cpu
n = forward_batch.global_num_token_non_padded_cpu
# log_scaling_tau is per-token, so narrow it to match the real tokens too.
hs, res = self._attn_block(
hidden_states[:n],
@@ -429,7 +429,7 @@ class InklingDecoderLayer(nn.Module):
"""Eager break for the final layer's deferred mlp_sconv: run on the real
tokens with the live forward_batch, write the padded output buffer."""
forward_batch = get_tc_piecewise_forward_context().forward_batch
n = forward_batch.num_token_non_padded_cpu
n = forward_batch.global_num_token_non_padded_cpu
y = self.mlp_sconv(hidden_states[:n], positions[:n], forward_batch)
if self.scattered_sconv:
# y is the [n, H/P] shard; the output buffer is post-gather [n, H].
python/sglang/srt/models/nemotron_h_utils.py
+4 -3
View File
@@ -27,9 +27,10 @@ def get_real_num_tokens(
) -> int:
"""Number of real (non DP-padding) rows in ``hidden_states``."""
real_tokens = hidden_states.shape[0]
num_token_non_padded_cpu = getattr(forward_batch, "num_token_non_padded_cpu", None)
if num_token_non_padded_cpu is not None:
real_tokens = min(real_tokens, int(num_token_non_padded_cpu))
if forward_batch.global_num_token_non_padded_cpu is not None:
real_tokens = min(
real_tokens, int(forward_batch.global_num_token_non_padded_cpu)
)
if (
forward_batch.forward_mode.is_extend()
and not forward_batch.forward_mode.is_mixed()
python/sglang/srt/speculative/dspark_components/dspark_draft.py
+6 -4
View File
@@ -426,8 +426,10 @@ class DraftBlockProposer:
spec_algorithm=SpeculativeAlgorithm.DSPARK,
spec_info=self._draft_block_spec_info,
capture_hidden_mode=CaptureHiddenMode.NULL,
num_token_non_padded=_make_num_token_non_padded(draft_num_tokens, device),
num_token_non_padded_cpu=draft_num_tokens,
global_num_token_non_padded=_make_num_token_non_padded(
draft_num_tokens, device
),
global_num_token_non_padded_cpu=draft_num_tokens,
)
self._fill_dp_moe_sync_metadata(draft_forward_batch, batch)
graph_runner = self.draft_model_runner.decode_cuda_graph_runner
@@ -487,8 +489,8 @@ class DraftBlockProposer:
num_tokens = forward_batch.input_ids.numel()
num_token_non_padded = _make_num_token_non_padded(num_tokens, device)
if num_token_non_padded is not None:
forward_batch.num_token_non_padded = num_token_non_padded
forward_batch.num_token_non_padded_cpu = num_tokens
forward_batch.global_num_token_non_padded = num_token_non_padded
forward_batch.global_num_token_non_padded_cpu = num_tokens
forward_batch.global_num_tokens_cpu = gnt
forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob
pin_memory = is_pin_memory_available(device)
python/sglang/srt/speculative/frozen_kv_mtp_utils.py
+5 -3
View File
@@ -129,9 +129,11 @@ def expand_for_topk_draft(forward_batch: ForwardBatch, topk: int) -> None:
positions = torch.clamp(forward_batch.seq_lens - 1, min=0).to(torch.int64)
forward_batch.positions = positions
forward_batch.num_token_non_padded_cpu = positions.numel()
if forward_batch.num_token_non_padded is not None:
forward_batch.num_token_non_padded.fill_(positions.numel())
forward_batch.global_num_token_non_padded_cpu = positions.numel()
# Bump the GLOBAL scalar; the LOCAL count is derived from it when the draft
# forward localizes (eager prep / graph replay).
if forward_batch.global_num_token_non_padded is not None:
forward_batch.global_num_token_non_padded.fill_(positions.numel())
if (
forward_batch.mrope_positions is not None
and forward_batch.mrope_positions.shape[-1] * topk == positions.numel()
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+1 -1
View File
@@ -354,7 +354,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
extend_seq_lens_cpu=extend_seq_lens_cpu,
extend_start_loc=extend_start_loc,
extend_num_tokens=self.captured_req_width * bs,
num_token_non_padded_cpu=self.captured_req_width * bs,
global_num_token_non_padded_cpu=self.captured_req_width * bs,
return_hidden_states_before_norm=True,
)
return forward_batch
python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py
+2 -2
View File
@@ -1629,7 +1629,7 @@ def run_dsa_forward(
input_hidden = inputs["input_hidden"]
# `input_hidden` may have trailing padding for split-op static-token
# contracts; project only the live token rows for QKV. The kernel
# respects `num_token_non_padded_cpu` via the metadata.
# respects `global_num_token_non_padded_cpu` via the metadata.
live_input_hidden = input_hidden[: case.num_input_tokens]
input_parts = _split_by_lens(live_input_hidden, case.input_lens)
kv_hidden = torch.cat(
@@ -1668,7 +1668,7 @@ def expected_dsa_output_from_inputs(
def dsa_attention_layers(fixture: DSAAttentionFixture) -> list:
"""Return the RadixAttention layers the backend forwards through. The
split-op runner uses this to install per-layer
`num_token_non_padded_cpu` metadata before forward."""
`global_num_token_non_padded_cpu` metadata before forward."""
return [fixture.actual_module.attn]
python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py
+1 -1
View File
@@ -1055,7 +1055,7 @@ def make_lightning_token_padded_inputs(
def lightning_attention_layers(fixture: LightningAttentionFixture) -> list:
"""Return the RadixAttention layers the backend forwards through. The
split-op runner uses this list to install per-layer
`num_token_non_padded_cpu` metadata before forward."""
`global_num_token_non_padded_cpu` metadata before forward."""
return [fixture.actual_module.attn]
python/sglang/test/kits/attention_unittest/runner_modes/split_op_runner.py
+2 -2
View File
@@ -223,13 +223,13 @@ def _make_static_forward_batch(raw_batch, static_num_tokens: int, device: str):
dim=0,
)
raw_batch.num_token_non_padded_cpu = raw_num_tokens
raw_batch.global_num_token_non_padded_cpu = raw_num_tokens
return replace(
raw_batch,
input_ids=input_ids,
positions=positions,
out_cache_loc=out_cache_loc,
num_token_non_padded_cpu=raw_num_tokens,
global_num_token_non_padded_cpu=raw_num_tokens,
)
python/sglang/test/kv_canary/fixtures.py
+2 -2
View File
@@ -175,7 +175,7 @@ def make_forward_batch(
input_ids: Optional[torch.Tensor] = None,
positions: Optional[torch.Tensor] = None,
out_cache_loc: Optional[torch.Tensor] = None,
num_token_non_padded_cpu: Optional[int] = None,
global_num_token_non_padded_cpu: Optional[int] = None,
) -> SimpleNamespace:
seq_lens_default = list(seq_lens_list[:bs])
if req_pool_indices is None:
@@ -216,7 +216,7 @@ def make_forward_batch(
input_ids=input_ids,
positions=positions,
out_cache_loc=out_cache_loc,
num_token_non_padded_cpu=num_token_non_padded_cpu,
global_num_token_non_padded_cpu=global_num_token_non_padded_cpu,
req_all_ids_flat=None,
req_all_ids_lens=None,
)