Keep compute-padding source offsets out of valid rerange output

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)
2026-06-08 00:17:28 +08:00
parent c9a39ccdd2
commit 0f9b445131
2 changed files with 117 additions and 9 deletions
@@ -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,
        *,