diff --git a/python/sglang/kernels/ops/attention/flash_attention.py b/python/sglang/kernels/ops/attention/flash_attention.py index c5e77b6f8..cfd8ca4cb 100644 --- a/python/sglang/kernels/ops/attention/flash_attention.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/flash_attention_v4.py b/python/sglang/kernels/ops/attention/flash_attention_v4.py index 62d86a89e..456036a58 100644 --- a/python/sglang/kernels/ops/attention/flash_attention_v4.py +++ b/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, diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index ca44b0021..8352086bb 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index ac440dbab..f552987d1 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 88028dd87..1d1f40c77 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index 5d160b282..7a3278465 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/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: diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 961b81e9e..725f4d6bd 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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") diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index de587ec5f..e8a1d7161 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 7e7508cc5..d6d9bd08d 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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] diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index 6e24068a1..f55634729 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/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): diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index eef411711..a7d374310 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/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", diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b914dccae..ba3b1c504 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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) ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1fbc56bb9..710011a6f 100644 --- a/python/sglang/srt/server_args.py +++ b/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: diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index c4725ebcc..97592cc99 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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 diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 8d0c33022..f99242d3f 100644 --- a/python/sglang/srt/utils/common.py +++ b/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) diff --git a/test/registered/kernels/ops/attention/test_flash_attention_4.py b/test/registered/kernels/ops/attention/test_flash_attention_4.py index 4e7195e82..5a16481f5 100644 --- a/test/registered/kernels/ops/attention/test_flash_attention_4.py +++ b/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"])) diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index 5a46125ef..efc8e9707 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/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) diff --git a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py new file mode 100644 index 000000000..b89b8d77e --- /dev/null +++ b/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() diff --git a/test/registered/unit/spec/test_eagle_seeded_coins.py b/test/registered/unit/spec/test_eagle_seeded_coins.py new file mode 100644 index 000000000..a089ca456 --- /dev/null +++ b/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()