chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Keep compute-padding source offsets out of valid rerange output

Browse Source 
Batch in-seq CP rerange has two lengths under compute padding: source payload length includes synthetic rows used only to keep CP compute well-shaped, while output length must expose only valid request rows.  The torch fallback now computes rank-major source offsets from compute splits and output placement from valid splits.

Constraint: Tiny extend batching can add compute-padding rows that must not appear in restored valid token order.

Rejected: Use valid splits for source offsets | following requests on the same rank are shifted when a previous request has padded mirror rows.

Confidence: medium

Scope-risk: narrow

Directive: Batch rerange implementations must distinguish source compute splits from valid output splits.

Tested: python -m py_compile on changed runtime files.

Not-tested: Local pytest blocked before collection by missing orjson dependency.
(cherry picked from commit 31e741477503caa52f3a23acdc1286f46079043c)
This commit is contained in:
laoyao0822
2026-06-07 23:53:29 +08:00
parent c9a39ccdd2
commit 0f9b445131
2 changed files with 117 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -2368,6 +2368,12 @@ def _torch_batch_in_seq_all_gather_rerange(
metadata = getattr(forward_batch, "nsa_cp_metadata", None)
batch_size = int(getattr(metadata, "batch_size", 1) or 1)
request_split_lists = getattr(metadata, "request_split_lists", None)
compute_padding_enabled = bool(
getattr(metadata, "compute_padding_enabled", False)
)
request_compute_split_lists = getattr(
metadata, "request_compute_split_lists", None
)
max_rank_len = getattr(metadata, "max_rank_len", None)
if metadata is None:
_raise_batch_rerange_error("missing_metadata", "nsa_cp_metadata is missing")
@@ -2384,6 +2390,17 @@ def _torch_batch_in_seq_all_gather_rerange(
batch_size,
request_split_lists,
)
if compute_padding_enabled and (
request_compute_split_lists is None
or len(request_compute_split_lists) != batch_size
):
_raise_batch_rerange_error(
"missing_request_compute_split_lists",
"compute-padding rerange requires request_compute_split_lists. "
"batch_size=%s value=%s",
batch_size,
request_compute_split_lists,
)
if max_rank_len is None or len(max_rank_len) < cp_size:
_raise_batch_rerange_error(
"missing_max_rank_len",
@@ -2394,7 +2411,7 @@ def _torch_batch_in_seq_all_gather_rerange(
if cp_size <= 0:
_raise_batch_rerange_error("bad_cp_size", "cp_size must be positive: %s", cp_size)
split_lists: List[List[int]] = []
output_split_lists: List[List[int]] = []
for req_id, split_list in enumerate(request_split_lists):
if split_list is None or len(split_list) != cp_size * 2:
_raise_batch_rerange_error(
@@ -2405,7 +2422,38 @@ def _torch_batch_in_seq_all_gather_rerange(
cp_size,
split_list,
)
split_lists.append([int(x) for x in split_list])
output_split_lists.append([int(x) for x in split_list])
if compute_padding_enabled:
source_split_lists: List[List[int]] = []
for req_id, split_list in enumerate(request_compute_split_lists):
if split_list is None or len(split_list) != cp_size * 2:
_raise_batch_rerange_error(
"bad_request_compute_split",
"compute split must have 2 * cp_size entries. req_id=%s "
"cp_size=%s split=%s",
req_id,
cp_size,
split_list,
)
source_split = [int(x) for x in split_list]
output_split = output_split_lists[req_id]
for segment_id, (source_len, output_len) in enumerate(
zip(source_split, output_split)
):
if source_len < output_len:
_raise_batch_rerange_error(
"compute_split_shorter_than_valid",
"compute source split must cover valid output split. "
"req_id=%s segment=%s source_len=%s output_len=%s",
req_id,
segment_id,
source_len,
output_len,
)
source_split_lists.append(source_split)
else:
source_split_lists = output_split_lists
max_rank_token = int(max_rank_len[0])
if max_rank_token < 0:
@@ -2430,7 +2478,7 @@ def _torch_batch_in_seq_all_gather_rerange(
mirror = cp_size * 2 - source_rank - 1
offsets: List[int] = []
cursor = 0
for split_list in split_lists:
for split_list in source_split_lists:
offsets.append(cursor)
cursor += split_list[source_rank] + split_list[mirror]
if cursor > max_rank_token:
@@ -2444,7 +2492,7 @@ def _torch_batch_in_seq_all_gather_rerange(
)
rank_request_offsets.append(offsets)
total_tokens = sum(sum(split_list) for split_list in split_lists)
total_tokens = sum(sum(split_list) for split_list in output_split_lists)
output_tensor = input_tensor_all.new_empty(
(total_tokens, *input_tensor_all.shape[1:])
)
@@ -2452,9 +2500,10 @@ def _torch_batch_in_seq_all_gather_rerange(
return output_tensor
output_request_base = 0
for req_id, split_list in enumerate(split_lists):
segment_prefix = [0] + list(accumulate(split_list))[:-1]
for segment_id, segment_len in enumerate(split_list):
for req_id, output_split in enumerate(output_split_lists):
source_split = source_split_lists[req_id]
segment_prefix = [0] + list(accumulate(output_split))[:-1]
for segment_id, segment_len in enumerate(output_split):
if segment_len <= 0:
continue
if segment_id < cp_size:
@@ -2462,7 +2511,7 @@ def _torch_batch_in_seq_all_gather_rerange(
source_segment_offset = 0
else:
source_rank = cp_size * 2 - segment_id - 1
source_segment_offset = split_list[source_rank]
source_segment_offset = source_split[source_rank]
source_start = (
source_rank * max_rank_token
+ rank_request_offsets[source_rank][req_id]
@@ -2472,7 +2521,7 @@ def _torch_batch_in_seq_all_gather_rerange(
output_tensor[output_start : output_start + segment_len].copy_(
input_tensor_all[source_start : source_start + segment_len]
)
output_request_base += sum(split_list)
output_request_base += sum(output_split)
return output_tensor