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

python/sglang/kernels/ops/attention/flash_attention.py

@@ -297,6 +297,7 @@ def flash_attn_varlen_func(
             cu_seqlens_k,
             max_seqlen_q=max_seqlen_q,
             max_seqlen_k=max_seqlen_k,
+            qv=qv,
             seqused_q=seqused_q,
             seqused_k=seqused_k,
             page_table=page_table,

python/sglang/kernels/ops/attention/flash_attention_v4.py

@@ -34,6 +34,7 @@ def flash_attn_varlen_func(
     v: torch.Tensor,
     cu_seqlens_q: Optional[torch.Tensor] = None,
     cu_seqlens_k: Optional[torch.Tensor] = None,
+    qv: Optional[torch.Tensor] = None,
     seqused_q: Optional[torch.Tensor] = None,
     seqused_k: Optional[torch.Tensor] = None,
     max_seqlen_q: Optional[int] = None,
@@ -75,7 +76,7 @@ def flash_attn_varlen_func(
             "vendored FA4 package is importable."
         ) from _flash_attn_import_error
 
-    q, k, v = [_maybe_contiguous(t) for t in (q, k, v)]
+    q, k, v, qv = [_maybe_contiguous(t) for t in (q, k, v, qv)]
     cu_seqlens_q, cu_seqlens_k = [
         _maybe_contiguous(t) for t in (cu_seqlens_q, cu_seqlens_k)
     ]
@@ -118,6 +119,7 @@ def flash_attn_varlen_func(
         q=q,
         k=k,
         v=v,
+        qv=qv,
         cu_seqlens_q=cu_seqlens_q,
         cu_seqlens_k=cu_seqlens_k,
         seqused_q=seqused_q,
@@ -189,7 +191,7 @@ def flash_attn_with_kvcache(
     return_softmax_lse: bool = False,
     **_: object,
 ):
-    if k is not None or v is not None or qv is not None:
+    if k is not None or v is not None:
         raise NotImplementedError("FA4 does not support updating KV cache in-place.")
     if rotary_cos is not None or rotary_sin is not None or rotary_seqlens is not None:
         raise NotImplementedError("FA4 path does not support rotary embedding.")
@@ -206,6 +208,7 @@ def flash_attn_with_kvcache(
         q=q,
         k=k_cache,
         v=v_cache,
+        qv=qv,
         cu_seqlens_q=cu_seqlens_q,
         seqused_k=cache_seqlens,
         max_seqlen_q=max_seqlen_q,

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -783,6 +783,7 @@ class TboForwardBatchPreparer:
                 original_global_num_tokens_cpu=None,
                 _original_batch_size=None,
                 _original_forward_mode=None,
+                _original_num_tokens=None,
                 global_num_tokens_gpu=None,
                 global_num_tokens_cpu=None,
                 global_dp_buffer_len=global_dp_buffer_len,

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -1535,7 +1535,10 @@ class FlashAttentionBackend(AttentionBackend):
                     return output, lse
                 return output
             else:
-                assert self.fa_impl_ver == 3, "Only FA3 support here"
+                # FA4 absorbed MLA is shared by extend and decode: once qv is
+                # threaded through the wrappers, decode's flash_attn_with_kvcache
+                # call takes the same qv/ver arguments as this extend path.
+                assert self.fa_impl_ver in (3, 4), "Only FA3/FA4 support here"
                 # Do absorbed multi-latent attention
                 kv_cache = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(
                     q.dtype

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -115,13 +115,12 @@ class DeepEPMoE(FusedMoE):
             and self.w13_weight.dtype == torch.bfloat16
             and get_moe_runner_backend().is_deep_gemm()
             and get_moe_a2a_backend().is_deepep()
-            and get_deepep_mode().enable_low_latency()
             and not _is_npu
             and not _is_hip
         ):
             assert (
                 deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
-            ), "Unquantized DeepEP low-latency MoE requires DeepGEMM BF16"
+            ), "Unquantized DeepEP MoE requires DeepGEMM BF16"
             self.deprecate_flag = True
         else:
             self.deprecate_flag = False

python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -25,6 +25,7 @@ from sglang.srt.layers.moe.moe_runner.base import (
     register_pre_permute,
 )
 from sglang.srt.layers.moe.utils import MoeRunnerBackend
+from sglang.srt.runtime_context import get_exec
 from sglang.srt.utils import (
     ceil_div,
     dispose_tensor,
@@ -828,7 +829,15 @@ def pre_permute_deepep_normal_to_deep_gemm(
     running_state["topk_ids"] = topk_ids
     running_state["topk_weights"] = topk_weights
 
-    input_tensor = torch.empty(
+    # Deterministic inference zero-fills the scatter buffers: expert-alignment
+    # padding leaves slots that ep_scatter never writes, and pad garbage in
+    # input_tensor would leak batch-dependent values into the grouped GEMM.
+    # The scale buffer only matters for FP8 activations sharing this
+    # pre-permute (ep_scatter skips scales entirely for BF16 dispatch).
+    deterministic = get_exec().deterministic.enable_deterministic_inference
+    buffer_init = torch.zeros if deterministic else torch.empty
+
+    input_tensor = buffer_init(
         (all_tokens, K),
         device=hidden_states.device,
         dtype=hidden_states.dtype,
@@ -841,12 +850,12 @@ def pre_permute_deepep_normal_to_deep_gemm(
             dtype=torch.int,
         ).transpose(0, 1)
     else:
-        input_tensor_scale = torch.empty(
+        input_tensor_scale = buffer_init(
             (all_tokens, K // 128),
             device=hidden_states.device,
             dtype=torch.float32,
         )
-    m_indices = torch.empty(all_tokens, device=hidden_states.device, dtype=torch.int32)
+    m_indices = buffer_init(all_tokens, device=hidden_states.device, dtype=torch.int32)
     output_index = torch.empty_like(topk_ids)
 
     if get_offloader().forbid_copy_engine_usage:

python/sglang/srt/layers/quantization/unquant.py

@@ -19,7 +19,6 @@ from sglang.srt.layers.moe import (
     MoeRunner,
     MoeRunnerBackend,
     MoeRunnerConfig,
-    get_deepep_mode,
     get_moe_a2a_backend,
     get_moe_runner_backend,
 )
@@ -367,7 +366,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
             self.use_deep_gemm
             and layer.w13_weight.dtype == torch.bfloat16
             and get_moe_a2a_backend().is_deepep()
-            and get_deepep_mode().enable_low_latency()
             and not _is_npu
             and not _is_hip
             and hasattr(layer, "dispatcher")

python/sglang/srt/layers/sampler.py

@@ -186,6 +186,21 @@ class Sampler(nn.Module):
                 # Standard path: do softmax and sample from probs.
                 logits.div_(sampling_info.temperatures)
 
+                # Deterministic inference must derive the returned logprobs
+                # from F.log_softmax — the same kernel prefill rescoring uses —
+                # not log(softmax(x)) below: the two disagree at ~1e-6 despite
+                # being mathematically equivalent, which breaks bitwise
+                # prefill/decode logprob alignment.
+                if (
+                    return_logprob
+                    and self.enable_deterministic
+                    and logprobs_via_logsoftmax_kernel is None
+                    and not SGLANG_RETURN_ORIGINAL_LOGPROB
+                ):
+                    logprobs_via_logsoftmax_kernel = torch.nn.functional.log_softmax(
+                        logits, dim=-1
+                    )
+
                 # In-place op to save memory
                 logits[:] = torch.softmax(logits, dim=-1)
                 probs = logits

python/sglang/srt/model_executor/forward_batch_info.py

@@ -501,6 +501,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     original_global_num_tokens_cpu: Optional[List[int]] = None
     _original_batch_size: Optional[int] = None
     _original_forward_mode: Optional[ForwardMode] = None
+    _original_num_tokens: Optional[int] = None
     global_num_tokens_cpu: Optional[List[int]] = None
     global_num_tokens_gpu: Optional[torch.Tensor] = None
     # Has to be None when cuda graph is captured.
@@ -1382,6 +1383,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
     def _pad_inputs_to_size(self, model_runner: ModelRunner, num_tokens, bs):
         # padding
+        self._original_num_tokens = self.positions.shape[0]
         self.input_ids = self._pad_tensor_to_size(self.input_ids, num_tokens)
         self.req_pool_indices = self._pad_tensor_to_size(self.req_pool_indices, bs)
         if self.lora_ids is not None:
@@ -1492,6 +1494,17 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
             self.batch_size = self._original_batch_size
         bs = self.batch_size
 
+        # MLP-sync padding appended dummy rows after the real ones; slice the
+        # per-request tensors back so post-forward consumers (seeded sampling,
+        # ngram token-table updates) never see the padding. The draft-decode
+        # branch below does the same for speculative batches.
+        if self.spec_info is None and self._original_num_tokens is not None:
+            self.positions = self.positions[: self._original_num_tokens]
+            self.seq_lens = self.seq_lens[:bs]
+            self.req_pool_indices = self.req_pool_indices[:bs]
+            if self.seq_lens_cpu is not None:
+                self.seq_lens_cpu = self.seq_lens_cpu[:bs]
+
         if self.spec_info is not None:
             if self.forward_mode.is_decode():  # draft
                 num_tokens = self.hidden_states_backup.shape[0]

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -11,8 +11,8 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods
     AttnForwardMethod,
 )
 from sglang.srt.models.deepseek_common.utils import _is_hip
-from sglang.srt.runtime_context import get_server_args
-from sglang.srt.utils import use_intel_amx_backend
+from sglang.srt.runtime_context import get_exec, get_server_args
+from sglang.srt.utils import is_sm100_or_sm110_supported, use_intel_amx_backend
 
 MHA_ONE_SHOT_SUPPORTED_BACKENDS = ["fa3", "flashinfer", "flashmla"]
 
@@ -133,8 +133,15 @@ def handle_attention_cutlass_mla(attn, forward_batch):
 
 
 def handle_attention_fa4(attn, forward_batch):
-    # TODO(cicirori): use FA4 MHA for DeepSeekV3 for now
-    return AttnForwardMethod.MHA_CHUNKED_KV
+    # FA4 absorbed MLA feeds q_nope through the qv argument, which
+    # flash_attn.cute only implements on SM100/SM110 (not SM120); keep the
+    # pre-existing MHA chunked-KV path elsewhere. Deterministic inference
+    # requires MLA and rejects fa4 on other archs at startup (server_args).
+    if not is_sm100_or_sm110_supported():
+        return AttnForwardMethod.MHA_CHUNKED_KV
+    if get_exec().deterministic.enable_deterministic_inference:
+        return _dispatch_mla_subtype(attn, forward_batch)
+    return _handle_attention_backend(attn, forward_batch, "fa4")
 
 
 def handle_attention_trtllm_mla(attn, forward_batch):

python/sglang/srt/models/deepseek_common/utils.py

@@ -59,6 +59,7 @@ NVFP4_CKPT_FP8_ATTN_QUANT_MODULES = ["q_b_proj"]
 
 FORWARD_ABSORB_CORE_ATTENTION_BACKENDS = [
     "fa3",
+    "fa4",
     "dsa",
     "nsa",  # Deprecated alias for "dsa"
     "flashinfer",

python/sglang/srt/models/deepseek_v2.py

@@ -184,6 +184,7 @@ from sglang.srt.models.deepseek_common.utils import (
     is_wint4afp8_or_wint4a16_config,
 )
 from sglang.srt.runtime_context import (
+    get_exec,
     get_flags,
     get_forward,
     get_model,
@@ -1783,6 +1784,7 @@ class DeepseekV2AttentionMLA(
         )
         self.use_min_latency_fused_a_gemm = (
             self.has_fused_proj
+            and not get_exec().deterministic.enable_deterministic_inference
             and not self.is_packed_weight
             and fused_a_gemm_weight_eligible(self.fused_qkv_a_proj_with_mqa)
         )

python/sglang/srt/server_args.py

@@ -82,6 +82,7 @@ from sglang.srt.utils.common import (
     is_npu,
     is_remote_url,
     is_sm90_supported,
+    is_sm100_or_sm110_supported,
     is_sm100_supported,
     is_sm120_supported,
     is_xpu,
@@ -7379,10 +7380,19 @@ class ServerArgs:
 
             attention_backend = resolved_view(self).attention_backend
             if is_deepseek_model:
-                deepseek_deterministic_attention_backends = ["fa3", "triton"]
-                if attention_backend not in deepseek_deterministic_attention_backends:
+                if (
+                    attention_backend
+                    not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND
+                ):
                     raise ValueError(
-                        f"Currently only {deepseek_deterministic_attention_backends} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
+                        f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
+                    )
+                if attention_backend == "fa4" and not is_sm100_or_sm110_supported():
+                    raise ValueError(
+                        "Deterministic inference with DeepSeek models on the fa4 "
+                        "attention backend requires SM100/SM110: it runs "
+                        "absorbed MLA, whose qv argument flash_attn.cute only "
+                        "implements on those archs."
                     )
 
             if attention_backend not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND:

python/sglang/srt/speculative/eagle_utils.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import logging
 import math
 from enum import IntEnum
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
 
@@ -33,6 +33,7 @@ if TYPE_CHECKING:
     from sglang.srt.managers.tp_worker import TpModelWorker
     from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
     from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
     from sglang.srt.speculative.eagle_info import EagleVerifyInput
 
 _is_cuda = is_cuda()
@@ -566,6 +567,76 @@ def eagle_prepare_for_verify(
     return verify_forward_batch, can_run_cuda_graph
 
 
+def _seeded_verify_coins(
+    *,
+    sampling_seed: torch.Tensor,
+    seq_lens: torch.Tensor,
+    draft_token_num: int,
+    device,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Derive deterministic verify-side coins from per-request sampling seeds.
+
+    Mirrors the main seeded-sampling path: murmur_hash32(seed, seq_lens,
+    column) mapped to [0, 1). Columns [0, draft_token_num) drive the
+    per-draft rejection coins; column draft_token_num drives the final
+    fallback-sampling coin.
+
+    Scope: this seeds only the verify-side RNG. With rejection sampling the
+    draft workers still pick candidates via unseeded multinomial
+    (fast_sample in eagle_worker_v2), so that mode stays non-deterministic
+    until the draft RNG is seeded in a follow-up; top-k/greedy draft
+    selection is already deterministic.
+    """
+    from sglang.kernels.ops.sampling.murmur_hash import murmur_hash32
+
+    cols = torch.arange(draft_token_num + 1, device=device, dtype=torch.int64)
+    hashed = murmur_hash32(
+        sampling_seed.to(torch.uint64), seq_lens.to(torch.uint64), cols
+    )
+    uniforms = hashed.to(torch.float64) / torch.iinfo(torch.uint32).max
+    # The float32 cast rounds the top 129 uint32 hashes to exactly 1.0, but
+    # the sampling kernels expect half-open [0, 1) coins: a 1.0 coin walks
+    # past the last CDF bucket and can return a zero-probability token.
+    # Clamp to the largest float32 below one; every other coin value is
+    # untouched, so previously verified bitwise baselines stay intact.
+    max_coin = 1.0 - 2**-24
+    coins = (
+        uniforms[:, :draft_token_num].to(torch.float32).clamp_(max=max_coin)
+    ).contiguous()
+    coins_for_final_sampling = (
+        uniforms[:, draft_token_num].to(torch.float32).clamp_(max=max_coin)
+    ).contiguous()
+    return coins, coins_for_final_sampling
+
+
+def _verify_coins(
+    *,
+    sampling_info: SamplingBatchInfo,
+    seq_lens: torch.Tensor,
+    draft_token_num: int,
+    candidates: torch.Tensor,
+    device,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Rejection and final-sampling coins for verify: deterministic seeded
+    coins when sampling_seed is set (see _seeded_verify_coins), torch.rand
+    otherwise.
+    """
+    if sampling_info.sampling_seed is not None:
+        return _seeded_verify_coins(
+            sampling_seed=sampling_info.sampling_seed,
+            seq_lens=seq_lens,
+            draft_token_num=draft_token_num,
+            device=device,
+        )
+    # coins for rejection sampling
+    coins = torch.rand_like(candidates, dtype=torch.float32, device=device)
+    # coins for final sampling
+    coins_for_final_sampling = torch.rand(
+        (candidates.shape[0],), dtype=torch.float32, device=device
+    )
+    return coins, coins_for_final_sampling
+
+
 def eagle_sample(
     verify_input: EagleVerifyInput,
     batch: ScheduleBatch,
@@ -717,10 +788,13 @@ def eagle_sample(
                 "does not produce one (draft_probs missing or vocab-mismatched)."
             )
 
-        # coins for rejection sampling
-        coins = torch.rand_like(candidates, dtype=torch.float32, device=device)
-        # coins for final sampling
-        coins_for_final_sampling = torch.rand((bs,), dtype=torch.float32, device=device)
+        coins, coins_for_final_sampling = _verify_coins(
+            sampling_info=sampling_info,
+            seq_lens=batch.seq_lens,
+            draft_token_num=verify_input.draft_token_num,
+            candidates=candidates,
+            device=device,
+        )
 
         sampling_fn = (
             chain_speculative_sampling_triton

python/sglang/srt/utils/common.py

@@ -288,6 +288,15 @@ is_sm100_supported = lru_cache(maxsize=1)(
         _check_cuda_device_version, device_capability_majors=[10], cuda_version=(12, 8)
     )
 )
+# Datacenter Blackwell (SM100) plus SM110; excludes consumer Blackwell (SM120).
+# This is the arch set flash_attn.cute accepts for the absorbed-MLA qv argument.
+is_sm100_or_sm110_supported = lru_cache(maxsize=1)(
+    partial(
+        _check_cuda_device_version,
+        device_capability_majors=[10, 11],
+        cuda_version=(12, 8),
+    )
+)
 is_sm80_supported = lru_cache(maxsize=1)(
     partial(
         _check_cuda_device_version, device_capability_majors=[8], cuda_version=(11, 0)

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