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/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,
)