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

test/registered/kernels/ops/attention/test_flash_attention_4.py

@@ -11,7 +11,11 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange, repeat
 
-from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
+from sglang.kernels.ops.attention.flash_attention import (
+    flash_attn_varlen_func,
+    flash_attn_with_kvcache,
+)
+from sglang.srt.utils import is_sm100_or_sm110_supported
 from sglang.test.ci.ci_register import register_cuda_ci
 
 register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
@@ -1507,5 +1511,112 @@ def _generate_block_kvcache(
     return k_cache, v_cache, page_table, k_cache_paged, v_cache_paged, num_blocks
 
 
+@pytest.mark.skipif(
+    not is_sm100_or_sm110_supported(),
+    reason="flash_attn.cute implements qv on SM100/SM110 only (not SM120).",
+)
+@pytest.mark.parametrize("mha_type", ["mqa", "gqa"])
+@pytest.mark.parametrize(
+    "seqlen_q,seqlen_k",
+    [
+        (1, 128),  # plain decode
+        (4, 1024),  # speculative decode (multiple q rows per request)
+        (64, 800),  # chunked extend
+        (16, 20000),  # long context
+    ],
+)
+def test_flash_attn_varlen_qv_deepseek_absorbed(seqlen_q, seqlen_k, mha_type):
+    """DeepSeek absorbed-MLA FA4 shape: rope q/k head_dim 64, latent v/qv
+    head_dim 512, varlen q over a paged KV cache, num_splits=1. Mirrors the
+    production calls in flashattention_backend.py, where extend
+    (flash_attn_varlen_func) and decode (flash_attn_with_kvcache) share this
+    qv-threaded path.
+    """
+    device = "cuda"
+    dtype = torch.bfloat16
+    torch.random.manual_seed(seqlen_q + seqlen_k)
+    batch_size = 5
+    nheads = 8
+    nheads_k = 1 if mha_type == "mqa" else 4
+    d, dv = 64, 512
+    page_size = 128
+
+    q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype)
+    qv = torch.randn(batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype)
+    k_cache, v_cache, page_table, k_cache_paged, v_cache_paged, _ = (
+        _generate_block_kvcache(
+            seqlen_k, page_size, batch_size, nheads_k, d, dv, device, dtype, dtype
+        )
+    )
+    cache_seqlens = torch.randint(
+        seqlen_q, seqlen_k + 1, (batch_size,), dtype=torch.int32, device=device
+    )
+    cache_seqlens[0] = seqlen_k
+    cu_seqlens_q = (
+        torch.arange(batch_size + 1, dtype=torch.int32, device=device) * seqlen_q
+    )
+
+    out_unpad = flash_attn_varlen_func(
+        rearrange(q, "b s h d -> (b s) h d"),
+        k_cache_paged,
+        v_cache_paged,
+        qv=rearrange(qv, "b s h d -> (b s) h d"),
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=None,  # KV comes from the paged cache via seqused_k
+        seqused_k=cache_seqlens,
+        page_table=page_table,
+        causal=True,
+        num_splits=1,
+        ver=4,
+    )
+    out = rearrange(out_unpad, "(b s) h d -> b s h d", b=batch_size)
+
+    # Decode enters through the flash_attn_with_kvcache wrapper; it must
+    # thread qv/num_splits down to the same varlen kernel call bit-for-bit.
+    out_kvcache = flash_attn_with_kvcache(
+        q=rearrange(q, "b s h d -> (b s) h d"),
+        k_cache=k_cache_paged,
+        v_cache=v_cache_paged,
+        qv=rearrange(qv, "b s h d -> (b s) h d"),
+        page_table=page_table,
+        cache_seqlens=cache_seqlens,
+        cu_seqlens_q=cu_seqlens_q,
+        max_seqlen_q=seqlen_q,
+        causal=True,
+        num_splits=1,
+        ver=4,
+    )
+    assert torch.equal(out_kvcache, out_unpad)
+
+    key_padding_mask = rearrange(
+        torch.arange(seqlen_k, device=device), "s -> 1 s"
+    ) < rearrange(cache_seqlens, "b -> b 1")
+    k_rep = repeat(k_cache, "b s h d -> b s (h g) d", g=nheads // nheads_k)
+    v_rep = repeat(v_cache, "b s h d -> b s (h g) d", g=nheads // nheads_k)
+    out_ref, _ = attention_ref(
+        q, k_rep, v_rep, None, key_padding_mask, causal=True, qv=qv
+    )
+    out_pt, _ = attention_ref(
+        q,
+        k_rep,
+        v_rep,
+        None,
+        key_padding_mask,
+        causal=True,
+        qv=qv,
+        upcast=False,
+        reorder_ops=True,
+    )
+
+    print(f"Output max diff: {(out - out_ref).abs().max().item()}")
+    print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}")
+    assert (out - out_ref).abs().max().item() <= 2 * (
+        out_pt - out_ref
+    ).abs().max().item() + 1e-5
+    assert (out - out_ref).abs().mean().item() <= 1.5 * (
+        out_pt - out_ref
+    ).abs().mean().item()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py

@@ -67,6 +67,18 @@ def _filter(batch: ForwardBatch, *, lo: int, hi: int) -> ForwardBatch:
 
 
 class TestTboFilterBatchMarker(CustomTestCase):
+    def test_filter_batch_clears_mlp_sync_unpad_fields_on_children(self):
+        # MLP-sync padding records _original_batch_size/_original_num_tokens
+        # before TBO splits the batch (prepare_mlp_sync_batch pads first, then
+        # runs TboForwardBatchPreparer); children carry no restore state — the
+        # parent performs the post-forward unpad.
+        parent = _make_target_verify_batch(8)
+        parent._original_batch_size = 8
+        parent._original_num_tokens = 8
+        child = _filter(parent, lo=0, hi=4)
+        self.assertIsNone(child._original_batch_size)
+        self.assertIsNone(child._original_num_tokens)
+
     def test_filter_batch_resets_plan_marker_on_children(self):
         child = _filter(_make_target_verify_batch(8), lo=0, hi=4)
         self.assertEqual(child.batch_size, 4)

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()

test/registered/unit/spec/test_eagle_seeded_coins.py

@@ -0,0 +1,163 @@
+"""Unit tests for the EAGLE verify coins (_verify_coins / _seeded_verify_coins).
+
+Locks the deterministic-coin contract behind seeded speculative sampling:
+identical (seed, seq_lens) inputs produce bitwise-identical coins, distinct
+seeds diverge, the column split (first draft_token_num columns -> rejection
+coins, last column -> final-sampling coin) holds, unseeded requests keep
+torch.rand, and the float32 conversion never emits a coin of exactly 1.0
+(the sampling kernels expect half-open [0, 1) coins — a 1.0 coin walks past
+the final CDF bucket and can return a zero-probability token).
+
+Requires a GPU: the coins hash through the murmur_hash32 Triton kernel.
+"""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.speculative.eagle_utils import _seeded_verify_coins, _verify_coins
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
+
+DRAFT_TOKEN_NUM = 4
+
+
+def _coins(seeds, seq_lens):
+    device = "cuda"
+    return _seeded_verify_coins(
+        sampling_seed=torch.tensor(seeds, device=device, dtype=torch.int64),
+        seq_lens=torch.tensor(seq_lens, device=device, dtype=torch.int64),
+        draft_token_num=DRAFT_TOKEN_NUM,
+        device=device,
+    )
+
+
+class TestSeededVerifyCoins(CustomTestCase):
+    def test_seeded_coins_are_reproducible(self):
+        coins_a, final_a = _coins([12345, 67890, 12345], [7, 9, 7])
+        coins_b, final_b = _coins([12345, 67890, 12345], [7, 9, 7])
+
+        self.assertEqual(coins_a.shape, (3, DRAFT_TOKEN_NUM))
+        self.assertEqual(final_a.shape, (3,))
+        self.assertTrue(torch.equal(coins_a, coins_b))
+        self.assertTrue(torch.equal(final_a, final_b))
+        # Same (seed, seq_len) pair hashes to the same coins regardless of row.
+        self.assertTrue(torch.equal(coins_a[0], coins_a[2]))
+        self.assertEqual(final_a[0].item(), final_a[2].item())
+        # Coins live in [0, 1).
+        self.assertTrue(bool((coins_a >= 0).all() and (coins_a < 1).all()))
+        self.assertTrue(bool((final_a >= 0).all() and (final_a < 1).all()))
+
+    def test_distinct_seeds_or_positions_diverge(self):
+        coins, final = _coins([12345, 67890, 12345], [7, 9, 11])
+        self.assertFalse(torch.equal(coins[0], coins[1]))  # different seed
+        self.assertFalse(torch.equal(coins[0], coins[2]))  # different seq_len
+
+    def test_column_split_maps_rejection_then_final(self):
+        # Structured hash: hashed[i, j] = i * 1000 + j, so each coin names its
+        # (row, column) origin. Locks the column-space contract: columns
+        # [0, draft_token_num) become the per-draft rejection coins and column
+        # draft_token_num becomes the final-sampling coin.
+        umax = torch.iinfo(torch.uint32).max
+
+        def _structured_hash(seed, positions, col_indices):
+            rows = torch.arange(seed.shape[0], device=seed.device).unsqueeze(1)
+            return (rows * 1000 + col_indices.unsqueeze(0)).to(torch.uint32)
+
+        with patch(
+            "sglang.kernels.ops.sampling.murmur_hash.murmur_hash32",
+            side_effect=_structured_hash,
+        ):
+            coins, final = _coins([1, 2], [3, 4])
+
+        def _expected(row, col):
+            return (
+                torch.tensor(row * 1000 + col, dtype=torch.float64)
+                .div(umax)
+                .to(torch.float32)
+                .item()
+            )
+
+        for row in range(2):
+            for col in range(DRAFT_TOKEN_NUM):
+                self.assertEqual(coins[row, col].item(), _expected(row, col))
+            self.assertEqual(final[row].item(), _expected(row, DRAFT_TOKEN_NUM))
+
+    def test_unseeded_requests_keep_torch_rand(self):
+        device = "cuda"
+        kwargs = dict(
+            sampling_info=SimpleNamespace(sampling_seed=None),
+            seq_lens=torch.tensor([3, 4, 5], device=device, dtype=torch.int64),
+            draft_token_num=DRAFT_TOKEN_NUM,
+            candidates=torch.zeros(
+                (3, DRAFT_TOKEN_NUM), device=device, dtype=torch.int64
+            ),
+            device=device,
+        )
+        with patch(
+            "sglang.kernels.ops.sampling.murmur_hash.murmur_hash32"
+        ) as mock_hash:
+            coins_a, final_a = _verify_coins(**kwargs)
+            coins_b, final_b = _verify_coins(**kwargs)
+
+        mock_hash.assert_not_called()
+        self.assertEqual(coins_a.shape, (3, DRAFT_TOKEN_NUM))
+        self.assertEqual(final_a.shape, (3,))
+        self.assertEqual(coins_a.dtype, torch.float32)
+        # torch.rand draws: two calls must not repeat.
+        self.assertFalse(torch.equal(coins_a, coins_b))
+        self.assertFalse(torch.equal(final_a, final_b))
+
+    def test_seeded_requests_dispatch_to_seeded_coins(self):
+        device = "cuda"
+        seeds = torch.tensor([12345, 67890], device=device, dtype=torch.int64)
+        seq_lens = torch.tensor([7, 9], device=device, dtype=torch.int64)
+        coins, final = _verify_coins(
+            sampling_info=SimpleNamespace(sampling_seed=seeds),
+            seq_lens=seq_lens,
+            draft_token_num=DRAFT_TOKEN_NUM,
+            candidates=torch.zeros(
+                (2, DRAFT_TOKEN_NUM), device=device, dtype=torch.int64
+            ),
+            device=device,
+        )
+        expected_coins, expected_final = _seeded_verify_coins(
+            sampling_seed=seeds,
+            seq_lens=seq_lens,
+            draft_token_num=DRAFT_TOKEN_NUM,
+            device=device,
+        )
+        self.assertTrue(torch.equal(coins, expected_coins))
+        self.assertTrue(torch.equal(final, expected_final))
+
+    def test_max_hash_clamps_coins_below_one(self):
+        # The top 129 uint32 hashes round to exactly 1.0 under the float32
+        # cast; force the worst case and assert the clamp holds the contract.
+        umax = torch.iinfo(torch.uint32).max
+
+        def _all_max_hash(seed, positions, col_indices):
+            return torch.full(
+                (seed.shape[0], col_indices.shape[0]),
+                umax,
+                dtype=torch.uint32,
+                device=seed.device,
+            )
+
+        with patch(
+            "sglang.kernels.ops.sampling.murmur_hash.murmur_hash32",
+            side_effect=_all_max_hash,
+        ):
+            coins, final = _coins([1, 2], [3, 4])
+
+        self.assertTrue(bool((coins < 1).all()))
+        self.assertTrue(bool((final < 1).all()))
+        # Clamped to the largest float32 strictly below one.
+        self.assertEqual(coins.max().item(), 1.0 - 2**-24)
+
+
+if __name__ == "__main__":
+    unittest.main()