diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 93eed48c8..c468087fb 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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 diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index aac6fbad5..55fbf5f25 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -1313,6 +1313,65 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(actual.dtype, torch.uint8) self.assertTrue(torch.equal(actual, expected)) + def test_batch_in_seq_all_gather_rerange_uses_compute_offsets_for_padded_source(self): + import torch + + cp_size = 2 + # Request 0 is a tiny suffix: valid output only has segment 0, but the + # rank-major source payload contains synthetic compute-padding rows in + # the rank-local mirror segment. Request 1 follows it on the same rank. + # Source offsets must therefore be computed from compute splits, while + # output rows must still be restored from valid splits only. + valid_split_lists = [ + [1, 0, 0, 0], + [2, 0, 1, 0], + ] + compute_split_lists = [ + [1, 1, 1, 1], + [2, 0, 1, 0], + ] + row_width = 2 + max_rank_token = 4 + input_tensor_all = torch.zeros((max_rank_token * cp_size, row_width)) + + # Build source rank-major payload by compute split. Values 900+ are + # synthetic padding rows and must never appear in the restored output. + req0_seg0 = torch.tensor([[10.0, 11.0]]) + req0_seg1_pad = torch.tensor([[900.0, 901.0]]) + req0_seg2_pad = torch.tensor([[902.0, 903.0]]) + req0_seg3_pad = torch.tensor([[904.0, 905.0]]) + req1_seg0 = torch.tensor([[20.0, 21.0], [22.0, 23.0]]) + req1_seg2 = torch.tensor([[24.0, 25.0]]) + + # rank0 owns segment 0 then mirror segment 3 for each request. + input_tensor_all[0:1] = req0_seg0 + input_tensor_all[1:2] = req0_seg3_pad + input_tensor_all[2:4] = req1_seg0 + # rank1 owns segment 1 then mirror segment 2 for each request. + rank1 = max_rank_token + input_tensor_all[rank1 : rank1 + 1] = req0_seg1_pad + input_tensor_all[rank1 + 1 : rank1 + 2] = req0_seg2_pad + input_tensor_all[rank1 + 2 : rank1 + 3] = req1_seg2 + + expected = torch.cat([req0_seg0, req1_seg0, req1_seg2], dim=0) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_split_lists=valid_split_lists, + request_compute_split_lists=compute_split_lists, + compute_padding_enabled=True, + max_rank_len=[max_rank_token, max_rank_token], + ) + ) + + actual = _torch_batch_in_seq_all_gather_rerange( + input_tensor_all, + forward_batch, + cp_size=cp_size, + ) + + self.assertTrue(torch.equal(actual, expected)) + def _build_batch_rerange_case( self, *,