[6/6][kimi-deterministic] Use deterministic seeded coins for EAGLE rejection sampling (#30822)

test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py

@@ -0,0 +1,108 @@
+"""Unit tests for the DP-attention MLP-sync pad/unpad round-trip.
+
+``prepare_mlp_sync_batch`` pads per-request tensors (positions / seq_lens /
+req_pool_indices) by appending dummy rows after the real ones so all DP ranks
+agree on tensor shapes. ``post_forward_mlp_sync_batch`` must slice them back so
+post-forward consumers — seeded sampling (which asserts positions rows ==
+sampling rows), ngram token-table updates — never see the padding.
+
+Pure dataclass logic — CPU only.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _mock_model_runner(seq_len_fill_value: int = 1) -> MagicMock:
+    runner = MagicMock()
+    runner.attn_backend.get_cuda_graph_seq_len_fill_value.return_value = (
+        seq_len_fill_value
+    )
+    return runner
+
+
+def _logits_output(num_rows: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        next_token_logits=torch.randn(num_rows, 16), hidden_states=None
+    )
+
+
+class TestMlpSyncPadUnpad(CustomTestCase):
+    def test_decode_post_forward_unpads_per_request_tensors(self):
+        fb = ForwardBatch(
+            forward_mode=ForwardMode.DECODE,
+            batch_size=3,
+            input_ids=torch.tensor([11, 12, 13]),
+            req_pool_indices=torch.tensor([5, 6, 7]),
+            seq_lens=torch.tensor([7, 8, 9]),
+            out_cache_loc=torch.tensor([0, 1, 2]),
+            seq_lens_sum=24,
+            positions=torch.tensor([6, 7, 8]),
+            seq_lens_cpu=torch.tensor([7, 8, 9]),
+            lora_ids=[None, None, None],
+        )
+        # Mirror the decode arm of prepare_mlp_sync_batch: record the original
+        # batch size, adopt the synced (padded) one, then pad the inputs.
+        padded = 5
+        fb._original_batch_size = fb.batch_size
+        fb.batch_size = padded
+        fb._pad_inputs_to_size(_mock_model_runner(), num_tokens=padded, bs=padded)
+
+        # Padding appends dummy rows after the real ones.
+        self.assertEqual(fb.positions.shape[0], padded)
+        self.assertEqual(fb.seq_lens.shape[0], padded)
+        self.assertEqual(fb.req_pool_indices.shape[0], padded)
+        torch.testing.assert_close(fb.positions[:3], torch.tensor([6, 7, 8]))
+
+        logits_output = _logits_output(padded)
+        fb.post_forward_mlp_sync_batch(logits_output)
+
+        self.assertEqual(fb.batch_size, 3)
+        torch.testing.assert_close(fb.positions, torch.tensor([6, 7, 8]))
+        torch.testing.assert_close(fb.seq_lens, torch.tensor([7, 8, 9]))
+        torch.testing.assert_close(fb.req_pool_indices, torch.tensor([5, 6, 7]))
+        torch.testing.assert_close(fb.seq_lens_cpu, torch.tensor([7, 8, 9]))
+        self.assertEqual(logits_output.next_token_logits.shape[0], 3)
+        # Seeded sampling asserts positions rows == sampled (real) rows.
+        self.assertEqual(fb.positions.shape[0], fb.batch_size)
+
+    def test_extend_post_forward_unpads_positions(self):
+        fb = ForwardBatch(
+            forward_mode=ForwardMode.EXTEND,
+            batch_size=2,
+            input_ids=torch.arange(7),
+            req_pool_indices=torch.tensor([1, 2]),
+            seq_lens=torch.tensor([3, 4]),
+            out_cache_loc=torch.arange(7),
+            seq_lens_sum=7,
+            positions=torch.tensor([0, 1, 2, 0, 1, 2, 3]),
+            seq_lens_cpu=torch.tensor([3, 4]),
+            lora_ids=[None, None],
+        )
+        # Extend keeps batch_size; only token-level tensors get padded.
+        fb._original_batch_size = fb.batch_size
+        fb._pad_inputs_to_size(_mock_model_runner(), num_tokens=10, bs=2)
+
+        self.assertEqual(fb.positions.shape[0], 10)
+
+        logits_output = _logits_output(10)
+        fb.post_forward_mlp_sync_batch(logits_output)
+
+        torch.testing.assert_close(fb.positions, torch.tensor([0, 1, 2, 0, 1, 2, 3]))
+        torch.testing.assert_close(fb.seq_lens, torch.tensor([3, 4]))
+        # sample() derives prefill sampling positions from seq_lens - 1, so the
+        # row count must match the real request count.
+        self.assertEqual((fb.seq_lens - 1).shape[0], fb.batch_size)
+
+
+if __name__ == "__main__":
+    unittest.main()