diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index b157eebea28b..45ff0e203286 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh @@ -505,10 +505,12 @@ inline PrefillPlan plan_compress_prefill( const auto f2s_ptr = static_cast(full_to_state.data_ptr()); const auto batch_size = static_cast(B.unwrap()); - constexpr auto kMaxTokens = static_cast(std::numeric_limits::max()); + // ragged_id is a zero-based uint16 index, so a 64K-token batch is valid. + constexpr auto kMaxTokens = static_cast(std::numeric_limits::max()) + 1; RuntimeCheck(compress_ratio == 4 || compress_ratio == 128); RuntimeCheck(!use_req_ring || compress_ratio == 4); - RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens); + // Keep batch_id below 65535: pack_w(65535, 65535, ...) is the invalid sentinel. + RuntimeCheck(batch_size < kMaxTokens && batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens); // `swa_page_size` >= `ring_size` >= `compress_ratio` RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0); // Write pad: trailing tokens kept resident so a verify batch's committed tail survives @@ -750,9 +752,9 @@ inline PrefillPlan plan_compress_prefill_legacy( const auto window_size = compress_ratio * (is_overlap ? 2 : 1); const auto batch_size = static_cast(B.unwrap()); - constexpr auto kMaxTokens = static_cast(std::numeric_limits::max()); + constexpr auto kMaxTokens = static_cast(std::numeric_limits::max()) + 1; RuntimeCheck(compress_ratio == 4 || compress_ratio == 128); - RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens); + RuntimeCheck(batch_size < kMaxTokens && batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens); uint32_t counter = 0; uint32_t counter_c = 0; diff --git a/python/sglang/kernels/ops/attention/dsv4/q_rope_store.py b/python/sglang/kernels/ops/attention/dsv4/q_rope_store.py index e89af4e47967..a539cc529d0f 100644 --- a/python/sglang/kernels/ops/attention/dsv4/q_rope_store.py +++ b/python/sglang/kernels/ops/attention/dsv4/q_rope_store.py @@ -21,6 +21,43 @@ def _q_rope_store(X, Y, F, POS, SX: tl.constexpr, SY: tl.constexpr): tl.store(Y + row * SY + head * 512 + r, tl.where(r >= 448, rotated, value)) +@triton.jit +def _q_rope_store_prefill( + X, + Y, + F, + POS, + M, + H: tl.constexpr, + SX: tl.constexpr, + SY: tl.constexpr, + BLOCK_HEADS: tl.constexpr, +): + # Keep token count dynamic to avoid compiling every prefill batch length. + heads = tl.program_id(0) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS) + row, head = heads // H, heads % H + r = tl.arange(0, 512) + # The padded output can exceed 2 GiB at the 64K prefill ceiling. + x_offset = row[:, None].to(tl.int64) * SX + head[:, None] * 512 + r[None, :] + value = tl.load(X + x_offset, row[:, None] < M, 0).to(tl.float32) + partner = tl.gather(value, tl.broadcast_to((r ^ 1)[None, :], (BLOCK_HEADS, 512)), 1) + position = tl.load(POS + row, row < M, 0) + freq_offset = position[:, None].to(tl.int64) * 64 + (r[None, :] - 448) // 2 * 2 + rope_mask = (row[:, None] < M) & (r[None, :] >= 448) + cos = tl.load(F + freq_offset, rope_mask, 0) + sin = tl.load(F + freq_offset + 1, rope_mask, 0) + # Keep the same arithmetic and BF16 rounding as the decode kernel. + even = tl.fma(value, cos, -partner * sin) + odd = tl.fma(partner, sin, value * cos) + rotated = tl.where((r[None, :] & 1) == 0, even, odd) + y_offset = row[:, None].to(tl.int64) * SY + head[:, None] * 512 + r[None, :] + tl.store( + Y + y_offset, + tl.where(r[None, :] >= 448, rotated, value), + row[:, None] < M, + ) + + def q_rope_store( q: torch.Tensor, output: torch.Tensor, @@ -35,6 +72,20 @@ def q_rope_store( assert freqs_cis.dtype == torch.complex64 and freqs_cis.is_contiguous() assert freqs_cis.shape[1] == 32 and positions.shape == (q.shape[0],) assert positions.dtype in (torch.int32, torch.int64) and positions.is_contiguous() + if q.shape[0] >= 4096 and q.shape[1] == 16: + _q_rope_store_prefill[(triton.cdiv(q.shape[0] * q.shape[1], 4),)]( + q, + output, + torch.view_as_real(freqs_cis), + positions, + q.shape[0], + q.shape[1], + q.stride(0), + output.stride(0), + BLOCK_HEADS=4, + num_warps=4, + ) + return _q_rope_store[(q.shape[0], q.shape[1])]( q, output, diff --git a/python/sglang/kernels/ops/layernorm/hc_combine_norm.py b/python/sglang/kernels/ops/layernorm/hc_combine_norm.py index e6c1c44f8e67..9234fbc67267 100644 --- a/python/sglang/kernels/ops/layernorm/hc_combine_norm.py +++ b/python/sglang/kernels/ops/layernorm/hc_combine_norm.py @@ -22,16 +22,40 @@ def _hc_combine_norm(X, P, W, Y, SX: tl.constexpr, SP: tl.constexpr, EPS: tl.con tl.store(Y + row * 5120 + h, value * inv_rms * weight, mask) +@triton.jit +def _hc_combine_norm_prefill( + X, P, W, Y, SX: tl.constexpr, SP: tl.constexpr, EPS: tl.constexpr +): + # Large batches have enough rows to use one CTA per row without repeating + # the combine and RMS reduction for each output partition. + row = tl.program_id(0).to(tl.int64) + h = tl.arange(0, 8192) + value = tl.full((8192,), 0, tl.float32) + for c in tl.static_range(4): + pre = tl.load(P + row * SP + c).to(tl.float32) + x = tl.load(X + row * SX + c * 5120 + h, h < 5120, 0).to(tl.float32) + value += x * pre + value = value.to(tl.bfloat16).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(value * value, 0) / 5120 + EPS) + weight = tl.load(W + h, h < 5120, 0).to(tl.float32) + tl.store(Y + row * 5120 + h, value * inv_rms * weight, h < 5120) + + def hc_combine_norm( x: torch.Tensor, pre: torch.Tensor, weight: torch.Tensor, eps: float ) -> torch.Tensor: - """Fuse four-stream combine and RMSNorm for small BF16 batches of width 5120.""" + """Fuse four-stream combine and RMSNorm for BF16 batches of width 5120.""" m = x.shape[0] - assert 0 < m <= 8 and x.shape == (m, 20480) + assert (0 < m <= 8 or 4096 <= m <= 65536) and x.shape == (m, 20480) assert pre.shape == (m, 4) and pre.stride(1) == 1 assert weight.shape == (5120,) and weight.is_contiguous() assert x.dtype == weight.dtype == torch.bfloat16 and x.stride(1) == 1 y = torch.empty((m, 5120), dtype=x.dtype, device=x.device) + if m >= 4096: + _hc_combine_norm_prefill[(m,)]( + x, pre, weight, y, x.stride(0), pre.stride(0), eps, num_warps=4 + ) + return y # Four CTAs per row trade redundant statistics for more concurrent loads # when only a few speculative tokens are being processed. _hc_combine_norm[(m, 4)]( diff --git a/python/sglang/kernels/ops/layernorm/hc_mix_stats_bf16x3.py b/python/sglang/kernels/ops/layernorm/hc_mix_stats_bf16x3.py new file mode 100644 index 000000000000..f6fe6000cd13 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/hc_mix_stats_bf16x3.py @@ -0,0 +1,100 @@ +"""Compensated mHC prefill projection with a shared activation load. + +Keep three BF16 components of the FP32 weights and accumulate their products +separately. The 16 fixed K slices bound FP32 accumulation error, as in the +compensated DeepGEMM path, while avoiding its second activation read/reduction. +""" + +import torch +import triton +import triton.language as tl + + +def split_bf16_hc_weight(weight: torch.Tensor): + assert weight.dtype == torch.float32 and weight.is_contiguous() + high = weight.bfloat16() + residual = weight - high.float() + middle = residual.bfloat16() + low = (residual - middle.float()).bfloat16() + return high, middle, low + + +@triton.jit +def _hc_mix_stats_bf16x3(X, W_HI, W_MID, W_LO, MIX, SQ, M, BLOCK_M: tl.constexpr): + # M stays runtime-valued so variable prefill lengths reuse the same binary. + rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, 32) + # 20480 input features / 16 independent slices. + start = tl.program_id(1) * 1280 + ks = start + tl.arange(0, 64) + hi = tl.zeros((BLOCK_M, 32), tl.float32) + mid = tl.zeros((BLOCK_M, 32), tl.float32) + lo = tl.zeros((BLOCK_M, 32), tl.float32) + sq = tl.zeros((BLOCK_M,), tl.float32) + for block in range(20): + k = ks + block * 64 + x = tl.load( + X + rows[:, None].to(tl.int64) * 20480 + k[None, :], + rows[:, None] < M, + 0, + ) + offsets = cols[None, :] * 20480 + k[:, None] + w_hi = tl.load(W_HI + offsets, cols[None, :] < 24, 0) + w_mid = tl.load(W_MID + offsets, cols[None, :] < 24, 0) + w_lo = tl.load(W_LO + offsets, cols[None, :] < 24, 0) + hi = tl.dot(x, w_hi, hi) + mid = tl.dot(x, w_mid, mid) + lo = tl.dot(x, w_lo, lo) + xf = x.to(tl.float32) + sq += tl.sum(xf * xf, 1) + offsets = (tl.program_id(1) * M + rows[:, None]) * 24 + cols[None, :] + tl.store(MIX + offsets, (hi + mid) + lo, (rows[:, None] < M) & (cols[None, :] < 24)) + tl.store(SQ + tl.program_id(1) * M + rows, sq, rows < M) + + +def hc_mix_stats_sinkhorn_bf16x3( + x: torch.Tensor, + weight_parts, + scale: torch.Tensor, + base: torch.Tensor, + sinkhorn_iters: int, + rms_eps: float, + hc_eps: float, +): + from sglang.kernels.ops.layernorm.mhc import _hc_mix_reduce_sinkhorn_kernel + + m = x.shape[0] + assert x.shape == (m, 20480) and x.is_contiguous() + assert x.dtype == torch.bfloat16 and 4096 <= m <= 65536 + assert len(weight_parts) == 3 + assert all( + w.shape == (24, 20480) and w.dtype == torch.bfloat16 and w.is_contiguous() + for w in weight_parts + ) + mix = torch.empty((16, m, 24), device=x.device, dtype=torch.float32) + sq = torch.empty((16, m), device=x.device, dtype=torch.float32) + pre = torch.empty((m, 4), device=x.device, dtype=torch.float32) + post = torch.empty_like(pre) + comb = torch.empty((m, 4, 4), device=x.device, dtype=torch.float32) + _hc_mix_stats_bf16x3[(triton.cdiv(m, 128), 16)]( + x, *weight_parts, mix, sq, m, 128, num_warps=4, num_stages=3 + ) + _hc_mix_reduce_sinkhorn_kernel[(m,)]( + mix, + sq, + scale, + base, + pre, + post, + comb, + m, + 1.0 / 20480, + rms_eps, + MIX=24, + HC=4, + NUM_SLICES=16, + ITERS=sinkhorn_iters, + EPS=hc_eps, + num_warps=1, + ) + return pre, post, comb diff --git a/python/sglang/kernels/ops/layernorm/hc_mix_stats_deepgemm.py b/python/sglang/kernels/ops/layernorm/hc_mix_stats_deepgemm.py new file mode 100644 index 000000000000..6cb245e65469 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/hc_mix_stats_deepgemm.py @@ -0,0 +1,67 @@ +"""Compensated FP32 mHC projections for SM100 batches with at least 128 rows. + +The small-row and batch-invariant paths remain in mhc.py. Native TF32 discards +too much of the FP32 projection weights, so evaluate their high and residual +components separately and bound accumulation length with a fixed split count. +""" + +import torch + +_NUM_SPLITS = 16 + + +def split_tf32_hc_weight(weight: torch.Tensor): + assert weight.dtype == torch.float32 and weight.is_contiguous() + high = (weight.view(torch.int32) & -8192).view(torch.float32) + return high, weight - high + + +def hc_mix_stats_sinkhorn_deepgemm( + x_flat: torch.Tensor, + weight_parts, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + sinkhorn_iters: int, + rms_eps: float, + hc_eps: float, +): + from sglang.kernels.ops.layernorm.mhc import _hc_mix_reduce_sinkhorn_kernel + from sglang.srt.layers.deep_gemm_wrapper.entrypoint import tf32_hc_prenorm_gemm + + assert x_flat.dtype == torch.bfloat16 and x_flat.is_contiguous() + m, k = x_flat.shape + high, low = weight_parts + assert k == 20480 and high.shape == low.shape == (24, k) + dev = x_flat.device + pre = torch.empty((m, 4), dtype=torch.float32, device=dev) + post = torch.empty_like(pre) + comb = torch.empty((m, 4, 4), dtype=torch.float32, device=dev) + if m == 0: + return pre, post, comb + + mix_hi = torch.empty((_NUM_SPLITS, m, 24), dtype=torch.float32, device=dev) + mix_lo = torch.empty_like(mix_hi) + sq = torch.empty((_NUM_SPLITS, m), dtype=torch.float32, device=dev) + unused_sq = torch.empty_like(sq) + tf32_hc_prenorm_gemm(x_flat, high, mix_hi, sq, _NUM_SPLITS) + tf32_hc_prenorm_gemm(x_flat, low, mix_lo, unused_sq, _NUM_SPLITS) + _hc_mix_reduce_sinkhorn_kernel[(m,)]( + mix_hi, + sq, + hc_scale, + hc_base, + pre, + post, + comb, + m, + 1.0 / k, + rms_eps, + MIX=24, + HC=4, + NUM_SLICES=_NUM_SPLITS, + ITERS=sinkhorn_iters, + EPS=hc_eps, + part_mix_residual_ptr=mix_lo, + num_warps=1, + ) + return pre, post, comb diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index bfe8f3106144..2c7a89dc1025 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -2246,6 +2246,7 @@ def _hc_mix_reduce_sinkhorn_kernel( NUM_SLICES: tl.constexpr, ITERS: tl.constexpr, EPS: tl.constexpr, + part_mix_residual_ptr=None, ): """One CTA per row keeps the sinkhorn reductions two-dimensional. Per-row arithmetic follows the slice reduction, then the Triton sinkhorn. @@ -2263,9 +2264,16 @@ def _hc_mix_reduce_sinkhorn_kernel( sq = tl.zeros([], dtype=tl.float32) for s in tl.static_range(NUM_SLICES): off = (s * m + row) * MIX - a_pre += tl.load(part_mix_ptr + off + j) - a_post += tl.load(part_mix_ptr + off + HC + j) - a_comb += tl.load(part_mix_ptr + off + 2 * HC + jj * HC + kk) + v_pre = tl.load(part_mix_ptr + off + j) + v_post = tl.load(part_mix_ptr + off + HC + j) + v_comb = tl.load(part_mix_ptr + off + 2 * HC + jj * HC + kk) + if part_mix_residual_ptr is not None: + v_pre += tl.load(part_mix_residual_ptr + off + j) + v_post += tl.load(part_mix_residual_ptr + off + HC + j) + v_comb += tl.load(part_mix_residual_ptr + off + 2 * HC + jj * HC + kk) + a_pre += v_pre + a_post += v_post + a_comb += v_comb sq += tl.load(part_sq_ptr + s * m + row) rsqrt = 1.0 / tl.sqrt(sq * inv_k + rms_eps) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 87fc13a2bf0d..e93b4f83cc85 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1452,6 +1452,9 @@ class Envs: # Kernels and indexer SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True) + # Compensated mHC caches FP32 weight splits; online weight updates must be + # disabled while this explicitly selected serving optimization is active. + SGLANG_DSV41_COMPENSATED_MHC = EnvBool(False) SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True) SGLANG_OPT_USE_FLASHINFER_MHC = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py index 83dba1a36647..f0a33a4815d3 100644 --- a/python/sglang/srt/layers/attention/graph_variants.py +++ b/python/sglang/srt/layers/attention/graph_variants.py @@ -64,7 +64,7 @@ def create_attention_graph_variants(hf_config) -> Optional[AttentionGraphVariant @dataclass(frozen=True) class Dsv41CandidateGraphVariants: - """DeepSeek-V4.1 candidate-indexer decode graphs. + """DeepSeek-V4.1 candidate-indexer decode and DSpark verify graphs. Selected by the longest request in the batch: while every request's positions fit a limit, the captured variant skips the low-ratio scoring or @@ -75,12 +75,22 @@ class Dsv41CandidateGraphVariants: # (label, max_seq_len it serves), ascending; the last label is the fallback. graph_limits: tuple[tuple[str, int], ...] capture_labels: tuple[str, ...] + verify_extra_tokens: int = 0 def select(self, forward_batch: ForwardBatch) -> str: - # Without the scheduler's CPU lengths, use the full graph without D2H. lengths = getattr(forward_batch, "seq_lens_cpu", None) + max_seq_len = None if lengths is not None and lengths.device.type == "cpu" and lengths.numel() > 0: max_seq_len = int(lengths.max()) + if max_seq_len is None and self.verify_extra_tokens: + # Includes acceptance still in flight, without a GPU-to-CPU copy. + max_seq_len = getattr( + getattr(forward_batch, "spec_info", None), + "candidate_max_seq_len_upper_bound", + None, + ) + if max_seq_len is not None: + max_seq_len += self.verify_extra_tokens for variant, limit in self.graph_limits: if max_seq_len <= limit: return variant @@ -88,17 +98,23 @@ def select(self, forward_batch: ForwardBatch) -> str: def create_dsv41_candidate_graph_variants( - model_runner, capture_forward_mode + model_runner, capture_forward_mode, captured_req_width: int = 0 ) -> Optional[Dsv41CandidateGraphVariants]: - """DeepSeek-V4.1 candidate graphs for plain decode on SM100+ CUDA, or None.""" + """DeepSeek-V4.1 candidate graphs for decode/DSpark verify on SM100+ CUDA.""" import torch from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.utils import is_hip text_config = model_runner.model_config.hf_text_config + dspark_target_verify = ( + capture_forward_mode == ForwardMode.TARGET_VERIFY + and model_runner.spec_algorithm.is_dspark() + and not model_runner.is_draft_worker + and captured_req_width > 0 + ) if not ( - capture_forward_mode == ForwardMode.DECODE + (capture_forward_mode == ForwardMode.DECODE or dspark_target_verify) and model_runner.device == "cuda" and not is_hip() and torch.cuda.get_device_capability(model_runner.gpu_id)[0] >= 10 @@ -112,7 +128,9 @@ def create_dsv41_candidate_graph_variants( ratios = set(text_config.compress_ratios) & {1, 2} topk = text_config.index_topk variants = [] - if topk > 0 and ratios: + # Verify still needs per-query causal top-k. Only remove candidate filtering + # when every possible block fits its budget; keep the low-ratio indexer. + if topk > 0 and ratios and not dspark_target_verify: variants.append(("candidate_all", topk * min(ratios))) if ratios == {1, 2}: variants.append(("candidate_c2_all", topk * 2)) @@ -130,4 +148,6 @@ def create_dsv41_candidate_graph_variants( return Dsv41CandidateGraphVariants( graph_limits=tuple(graph_limits), capture_labels=tuple(v for v, _ in graph_limits) + (DSV41_CANDIDATE_FILTERED,), + # The verify backend adds this width to committed CPU lengths. + verify_extra_tokens=captured_req_width if dspark_target_verify else 0, ) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 2d1d1086924c..acbc385ed828 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -506,6 +506,9 @@ def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]): self.w8a8_block_fp8_linear = None self.w8a8_mxfp8_linear = None self.mxfp8_dense_backend = None + # Set by a model-owned startup hook after opting into prefill tuning. + # Other block-FP8 models retain their fixed tactic at every batch size. + self.mxfp8_prefill_autotune_min_tokens = None if self.use_mxfp8 and not self.convert_mxfp8_to_block: self.mxfp8_dense_backend = resolve_mxfp8_dense_gemm_backend() self.w8a8_mxfp8_linear = dispatch_w8a8_mxfp8_linear() @@ -1099,6 +1102,19 @@ def apply( "128x4 MXFP8 input requires a FlashInfer CUTLASS backend" ) extra_kwargs = {} + if self.mxfp8_prefill_autotune_min_tokens is not None: + input_tensor = x[0] if isinstance(x, tuple) else x + num_tokens = input_tensor.numel() // input_tensor.shape[-1] + if num_tokens >= self.mxfp8_prefill_autotune_min_tokens: + from sglang.srt.batch_invariant_ops import ( + is_batch_invariant_mode_enabled, + ) + from sglang.srt.runtime_context import get_exec + + extra_kwargs["pin_tactic"] = ( + is_batch_invariant_mode_enabled() + or get_exec().deterministic.enable_deterministic_inference + ) if backend.is_flashinfer_cutlass() or backend.is_flashinfer_cutedsl(): weight_scale = layer.weight_scale_inv_swizzled elif backend.is_flashinfer_trtllm(): diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index e9dd1c7023f9..20ea96de2443 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py @@ -8,6 +8,7 @@ import torch from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.environ import envs from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader from sglang.srt.model_loader.utils import set_default_torch_dtype from sglang.srt.model_loader.weight_utils import default_weight_loader @@ -41,6 +42,14 @@ def _unsupported_derived_weight_cache_error() -> Optional[str]: old weights. The check is startup-determined and rank-uniform, so an update never proceeds on some workers while rejected on others. """ + if envs.SGLANG_DSV41_COMPENSATED_MHC.get(): + return ( + "Online weight updates are not supported with " + "SGLANG_DSV41_COMPENSATED_MHC=1: captured CUDA graphs retain the " + "derived mHC weight splits. Restart with this flag disabled to " + "use online weight updates." + ) + from sglang.kernels.ops.attention.dsv4.gemm import hpc_bf16xfp32_gemm_enabled if hpc_bf16xfp32_gemm_enabled(): @@ -307,7 +316,7 @@ def _update_bucketed_weights_from_distributed( ) reconstructed_tensors = bucket.reconstruct_tensors() self.get_model().load_weights(reconstructed_tensors) - return True, f"Succeeded to update parameter online." + return True, "Succeeded to update parameter online." except Exception as e: error_msg = ( f"Failed to update parameter online: {e}. " diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 5bd90e343e7d..f7dc6ac40cf9 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -301,7 +301,7 @@ def __init__( self.attention_graph_variants: Optional[AttentionGraphVariants] = ( create_attention_graph_variants(model_runner.model_config.hf_config) or create_dsv41_candidate_graph_variants( - model_runner, self.capture_forward_mode + model_runner, self.capture_forward_mode, self.captured_req_width ) ) diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 51fa74ecd41d..28e6b1543c0f 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -249,12 +249,13 @@ def _drop_diverged_autotune_cache( @contextlib.contextmanager def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool): # The gate below decides on the same inputs load_configs does. - from flashinfer.autotuner import _collect_metadata, autotune + from flashinfer.autotuner import AutoTuner, _collect_metadata, autotune mr = model_runner cache_path = flashinfer_autotune_cache_path(mr) sync_group = _autotune_tactic_sync_group(mr.tp_group) - if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get(): + reuse_cache = envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get() + if reuse_cache: autotune_cache = cache_path if sync_group is not None: _drop_diverged_autotune_cache(cache_path, sync_group, _collect_metadata()) @@ -277,16 +278,26 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool) from sglang.srt.layers.logits_processor import autotune_dummy_run_mode skip_ops = get_flashinfer_autotune_skip_ops(mr) + # autotune(cache=...) clears all file-loaded tactics on entry. In a + # speculative worker, loading the draft cache would then discard the + # target's prefill tactics after a restart (freshly profiled tactics + # live in a different cache and mask this on the first startup). + # The public load/save API merges the target and draft entries instead. + tuner = AutoTuner.get() + if reuse_cache and autotune_cache.is_file(): + tuner.load_configs(str(autotune_cache)) with ( _autotune_process_group(sync_group), autotune( True, - cache=str(autotune_cache), + cache=None if reuse_cache else str(autotune_cache), skip_ops=skip_ops, ), autotune_dummy_run_mode(run_lm_head=run_lm_head), ): yield + if reuse_cache: + tuner.save_configs(str(autotune_cache)) torch.cuda.current_stream().wait_stream(mr.forward_stream) logger.info("FlashInfer autotune completed.") @@ -334,7 +345,7 @@ def run_and_reset(): def maybe_flashinfer_autotune_extend( runner: BaseRunner, *, decode_num_tokens: int ) -> None: - """Also autotune one EXTEND-shaped dummy forward. + """Also autotune kernels at the prefill token ceiling. The decode-shaped autotune only covers token counts up to the decode batch size, so larger prefill/extend batches fall outside the tuned @@ -343,14 +354,25 @@ def maybe_flashinfer_autotune_extend( untuned at >=8k tokens on sm100). One extra forward at the largest per-rank extend token count tunes all buckets up to it. """ - if not envs.SGLANG_FLASHINFER_AUTOTUNE_EXTEND.get(): - return mr = runner.model_runner # Prefer the per-rank scheduler buffer while preserving the legacy ceiling # when chunked prefill is disabled. num_tokens = max_prefill_buffer_tokens() or get_schedule().max_prefill_tokens if num_tokens <= (decode_num_tokens or 0): return # decode-shaped autotune already covered these buckets + # A model can warm up its prefill kernels without constructing a dummy + # attention batch. In particular, DSpark's ordinary dummy forward uses + # TARGET_VERIFY and cannot cover large prefill GEMMs. Keep the existing + # cross-rank tactic synchronization, cache and skip policy for this hook. + prefill_autotune = getattr(mr.model, "autotune_prefill_kernels", None) + if prefill_autotune is not None and mr.is_generation and not mr.is_draft_worker: + with flashinfer_autotune_context(mr, run_lm_head=False): + tuned = prefill_autotune(num_tokens, dtype=mr.dtype) + if tuned: + return + + if not envs.SGLANG_FLASHINFER_AUTOTUNE_EXTEND.get(): + return is_pd_prefill_target = ( get_disagg().disaggregation_mode == "prefill" and not mr.is_draft_worker ) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 9e3aeef8a201..85f49818f19f 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -116,7 +116,7 @@ is_shared_experts_fusion_disabled, uses_per_rank_fused_shared_slots, ) -from sglang.srt.layers.quantization.fp8 import Fp8Config +from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod from sglang.srt.layers.quantization.fp8_utils import ( Mxfp8DenseGemmBackend, view_aiter_fused_rms_transposed_fp8_scale, @@ -451,13 +451,14 @@ def _apply_wo_a_bf16_matmul( is_decode: bool, is_target_verify: bool = False, fuse_mxfp8_quant: bool = False, + is_prefill: bool = False, ) -> torch.Tensor | Mxfp8SwizzledInput: """Compute bf16 wo_a: o [T, G, D] @ wo_a [G, R, D] -> [T, G, R]. Single-token decode uses a GEMV for the validated TP4 shape. Blackwell - verify batches up to 384 rows write token-major output directly to avoid - the layout copy before wo_b. ROCm decode can use aiter batched GEMM; - other cases use torch.einsum. + verify batches up to 384 rows and large prefill batches write token-major + output directly to avoid the layout copy before wo_b. ROCm decode can use + aiter batched GEMM; other cases use torch.einsum. """ global _wo_a_aiter_batched_gemm_disabled if ( @@ -473,6 +474,11 @@ def _apply_wo_a_bf16_matmul( and 0 < o.shape[0] <= 384 and get_platform().is_blackwell ) + or ( + is_prefill + and 4096 <= o.shape[0] <= 65536 + and get_platform().is_blackwell + ) ) and o.shape[1:] == (2, 4096) and wo_a.shape == (2, 1024, 4096) @@ -1215,7 +1221,15 @@ def _compute_q_b( if ( _is_cuda and q_out is not None - and 0 < q.shape[0] <= 8 + and ( + 0 < q.shape[0] <= 8 + or ( + self.is_dsv41 + and get_platform().is_blackwell + and self.n_local_heads == 16 + and 4096 <= q.shape[0] <= 65536 + ) + ) and self.head_dim == 512 and self.qk_rope_head_dim == 64 and q.dtype == q_out.dtype == torch.bfloat16 @@ -2145,6 +2159,7 @@ def forward( wo_a, is_decode=forward_batch.forward_mode.is_decode(), is_target_verify=forward_batch.forward_mode.is_target_verify(), + is_prefill=forward_batch.forward_mode.is_extend_without_speculative(), fuse_mxfp8_quant=( not get_forward().sp_active and getattr( @@ -2322,6 +2337,39 @@ def refresh_mhc_norm_weight_cache(self): self._post_attention_layernorm_weight_bf16 = ( self.post_attention_layernorm.weight.data.bfloat16().contiguous() ) + # Rebuilt after weight loading, like the norm cache above. Keep the + # original FP32 parameters intact for small rows and invariant mode. + self._hc_attn_tf32_parts = self._hc_ffn_tf32_parts = None + self._hc_attn_bf16_parts = self._hc_ffn_bf16_parts = None + if ( + self.hc_pre_from_prev_sublayer + and get_platform().is_sm100 + and self.hc_attn_fn.shape == (24, 20480) + and envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get() + and envs.SGLANG_DSV41_COMPENSATED_MHC.get() + ): + from sglang.kernels.ops.layernorm.hc_mix_stats_deepgemm import ( + split_tf32_hc_weight, + ) + from sglang.srt.layers.deep_gemm_wrapper.configurer import ( + ENABLE_JIT_DEEPGEMM, + ) + + if ENABLE_JIT_DEEPGEMM: + self._hc_attn_tf32_parts = split_tf32_hc_weight(self.hc_attn_fn.data) + self._hc_ffn_tf32_parts = split_tf32_hc_weight(self.hc_ffn_fn.data) + if ( + getattr(getattr(self, "config", None), "model_type", None) + == "deepseek_v41" + ): + from sglang.kernels.ops.layernorm.hc_mix_stats_bf16x3 import ( + split_bf16_hc_weight, + ) + + self._hc_attn_bf16_parts = split_bf16_hc_weight( + self.hc_attn_fn.data + ) + self._hc_ffn_bf16_parts = split_bf16_hc_weight(self.hc_ffn_fn.data) def hc_pre( self, @@ -2783,7 +2831,13 @@ def combine_and_norm(): if ( x.is_cuda and get_platform().is_blackwell - and 0 < x.shape[0] <= 8 + and ( + 0 < x.shape[0] <= 8 + or ( + self.config.model_type == "deepseek_v41" + and 4096 <= x.shape[0] <= 65536 + ) + ) and self.hc_mult == 4 and x_flat.shape[1] == 20480 and x.dtype == norm.weight.dtype == torch.bfloat16 @@ -2826,16 +2880,63 @@ def combine_and_norm(): if stats_stream is not None else nullcontext() ): - pre, post, comb = hc_mix_stats_sinkhorn( - x_flat, - hc_fn, - hc_scale, - hc_base, - self.hc_mult, - self.hc_sinkhorn_iters, - self.rms_norm_eps, - self.hc_eps, + from sglang.srt.batch_invariant_ops import ( + is_batch_invariant_mode_enabled, ) + + parts = bf16_parts = None + if ( + x_flat.shape[0] >= 128 + and x_flat.is_contiguous() + and get_platform().is_sm100 + and envs.SGLANG_OPT_DEEPGEMM_HC_PRENORM.get() + and not is_batch_invariant_mode_enabled() + ): + if hc_fn is self.hc_attn_fn: + parts = getattr(self, "_hc_attn_tf32_parts", None) + bf16_parts = getattr(self, "_hc_attn_bf16_parts", None) + elif hc_fn is self.hc_ffn_fn: + parts = getattr(self, "_hc_ffn_tf32_parts", None) + bf16_parts = getattr(self, "_hc_ffn_bf16_parts", None) + if bf16_parts is not None and 4096 <= x_flat.shape[0] <= 65536: + from sglang.kernels.ops.layernorm.hc_mix_stats_bf16x3 import ( + hc_mix_stats_sinkhorn_bf16x3, + ) + + pre, post, comb = hc_mix_stats_sinkhorn_bf16x3( + x_flat, + bf16_parts, + hc_scale, + hc_base, + self.hc_sinkhorn_iters, + self.rms_norm_eps, + self.hc_eps, + ) + elif parts is not None: + from sglang.kernels.ops.layernorm.hc_mix_stats_deepgemm import ( + hc_mix_stats_sinkhorn_deepgemm, + ) + + pre, post, comb = hc_mix_stats_sinkhorn_deepgemm( + x_flat, + parts, + hc_scale, + hc_base, + self.hc_sinkhorn_iters, + self.rms_norm_eps, + self.hc_eps, + ) + else: + pre, post, comb = hc_mix_stats_sinkhorn( + x_flat, + hc_fn, + hc_scale, + hc_base, + self.hc_mult, + self.hc_sinkhorn_iters, + self.rms_norm_eps, + self.hc_eps, + ) if stats_stream is not None: # These allocations originate on the side stream and are read # after the caller joins it, on the main stream. @@ -4061,6 +4162,67 @@ def __init__( # its barrier must only run on the first (startup) load. self._mhc_prewarmed_at_load = False + @torch.inference_mode() + def autotune_prefill_kernels(self, num_tokens: int, *, dtype: torch.dtype) -> int: + """Tune resident MXFP8 linears without touching request/KV/draft state. + + FlashInfer covers all M buckets through ``num_tokens`` from one call. + Decode/verify warmup only covers small M; the untuned large-M heuristic + can be substantially slower. Tune each distinct weight layout once and + call the quantization method directly to avoid TP collectives and model + side effects. The runner owns the synchronized autotune context. + """ + if getattr(self.config, "model_type", None) != "deepseek_v41": + return 0 + seen = set() + # The backbone excludes vision and lm_head, whose prefill shapes differ. + for layer in self.model.modules(): + method = getattr(layer, "quant_method", None) + if not isinstance(method, Fp8LinearMethod): + continue + if not (method.use_mxfp8 or method.block_fp8_as_mxfp8): + continue + if method.block_fp8_as_mxfp8 and not getattr( + layer, "block_fp8_mxfp8_ready", False + ): + # Some weights have a model-specific consumer or retain the + # block-FP8 fallback; they have no swizzled MXFP8 scale buffer. + continue + backend = method.mxfp8_dense_backend + if backend is None or not backend.is_flashinfer_cutedsl(): + continue + if method.block_fp8_as_mxfp8: + # Tune large row counts; small decode/verify shapes and + # deterministic execution retain their pinned tactic. + method.mxfp8_prefill_autotune_min_tokens = 4096 + weight = layer.weight + scale = layer.weight_scale_inv_swizzled + key = ( + weight.shape, + weight.stride(), + weight.dtype, + scale.shape, + scale.stride(), + scale.dtype, + ) + if key in seen: + continue + seen.add(key) + x = torch.zeros( + (num_tokens, weight.shape[1]), + dtype=dtype, + device=weight.device, + ) + method.apply(layer, x) + del x + if seen: + logger.info( + "FlashInfer prefill autotune: %d MXFP8 weight layouts at M=%d.", + len(seen), + num_tokens, + ) + return len(seen) + @property def routed_experts_weights_of_layer(self): return self._routed_experts_weights_of_layer.value diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 8390480af01c..2415a40d6a94 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -49,6 +49,9 @@ class DFlashVerifyInput(SpecInput): # Committed/live lengths before the verify caller temporarily expands # batch.seq_lens_cpu to the target-attention KV lengths. live_seq_lens_cpu: Optional[torch.Tensor] = None + # Conservative request-lifetime bound for candidate graph dispatch when + # DSpark keeps committed lengths on the GPU only. + candidate_max_seq_len_upper_bound: Optional[int] = None def __post_init__(self): super().__init__(spec_input_type=SpecInputType.DFLASH_VERIFY) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 83606b51113c..c09a22b1cacd 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -78,6 +78,34 @@ class TargetVerifyResult(msgspec.Struct, frozen=True): can_run_cuda_graph: bool +def candidate_request_length_bound( + reqs, pending_verify_tokens: int = 0 +) -> Optional[int]: + """Bound committed positions without reading asynchronous acceptance results. + + The overlap loop can have one unprocessed result, which may overshoot the + output budget. Reserve its full width here; the runner adds the current + verify width as well. Neither bound depends on CPU acceptance results. + Aborted/embedding/multimodal requests keep the general graph because their + visible token IDs may not represent the actual cache position space. + """ + if not reqs: + return None + longest = 0 + for req in reqs: + budget = req.sampling_params.max_new_tokens + if ( + not isinstance(budget, int) + or budget < 0 + or getattr(req, "to_finish", None) is not None + or getattr(req, "input_embeds", None) is not None + or getattr(req, "multimodal_inputs", None) is not None + ): + return None + longest = max(longest, len(req.origin_input_ids) + budget) + return longest + pending_verify_tokens + + class TargetVerifyExecutor: def __init__( self, @@ -296,6 +324,10 @@ def _forward_prepared_verify( seq_lens_cpu_backup, seq_lens_sum_backup, ) -> TargetVerifyResult: + if verify_input.live_seq_lens_cpu is None: + verify_input.candidate_max_seq_len_upper_bound = ( + candidate_request_length_bound(batch.reqs, self.verify_num_draft_tokens) + ) verify_forward_batch, _ = verify_input.prepare_for_verify( batch, self.target_worker ) diff --git a/test/registered/kernel/ops/attention/test_dsv41_prefill_wo_a.py b/test/registered/kernel/ops/attention/test_dsv41_prefill_wo_a.py new file mode 100644 index 000000000000..43fb29ee4399 --- /dev/null +++ b/test/registered/kernel/ops/attention/test_dsv41_prefill_wo_a.py @@ -0,0 +1,37 @@ +import unittest + +import torch + +from sglang.srt.models.deepseek_v4 import _apply_wo_a_bf16_matmul +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +@unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10, + "The prefill WO-A path targets Blackwell", +) +class TestPrefillWoA(CustomTestCase): + def test_exact_output_and_contiguous_layout(self): + torch.manual_seed(911) + weight = torch.randn(2, 1024, 4096, device="cuda", dtype=torch.bfloat16) + for rows in (4096, 4097, 65536): + with self.subTest(rows=rows): + # Match the attention backend's 64 padded heads, 16 local heads. + backing = torch.randn( + rows, 64, 512, device="cuda", dtype=torch.bfloat16 + ) + x = backing[:, :16].view(rows, 2, 4096) + expected = torch.einsum("tgd,grd->tgr", x, weight) + actual = _apply_wo_a_bf16_matmul( + x, weight, is_decode=False, is_prefill=True + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + self.assertTrue(actual.is_contiguous()) + self.assertEqual(actual.flatten(1).data_ptr(), actual.data_ptr()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernel/ops/layernorm/test_hc_combine_norm.py b/test/registered/kernel/ops/layernorm/test_hc_combine_norm.py new file mode 100644 index 000000000000..032ee1ef5176 --- /dev/null +++ b/test/registered/kernel/ops/layernorm/test_hc_combine_norm.py @@ -0,0 +1,60 @@ +import unittest + +import torch + +from sglang.kernels.ops.layernorm.hc_combine_norm import hc_combine_norm +from sglang.kernels.ops.layernorm.mhc import hc_combine +from sglang.srt.layers.layernorm import RMSNorm +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +@unittest.skipUnless(torch.cuda.is_available(), "Requires CUDA") +class TestHcCombineNorm(CustomTestCase): + def _check(self, x, pre, norm, actual): + expected = norm(hc_combine(x, pre, 4, torch.bfloat16)) + self.assertTrue(torch.isfinite(actual).all().item()) + # The same BF16 combine intermediate is preserved. RMS reductions may + # round differently, so check the aggregate error as well as each value. + torch.testing.assert_close(actual, expected, rtol=0.008, atol=0.001) + relative_l2 = ( + actual.float() - expected.float() + ).norm() / expected.float().norm().clamp_min(1e-20) + self.assertLess(relative_l2.item(), 5e-5) + + def test_prefill_and_small_batches(self): + torch.manual_seed(514) + norm = RMSNorm(5120, eps=1e-6).to(device="cuda", dtype=torch.bfloat16) + norm.weight.data.normal_(1, 0.1) + for rows in (1, 6, 4096, 4097, 65536): + with self.subTest(rows=rows): + # Exercise row strides and offsets beyond 2 GiB in the 64K case. + x = torch.randn(rows, 20488, device="cuda", dtype=torch.bfloat16)[ + :, :20480 + ] + pre = torch.rand(rows, 8, device="cuda")[:, :4] + actual = hc_combine_norm(x, pre, norm.weight, norm.variance_epsilon) + self._check(x, pre, norm, actual) + + def test_graph_replay_and_zero_input(self): + torch.manual_seed(516) + rows = 4097 + norm = RMSNorm(5120, eps=1e-6).to(device="cuda", dtype=torch.bfloat16) + x = torch.randn(rows, 20480, device="cuda", dtype=torch.bfloat16) + pre = torch.rand(rows, 4, device="cuda") + hc_combine_norm(x, pre, norm.weight, norm.variance_epsilon) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = hc_combine_norm(x, pre, norm.weight, norm.variance_epsilon) + for scale in (1.0, 0.01, 0.0): + with self.subTest(scale=scale): + x.normal_().mul_(scale) + pre.uniform_() + graph.replay() + self._check(x, pre, norm, actual) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernel/ops/layernorm/test_hc_mix_stats_bf16x3.py b/test/registered/kernel/ops/layernorm/test_hc_mix_stats_bf16x3.py new file mode 100644 index 000000000000..a2f851d14a76 --- /dev/null +++ b/test/registered/kernel/ops/layernorm/test_hc_mix_stats_bf16x3.py @@ -0,0 +1,156 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.kernels.ops.layernorm.hc_mix_stats_bf16x3 import ( + hc_mix_stats_sinkhorn_bf16x3, + split_bf16_hc_weight, +) +from sglang.kernels.ops.layernorm.hc_mix_stats_deepgemm import ( + hc_mix_stats_sinkhorn_deepgemm, + split_tf32_hc_weight, +) +from sglang.srt.environ import envs +from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=90, stage="base-b-kernel-unit", runner_config="4-gpu-b200") +EPS = 1e-6 + + +def reference(x, weight, scale, base): + x, weight, scale, base = [v.double() for v in (x, weight, scale, base)] + mixes = (x @ weight.T) * torch.rsqrt(x.square().mean(-1, keepdim=True) + EPS) + pre = torch.sigmoid(mixes[:, :4] * scale[0] + base[:4]) + EPS + post = 2 * torch.sigmoid(mixes[:, 4:8] * scale[1] + base[4:8]) + comb = torch.softmax((mixes[:, 8:] * scale[2] + base[8:]).view(-1, 4, 4), -1) + EPS + comb = comb / (comb.sum(-2, keepdim=True) + EPS) + for _ in range(19): + comb = comb / (comb.sum(-1, keepdim=True) + EPS) + comb = comb / (comb.sum(-2, keepdim=True) + EPS) + return pre, post, comb + + +@unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10, + "Compensated prefill projection targets Blackwell", +) +class TestHcMixStatsBf16x3(CustomTestCase): + def _inputs(self, rows, seed): + torch.manual_seed(seed) + x = torch.randn(rows, 20480, device="cuda", dtype=torch.bfloat16) + w = torch.randn(24, 20480, device="cuda") * 0.02 + scale = torch.tensor([0.1, 0.2, 0.3], device="cuda") + base = torch.randn(24, device="cuda") * 0.2 + return x, w, scale, base + + def test_prefill_coefficients(self): + for rows in (4096, 4097, 16384, 65536): + for seed in (0, 42): + with self.subTest(rows=rows, seed=seed): + x, w, scale, base = self._inputs(rows, seed) + parts = split_bf16_hc_weight(w) + torch.testing.assert_close( + sum(p.float() for p in parts), w, rtol=2e-7, atol=0 + ) + actual = hc_mix_stats_sinkhorn_bf16x3( + x, parts, scale, base, 20, EPS, EPS + ) + old = hc_mix_stats_sinkhorn_deepgemm( + x, split_tf32_hc_weight(w), scale, base, 20, EPS, EPS + ) + # All rows against the existing compensated implementation. + for a, b in zip(actual, old): + self.assertTrue(torch.isfinite(a).all().item()) + torch.testing.assert_close(a, b, rtol=2e-5, atol=2e-6) + # Independent FP64 reference, including a masked final tile. + indices = torch.cat( + ( + torch.arange(32, device="cuda"), + torch.arange(rows - 32, rows, device="cuda"), + ) + ) + expected = reference(x[indices], w, scale, base) + for a, b in zip(actual, expected): + torch.testing.assert_close( + a[indices].double(), b, rtol=2e-5, atol=2e-6 + ) + + def test_graph_replay(self): + x, w, scale, base = self._inputs(4097, 13) + parts = split_bf16_hc_weight(w) + hc_mix_stats_sinkhorn_bf16x3(x, parts, scale, base, 20, EPS, EPS) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = hc_mix_stats_sinkhorn_bf16x3(x, parts, scale, base, 20, EPS, EPS) + x.normal_() + graph.replay() + expected = reference(x[-32:], w, scale, base) + for a, b in zip(actual, expected): + torch.testing.assert_close(a[-32:].double(), b, rtol=2e-5, atol=2e-6) + + def test_model_dispatch_and_weight_refresh(self): + x, w, scale, base = self._inputs(4096, 17) + layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + torch.nn.Module.__init__(layer) + layer.config = SimpleNamespace(model_type="deepseek_v41") + layer.hc_attn_fn = torch.nn.Parameter(w) + layer.hc_ffn_fn = torch.nn.Parameter(w.clone()) + layer.hc_pre_from_prev_sublayer = True + layer.hc_mult, layer.hc_sinkhorn_iters = 4, 20 + layer.rms_norm_eps = layer.hc_eps = EPS + layer.input_layernorm = torch.nn.LayerNorm(5120, device="cuda") + layer.post_attention_layernorm = torch.nn.LayerNorm(5120, device="cuda") + with ( + envs.SGLANG_DSV41_COMPENSATED_MHC.override(True), + patch( + "sglang.srt.layers.deep_gemm_wrapper.configurer.ENABLE_JIT_DEEPGEMM", + True, + ), + ): + layer.refresh_mhc_norm_weight_cache() + previous = layer._hc_attn_bf16_parts + with torch.no_grad(): + layer.hc_attn_fn.add_(0.1) + layer.refresh_mhc_norm_weight_cache() + self.assertFalse(torch.equal(previous[0], layer._hc_attn_bf16_parts[0])) + torch.testing.assert_close( + sum(p.float() for p in layer._hc_attn_bf16_parts), + layer.hc_attn_fn, + rtol=2e-7, + atol=0, + ) + target = "sglang.kernels.ops.layernorm.hc_mix_stats_bf16x3.hc_mix_stats_sinkhorn_bf16x3" + for rows, invariant, expected in [ + (384, False, False), + (4096, True, False), + (4096, False, True), + ]: + with ( + self.subTest(rows=rows, invariant=invariant), + patch( + "sglang.srt.batch_invariant_ops.is_batch_invariant_mode_enabled", + return_value=invariant, + ), + patch(target, wraps=hc_mix_stats_sinkhorn_bf16x3) as fast, + ): + layer._hc_mix_and_combine( + x[:rows].view(rows, 4, 5120), + layer.hc_attn_fn, + scale, + base, + None, + lambda v: v, + ) + self.assertEqual(fast.called, expected) + with envs.SGLANG_DSV41_COMPENSATED_MHC.override(False): + layer.refresh_mhc_norm_weight_cache() + self.assertIsNone(layer._hc_attn_bf16_parts) + self.assertIsNone(layer._hc_ffn_bf16_parts) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernel/ops/layernorm/test_hc_mix_stats_deepgemm.py b/test/registered/kernel/ops/layernorm/test_hc_mix_stats_deepgemm.py new file mode 100644 index 000000000000..ca0b2222b96a --- /dev/null +++ b/test/registered/kernel/ops/layernorm/test_hc_mix_stats_deepgemm.py @@ -0,0 +1,203 @@ +import sys +from unittest.mock import patch + +import pytest +import torch + +from sglang.kernels.ops.layernorm.hc_mix_stats_deepgemm import ( + hc_mix_stats_sinkhorn_deepgemm, + split_tf32_hc_weight, +) +from sglang.srt.environ import envs +from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200") +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="Compensated mHC path targets datacenter Blackwell", +) +EPS = 1e-6 + + +def inputs(m, seed): + torch.manual_seed(seed) + x = torch.randn((m, 20480), device="cuda", dtype=torch.bfloat16) + w = torch.randn((24, 20480), device="cuda", dtype=torch.float32) * 0.02 + scale = torch.tensor([0.1, 0.2, 0.3], device="cuda") + base = torch.randn(24, device="cuda", dtype=torch.float32) * 0.2 + return x, w, scale, base + + +def reference(x, w, scale, base): + x, w, scale, base = (v.double() for v in (x, w, scale, base)) + mixes = (x @ w.T) * torch.rsqrt(x.square().mean(-1, keepdim=True) + EPS) + pre = torch.sigmoid(mixes[:, :4] * scale[0] + base[:4]) + EPS + post = 2 * torch.sigmoid(mixes[:, 4:8] * scale[1] + base[4:8]) + comb = (mixes[:, 8:] * scale[2] + base[8:]).view(-1, 4, 4) + comb = torch.softmax(comb, dim=-1) + EPS + comb = comb / (comb.sum(-2, keepdim=True) + EPS) + for _ in range(19): + comb = comb / (comb.sum(-1, keepdim=True) + EPS) + comb = comb / (comb.sum(-2, keepdim=True) + EPS) + return pre, post, comb + + +@pytest.mark.parametrize("m", [0, 128, 384, 2049, 4096, 16384, 32768, 65536]) +@pytest.mark.parametrize("seed", [0, 42]) +def test_compensated_coefficients_match_fp64(m, seed): + x, w, scale, base = inputs(m, seed) + parts = split_tf32_hc_weight(w) + assert torch.equal(parts[0] + parts[1], w) + got = hc_mix_stats_sinkhorn_deepgemm(x, parts, scale, base, 20, EPS, EPS) + expected = reference(x, w, scale, base) + for actual, ref in zip(got, expected): + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual.double(), ref, rtol=2e-5, atol=2e-6) + + +@pytest.mark.parametrize("m", [384, 4096]) +def test_graph_replay_reads_updated_input(m): + x, w, scale, base = inputs(m, 13) + parts = split_tf32_hc_weight(w) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + hc_mix_stats_sinkhorn_deepgemm(x, parts, scale, base, 20, EPS, EPS) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = hc_mix_stats_sinkhorn_deepgemm(x, parts, scale, base, 20, EPS, EPS) + # Scaling alone almost cancels under RMS normalization and would not expose + # a replay that accidentally kept using the capture-time activations. + x.normal_() + graph.replay() + expected = reference(x, w, scale, base) + for actual, ref in zip(captured, expected): + torch.testing.assert_close(actual.double(), ref, rtol=2e-5, atol=2e-6) + + +@pytest.mark.parametrize("m", [1, 6, 64, 127]) +def test_original_sinkhorn_path_without_residual_matches_fp64(m): + from sglang.kernels.ops.layernorm.mhc import hc_mix_stats_sinkhorn + + x, w, scale, base = inputs(m, 7) + got = hc_mix_stats_sinkhorn(x, w, scale, base, 4, 20, EPS, EPS) + expected = reference(x, w, scale, base) + for actual, ref in zip(got, expected): + torch.testing.assert_close(actual.double(), ref, rtol=2e-5, atol=2e-6) + + +@pytest.mark.parametrize( + "m,invariant,use_fast", + [ + (6, False, False), + (64, False, False), + (127, False, False), + (128, False, True), + (384, True, False), + (384, False, True), + (2049, True, False), + (2049, False, True), + ], +) +def test_model_dispatch_preserves_invariant_and_small_rows(m, invariant, use_fast): + x, w, scale, base = inputs(m, 1) + layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + torch.nn.Module.__init__(layer) + layer.hc_attn_fn = torch.nn.Parameter(w) + layer.hc_ffn_fn = torch.nn.Parameter(w.clone()) + layer._hc_attn_tf32_parts = split_tf32_hc_weight(w) + layer.hc_mult, layer.hc_sinkhorn_iters = 4, 20 + layer.rms_norm_eps = layer.hc_eps = EPS + target = "sglang.kernels.ops.layernorm.hc_mix_stats_deepgemm.hc_mix_stats_sinkhorn_deepgemm" + with ( + patch( + "sglang.srt.batch_invariant_ops.is_batch_invariant_mode_enabled", + return_value=invariant, + ), + patch(target, wraps=hc_mix_stats_sinkhorn_deepgemm) as fast, + ): + layer._hc_mix_and_combine( + x.view(m, 4, 5120), layer.hc_attn_fn, scale, base, None, lambda v: v + ) + assert fast.called == use_fast + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_weight_cache_honors_deepgemm_availability(enabled): + _, w, _, _ = inputs(0, 3) + layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + torch.nn.Module.__init__(layer) + layer.hc_pre_from_prev_sublayer = True + layer.hc_attn_fn = torch.nn.Parameter(w) + layer.hc_ffn_fn = torch.nn.Parameter(w.clone()) + layer.input_layernorm = torch.nn.LayerNorm(5120, device="cuda") + layer.post_attention_layernorm = torch.nn.LayerNorm(5120, device="cuda") + with ( + patch( + "sglang.srt.layers.deep_gemm_wrapper.configurer.ENABLE_JIT_DEEPGEMM", + enabled, + ), + envs.SGLANG_DSV41_COMPENSATED_MHC.override(True), + ): + layer.refresh_mhc_norm_weight_cache() + if enabled: + assert torch.equal(sum(layer._hc_attn_tf32_parts), w) + assert torch.equal(sum(layer._hc_ffn_tf32_parts), w) + else: + assert layer._hc_attn_tf32_parts is None + assert layer._hc_ffn_tf32_parts is None + + # The default opt-out also avoids allocating derived weight buffers. + with envs.SGLANG_DSV41_COMPENSATED_MHC.override(False): + layer.refresh_mhc_norm_weight_cache() + assert layer._hc_attn_tf32_parts is None + assert layer._hc_ffn_tf32_parts is None + + +@pytest.mark.parametrize("m", [128, 384, 4096]) +def test_fused_compensation_preserves_epilogue_bits(m): + from sglang.kernels.ops.layernorm.mhc import _hc_mix_reduce_sinkhorn_kernel + + torch.manual_seed(1) + hi = torch.randn(16, m, 24, device="cuda") + lo = torch.randn_like(hi) * 0.001 + sq = torch.rand(16, m, device="cuda") * 1280 + scale = torch.tensor([0.1, 0.2, 0.3], device="cuda") + base = torch.randn(24, device="cuda") + + def reduce(partial, residual=None): + pre = torch.empty(m, 4, device="cuda") + post = torch.empty_like(pre) + comb = torch.empty(m, 4, 4, device="cuda") + _hc_mix_reduce_sinkhorn_kernel[(m,)]( + partial, + sq, + scale, + base, + pre, + post, + comb, + m, + 1.0 / 20480, + EPS, + MIX=24, + HC=4, + NUM_SLICES=16, + ITERS=20, + EPS=EPS, + part_mix_residual_ptr=residual, + num_warps=1, + ) + return pre, post, comb + + expected = reduce(hi + lo) + actual = reduce(hi, lo) + for got, ref in zip(actual, expected): + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/kernel/quantization/test_flashinfer_autotune_cache_phases.py b/test/registered/kernel/quantization/test_flashinfer_autotune_cache_phases.py new file mode 100644 index 000000000000..2a7ab264df1b --- /dev/null +++ b/test/registered/kernel/quantization/test_flashinfer_autotune_cache_phases.py @@ -0,0 +1,116 @@ +"""A cached target warmup must survive a subsequent draft warmup.""" + +import json +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.model_executor.runner import flashinfer_autotune as warmup +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +@unittest.skipUnless(torch.cuda.is_available(), "FlashInfer requires CUDA") +class TestAutotuneCachePhases(CustomTestCase): + def test_disabling_cache_reuse_drops_file_loaded_tactics(self): + from flashinfer.autotuner import AutoTuner, _collect_metadata + + tuner = AutoTuner.get() + tuner.clear_cache() + self.addCleanup(tuner.clear_cache) + runner = SimpleNamespace( + device="cuda", + forward_stream=torch.cuda.Stream(), + tp_group=SimpleNamespace(world_size=1), + ) + with tempfile.TemporaryDirectory() as directory: + cache = Path(directory) / "cached.json" + cache.write_text( + json.dumps( + { + "_metadata": _collect_metadata(), + "old_tactic": ["TestRunner", 7], + } + ) + ) + tuner.load_configs(str(cache)) + self.assertIn("old_tactic", tuner._file_configs) + with ( + patch.object( + warmup, "flashinfer_autotune_cache_path", return_value=cache + ), + patch.object( + warmup, "get_flashinfer_autotune_skip_ops", return_value=set() + ), + patch.object( + warmup.envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE, + "get", + return_value=False, + ), + warmup.flashinfer_autotune_context(runner, run_lm_head=False), + ): + self.assertNotIn("old_tactic", tuner._file_configs) + + def test_loaded_target_tactics_survive_draft_cache(self): + from flashinfer.autotuner import AutoTuner, _collect_metadata + + tuner = AutoTuner.get() + tuner.clear_cache() + self.addCleanup(tuner.clear_cache) + runner = SimpleNamespace( + device="cuda", + forward_stream=torch.cuda.Stream(), + tp_group=SimpleNamespace(world_size=1), + ) + with tempfile.TemporaryDirectory() as directory: + target = Path(directory) / "target.json" + draft = Path(directory) / "draft.json" + metadata = _collect_metadata() + target.write_text( + json.dumps({"_metadata": metadata, "target_prefill": ["TestRunner", 7]}) + ) + draft.write_text( + json.dumps({"_metadata": metadata, "draft_decode": ["TestRunner", 3]}) + ) + with ( + patch.object( + warmup, + "flashinfer_autotune_cache_path", + side_effect=[target, draft], + ), + patch.object( + warmup, "get_flashinfer_autotune_skip_ops", return_value=set() + ), + patch.object( + warmup.envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE, + "get", + return_value=True, + ), + ): + with warmup.flashinfer_autotune_context(runner, run_lm_head=False): + self.assertEqual( + tuner._file_configs["target_prefill"], ("TestRunner", 7) + ) + # No operation was profiled: this simulates a restart where all + # target tactics came from disk, not the in-memory profile cache. + self.assertFalse(tuner.profiling_cache) + with warmup.flashinfer_autotune_context(runner, run_lm_head=False): + self.assertEqual( + tuner._file_configs["target_prefill"], ("TestRunner", 7) + ) + self.assertEqual( + tuner._file_configs["draft_decode"], ("TestRunner", 3) + ) + saved = json.loads(draft.read_text()) + self.assertEqual(saved["target_prefill"], ["TestRunner", 7]) + self.assertEqual(saved["draft_decode"], ["TestRunner", 3]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_draft_pad.py b/test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_draft_pad.py index 8765cb7c8df4..d00a8323de1d 100644 --- a/test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_draft_pad.py +++ b/test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_draft_pad.py @@ -21,7 +21,11 @@ import torch from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.kernels.deepseek_v4.common import make_paged_context, to_seq_extend +from sglang.test.kernels.deepseek_v4.common import ( + make_legacy_context, + make_paged_context, + to_seq_extend, +) from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") @@ -58,6 +62,55 @@ def _written_positions(plan_w: torch.Tensor, prefix_len: int) -> set[int]: class TestCompressWritePlanDraftPad(CustomTestCase): + def test_64k_prefill_preserves_last_token(self): + """65536 tokens fit uint16 indices; the last token must not wrap or vanish.""" + for cr in (4, 128): + paged = make_paged_context( + bs=16, compress_ratio=cr, num_swa_pages_per_req=16 + ) + legacy = make_legacy_context(bs=16, compress_ratio=cr) + seq_lens, extend_lens, num_q = to_seq_extend([(4096, 4096)] * 16) + for ctx, on_gpu in ((paged, False), (paged, True), (legacy, False)): + with self.subTest(cr=cr, paged=ctx is paged, on_gpu=on_gpu): + device = "cuda" if on_gpu else "cpu" + plan = ctx.make_prefill_plan( + seq_lens.to(device), extend_lens.to(device), num_q + ) + c = plan.plan_c.cpu().view(torch.int32).view(-1, 4) + valid_c = c[:, 0] != -1 + ids = c[valid_c, 1].bitwise_and(0xFFFF).sort().values + torch.testing.assert_close( + ids, torch.arange(cr - 1, num_q, cr, dtype=torch.int32) + ) + w = plan.plan_w.cpu().view(torch.int32).view(-1, 2) + last = w[w[:, 0] == 65535] + if cr == 4: + self.assertEqual(len(last), 1) + self.assertEqual(int(last[0, 1]), ctx.state_loc(15, 4095)) + else: + # Non-overlapping C128 consumed the complete final block; + # no raw tail remains to persist into the state ring. + self.assertEqual(len(w[w[:, 0] != -1]), 0) + + def test_prefill_rejects_uint16_index_overflow(self): + for ctx in ( + make_paged_context(bs=16, compress_ratio=4, num_swa_pages_per_req=17), + make_legacy_context(bs=16, compress_ratio=4), + ): + seq_lens, extend_lens, num_q = to_seq_extend( + [(4096, 4096)] * 15 + [(4097, 4097)] + ) + with self.assertRaisesRegex(RuntimeError, "plan_compress_prefill"): + ctx.make_prefill_plan(seq_lens, extend_lens, num_q) + + def test_prefill_rejects_packed_invalid_sentinel(self): + # A 65536-request, one-token-per-request batch makes the last packed + # (batch_id, ragged_id) equal (65535, 65535), the invalid write sentinel. + ctx = make_legacy_context(bs=65536, compress_ratio=4) + seq_lens, extend_lens, num_q = to_seq_extend([(1, 1)] * 65536) + with self.assertRaisesRegex(RuntimeError, "plan_compress_prefill"): + ctx.make_prefill_plan(seq_lens, extend_lens, num_q) + def _make_plan_positions( self, *, diff --git a/test/registered/kernels/test_dsv4_indexer_postprocess.py b/test/registered/kernels/test_dsv4_indexer_postprocess.py index 42a6558346f3..bccda0f7b177 100644 --- a/test/registered/kernels/test_dsv4_indexer_postprocess.py +++ b/test/registered/kernels/test_dsv4_indexer_postprocess.py @@ -24,6 +24,55 @@ def reference_pages(scores, indices, pages, page_size): class TestIndexerPostprocess(CustomTestCase): + def test_unfiltered_verify_matches_candidate_chain(self): + from sglang.kernels.ops.attention.dsv4.topk import ( + plan_topk_v2, + topk_transform_paged_v2, + ) + + # Six causal rows per request. The last request ends exactly at the + # candidate budget; capacity and unread logits extend well past it. + torch.manual_seed(123) + rows, width, page_size = 384, 32768, 64 + base = torch.tensor([0, 506, 4096, 16378], device="cuda", dtype=torch.int32) + lens = (base.repeat(16)[:, None] + torch.arange(1, 7, device="cuda")).flatten() + lens = lens.to(torch.int32) + source = torch.randn(rows, width, device="cuda").relu_() + consumer = torch.randn_like(source).relu_() + cols = torch.arange(width, device="cuda")[None, :] + source.masked_fill_(cols >= lens[:, None], 1e6) + consumer.masked_fill_(cols >= lens[:, None], 1e6) + pages = torch.randint( + 0, 100000, (rows, width // page_size), device="cuda", dtype=torch.int32 + ) + source_masked, keep = candidate_block_logits( + source, lens, topk_blocks=2048, block_size=8, published=None + ) + consumer_masked, _ = candidate_block_logits( + consumer, lens, topk_blocks=2048, block_size=8, published=keep + ) + plan = plan_topk_v2(lens) + for original, masked in ((source, source_masked), (consumer, consumer_masked)): + with self.subTest(source=original is source): + old = torch.empty((rows, 512), dtype=torch.int32, device="cuda") + new = torch.empty_like(old) + raw = torch.empty_like(old) + new_raw = torch.empty_like(old) + topk_transform_paged_v2(masked, lens, pages, old, page_size, plan, raw) + filter_topk_pages(masked, raw, pages, old, page_size) + topk_transform_paged_v2( + original, lens, pages, new, page_size, plan, new_raw + ) + # Top-k v2 uses a persistent work queue and does not promise + # output order. Compare the selected multiset, including -1 + # padding, in both logical-position and physical-page space. + torch.testing.assert_close( + new.sort(-1).values, old.sort(-1).values, rtol=0, atol=0 + ) + torch.testing.assert_close( + new_raw.sort(-1).values, raw.sort(-1).values, rtol=0, atol=0 + ) + def test_candidate_row_lens(self): lens = torch.tensor( [1, 7, 8, 9, 300, 16383, 16384, 16385, 16392, 40000, 1048576, 1048571], diff --git a/test/registered/kernels/test_dsv4_q_rope_store.py b/test/registered/kernels/test_dsv4_q_rope_store.py index c14733865757..c4943112a83b 100644 --- a/test/registered/kernels/test_dsv4_q_rope_store.py +++ b/test/registered/kernels/test_dsv4_q_rope_store.py @@ -7,7 +7,7 @@ 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-kernel-unit", runner_config="1-gpu-large") +register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") class TestQRopeStore(CustomTestCase): @@ -37,13 +37,40 @@ def test_exact_output_and_padding(self): torch.testing.assert_close(q, original, rtol=0, atol=0) self.assertTrue((padding[:, heads:] == 7).all().item()) - def test_graph_replay(self): - q = torch.randn(6, 16, 512, device="cuda", dtype=torch.bfloat16) - output = torch.zeros(6, 64, 512, device="cuda", dtype=q.dtype)[:, :16] + def test_large_prefill_exact_output_and_padding(self): + torch.manual_seed(911) + freqs = torch.polar( + torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda") + ) + for rows in (4096, 4097, 65536): + for dtype in (torch.int32, torch.int64): + with self.subTest(rows=rows, dtype=dtype): + q = torch.randn(rows, 17, 512, device="cuda", dtype=torch.bfloat16)[ + :, :16 + ] + original = q.clone() + expected = q.clone() + padding = torch.full( + (rows, 64, 512), 7.0, device="cuda", dtype=q.dtype + ) + positions = torch.randint( + 0, 8192, (rows,), device="cuda", dtype=dtype + ) + fused_rope_inplace(expected[..., 448:], None, freqs, positions) + q_rope_store(q, padding[:, :16], freqs, positions) + torch.testing.assert_close( + padding[:, :16], expected, rtol=0, atol=0 + ) + torch.testing.assert_close(q, original, rtol=0, atol=0) + self.assertTrue((padding[:, 16:] == 7).all().item()) + + def _check_graph_replay(self, rows): + q = torch.randn(rows, 16, 512, device="cuda", dtype=torch.bfloat16) + output = torch.zeros(rows, 64, 512, device="cuda", dtype=q.dtype)[:, :16] freqs = torch.polar( torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda") ) - positions = torch.arange(4096, 4102, device="cuda") + positions = torch.arange(rows, device="cuda") % 8192 q_rope_store(q, output, freqs, positions) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): @@ -56,6 +83,12 @@ def test_graph_replay(self): fused_rope_inplace(expected[..., 448:], None, freqs, positions) torch.testing.assert_close(output, expected, rtol=0, atol=0) + def test_graph_replay(self): + self._check_graph_replay(6) + + def test_large_prefill_graph_replay(self): + self._check_graph_replay(4097) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/quant/test_block_fp8_as_mxfp8.py b/test/registered/quant/test_block_fp8_as_mxfp8.py index 3fa721555639..aceab50318b5 100644 --- a/test/registered/quant/test_block_fp8_as_mxfp8.py +++ b/test/registered/quant/test_block_fp8_as_mxfp8.py @@ -1,7 +1,8 @@ """Check the MXFP8 linear layer against an FP32 dequantization reference.""" import unittest -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import Mock, patch import torch @@ -207,5 +208,125 @@ def test_against_dequantized_reference(self): self.assertLess(error, 1e-2, (n, k, m, error)) +class TestPrefillAutotune(_OptInCase): + def test_model_hook_deduplicates_ready_block_fp8_weights(self): + from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM + + layers = torch.nn.ModuleList() + methods = [] + q, s = _quant_block32(torch.randn(128, 128, device=DEVICE)) + for _ in range(2): + method = Fp8LinearMethod(_block32_config()) + layer = _build_layer(method, q, s) + layer.quant_method = method + method.apply = Mock() + methods.append(method) + layers.append(layer) + # An unprepared layer intentionally has no swizzled scale buffer. + fallback = torch.nn.Module() + fallback.quant_method = Fp8LinearMethod(_block32_config()) + fallback.block_fp8_mxfp8_ready = False + layers.append(fallback) + model = SimpleNamespace( + config=SimpleNamespace(model_type="deepseek_v41"), model=layers + ) + count = DeepseekV4ForCausalLM.autotune_prefill_kernels( + model, 4096, dtype=torch.bfloat16 + ) + self.assertEqual(count, 1) + methods[0].apply.assert_called_once() + self.assertEqual(methods[0].apply.call_args.args[1].shape, (4096, 128)) + methods[1].apply.assert_not_called() + for method in methods: + self.assertEqual(method.mxfp8_prefill_autotune_min_tokens, 4096) + self.assertIsNone(fallback.quant_method.mxfp8_prefill_autotune_min_tokens) + + def test_block_fp8_dispatch_keeps_decode_and_determinism_pinned(self): + method = Fp8LinearMethod(_block32_config()) + q, scale = _quant_block32(torch.randn(128, 128, device=DEVICE)) + layer = _build_layer(method, q, scale) + method.mxfp8_prefill_autotune_min_tokens = 4096 + call = Mock(return_value=torch.empty(0)) + method.w8a8_mxfp8_linear = call + for rows, invariant, deterministic, expected in ( + (6, False, False, None), + (384, False, False, None), + (4096, False, False, False), + (65536, False, False, False), + (4096, True, False, True), + (4096, False, True, True), + ): + with self.subTest( + rows=rows, invariant=invariant, deterministic=deterministic + ): + with ( + patch( + "sglang.srt.batch_invariant_ops.is_batch_invariant_mode_enabled", + return_value=invariant, + ), + patch( + "sglang.srt.runtime_context.get_exec", + return_value=SimpleNamespace( + deterministic=SimpleNamespace( + enable_deterministic_inference=deterministic + ) + ), + ), + ): + method.apply(layer, torch.empty(rows, 128, device=DEVICE)) + self.assertEqual(call.call_args.kwargs.get("pin_tactic"), expected) + + def test_tuned_prefill_against_fp32_reference(self): + self.enterContext( + patch( + "sglang.srt.runtime_context.get_exec", + return_value=SimpleNamespace( + deterministic=SimpleNamespace(enable_deterministic_inference=False) + ), + ) + ) + from flashinfer.autotuner import autotune + + from sglang.kernels.ops.quantization.fp8_kernel import ( + sglang_per_token_group_quant_fp8, + ) + + method = Fp8LinearMethod(_block32_config()) + n, k = 1792, 5120 + original_tf32 = torch.backends.cuda.matmul.allow_tf32 + self.addCleanup( + setattr, torch.backends.cuda.matmul, "allow_tf32", original_tf32 + ) + torch.backends.cuda.matmul.allow_tf32 = False + q, s = _quant_block32( + torch.randn(n, k, device=DEVICE, dtype=torch.bfloat16) / k**0.5 + ) + layer = _build_layer(method, q, s) + w_deq = _dequant_block32(q, s) + x = torch.randn(65536, k, device=DEVICE, dtype=torch.bfloat16) + method.mxfp8_prefill_autotune_min_tokens = 4096 + with autotune(True): + method.apply(layer, x) + for rows in (6, 384, 4096, 65536): + with self.subTest(rows=rows): + out = method.apply(layer, x[:rows]) + self.assertTrue(torch.isfinite(out).all().item()) + # Independently quantize/dequantize the first 64 rows. + xr = x[: min(rows, 64)] + xq, xs = sglang_per_token_group_quant_fp8(xr, BLOCK, scale_ue8m0=True) + x_deq = ( + xq.float().view(-1, k // BLOCK, BLOCK) + * xs.view(-1, k // BLOCK, 1).float() + ).view(-1, k) + ref = x_deq @ w_deq.t() + error = out[: ref.shape[0]].float() - ref + self.assertLess((error.norm() / ref.norm()).item(), 0.004) + if rows < 4096: + method.mxfp8_prefill_autotune_min_tokens = None + original = method.apply(layer, x[:rows]) + method.mxfp8_prefill_autotune_min_tokens = 4096 + torch.testing.assert_close(out, original, rtol=0, atol=0) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_executor/runner/test_model_prefill_autotune.py b/test/registered/unit/model_executor/runner/test_model_prefill_autotune.py new file mode 100644 index 000000000000..8e41a1882615 --- /dev/null +++ b/test/registered/unit/model_executor/runner/test_model_prefill_autotune.py @@ -0,0 +1,69 @@ +"""Model kernel warmup must cover prefill without a speculative dummy batch.""" + +import unittest +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.model_executor.runner import flashinfer_autotune as autotune +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, stage="base-a", runner_config="cpu") + + +class TestModelPrefillAutotune(CustomTestCase): + def setUp(self): + self.hook = Mock(return_value=1) + self.mr = SimpleNamespace( + model=SimpleNamespace(autotune_prefill_kernels=self.hook), + is_generation=True, + is_draft_worker=False, + dtype=torch.bfloat16, + ) + self.runner = SimpleNamespace(model_runner=self.mr) + # Deliberately no dummy-buffer or attention APIs: this path must not + # construct a TARGET_VERIFY batch or mutate request/KV state. + for target, kwargs in ( + ("max_prefill_buffer_tokens", {"return_value": 65536}), + ( + "flashinfer_autotune_context", + {"side_effect": lambda *a, **k: nullcontext()}, + ), + ): + p = patch.object(autotune, target, **kwargs) + setattr(self, target, p.start()) + self.addCleanup(p.stop) + p = patch.object( + autotune.envs.SGLANG_FLASHINFER_AUTOTUNE_EXTEND, "get", return_value=False + ) + p.start() + self.addCleanup(p.stop) + + def test_prefill_kernel_hook_uses_large_m_and_runner_dtype(self): + autotune.maybe_flashinfer_autotune_extend(self.runner, decode_num_tokens=384) + self.hook.assert_called_once_with(65536, dtype=torch.bfloat16) + self.flashinfer_autotune_context.assert_called_once_with( + self.mr, run_lm_head=False + ) + + def test_draft_worker_keeps_its_own_warmup(self): + self.mr.is_draft_worker = True + autotune.maybe_flashinfer_autotune_extend(self.runner, decode_num_tokens=384) + self.hook.assert_not_called() + self.flashinfer_autotune_context.assert_not_called() + + def test_no_extra_pass_when_decode_already_covers_prefill(self): + autotune.maybe_flashinfer_autotune_extend(self.runner, decode_num_tokens=65536) + self.hook.assert_not_called() + + def test_other_models_keep_extend_opt_in(self): + del self.mr.model.autotune_prefill_kernels + autotune.maybe_flashinfer_autotune_extend(self.runner, decode_num_tokens=384) + self.flashinfer_autotune_context.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_compensated_mhc_update_guard.py b/test/registered/unit/model_executor/test_compensated_mhc_update_guard.py new file mode 100644 index 000000000000..677d0bb31b6b --- /dev/null +++ b/test/registered/unit/model_executor/test_compensated_mhc_update_guard.py @@ -0,0 +1,43 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.environ import envs +from sglang.srt.model_executor.model_runner_components.weight_updater import ( + WeightUpdater, + _unsupported_derived_weight_cache_error, +) +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") + + +class TestCompensatedMhcUpdateGuard(CustomTestCase): + def test_direct_update_rejected_before_writes(self): + with ( + envs.SGLANG_DSV41_COMPENSATED_MHC.override(True), + patch( + "sglang.srt.model_executor.model_runner_components.weight_updater.default_weight_loader" + ) as loader, + ): + ok, message = WeightUpdater.update_weights_from_tensor( + SimpleNamespace(), [], load_format="direct" + ) + self.assertFalse(ok) + self.assertIn("SGLANG_DSV41_COMPENSATED_MHC", message) + loader.assert_not_called() + + def test_opt_out_keeps_existing_update_support(self): + with ( + envs.SGLANG_DSV41_COMPENSATED_MHC.override(False), + patch( + "sglang.kernels.ops.attention.dsv4.gemm.hpc_bf16xfp32_gemm_enabled", + return_value=False, + ), + ): + self.assertIsNone(_unsupported_derived_weight_cache_error()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_dsv41_verify_candidate_graph.py b/test/registered/unit/model_executor/test_dsv41_verify_candidate_graph.py new file mode 100644 index 000000000000..f857e026ad48 --- /dev/null +++ b/test/registered/unit/model_executor/test_dsv41_verify_candidate_graph.py @@ -0,0 +1,179 @@ +"""Verify graph selection must bound all speculative query positions.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from sglang.srt.layers.attention.graph_variants import ( + Dsv41CandidateGraphVariants, + create_dsv41_candidate_graph_variants, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.speculative.dspark_components.dspark_verify import ( + candidate_request_length_bound, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class TestVerifyCandidateGraph(CustomTestCase): + def make_policy(self, width=6): + return Dsv41CandidateGraphVariants( + graph_limits=(("candidate_unfiltered", 16384),), + capture_labels=("candidate_unfiltered", "candidate_filtered"), + verify_extra_tokens=width, + ) + + def test_longest_query_controls_selection(self): + policy = self.make_policy() + for lengths, expected in ( + ([4096] * 64, "candidate_unfiltered"), + ([4096, 16378], "candidate_unfiltered"), + ([4096, 16379], "candidate_filtered"), + ([16384], "candidate_filtered"), + ([1000000], "candidate_filtered"), + ): + with self.subTest(lengths=lengths): + batch = SimpleNamespace(seq_lens_cpu=torch.tensor(lengths)) + self.assertEqual(policy.select(batch), expected) + + def test_missing_or_non_cpu_lengths_use_full_graph(self): + policy = self.make_policy() + for lengths in ( + None, + torch.empty(0, dtype=torch.int64), + torch.empty(1, device="meta"), + ): + with self.subTest(lengths=lengths): + self.assertEqual( + policy.select(SimpleNamespace(seq_lens_cpu=lengths)), + "candidate_filtered", + ) + + def test_plain_decode_keeps_existing_boundary(self): + policy = self.make_policy(width=0) + batch = SimpleNamespace(seq_lens_cpu=torch.tensor([16384])) + self.assertEqual(policy.select(batch), "candidate_unfiltered") + + def test_gpu_only_lengths_use_request_budget(self): + policy = self.make_policy() + for bound, expected in ( + (5120, "candidate_unfiltered"), + (16378, "candidate_unfiltered"), + (16379, "candidate_filtered"), + (None, "candidate_filtered"), + ): + with self.subTest(bound=bound): + batch = SimpleNamespace( + seq_lens_cpu=None, + spec_info=SimpleNamespace(candidate_max_seq_len_upper_bound=bound), + ) + self.assertEqual(policy.select(batch), expected) + + def test_factory_keeps_causal_indexer_for_verify(self): + config = SimpleNamespace( + model_type="deepseek_v41", + candidate_source_layer_id=1, + candidate_topk_blocks=128, + candidate_block_size=128, + compress_ratios=[1, 2], + index_topk=2048, + ) + runner = SimpleNamespace( + model_config=SimpleNamespace(hf_text_config=config), + device="cuda", + gpu_id=0, + spec_algorithm=SimpleNamespace(is_dspark=lambda: True), + is_draft_worker=False, + ) + with ( + patch("torch.cuda.get_device_capability", return_value=(10, 0)), + patch("sglang.srt.utils.is_hip", return_value=False), + ): + verify = create_dsv41_candidate_graph_variants( + runner, ForwardMode.TARGET_VERIFY, 6 + ) + self.assertEqual( + verify.capture_labels, ("candidate_unfiltered", "candidate_filtered") + ) + self.assertEqual(verify.verify_extra_tokens, 6) + decode = create_dsv41_candidate_graph_variants( + runner, ForwardMode.DECODE, 1 + ) + self.assertEqual( + decode.capture_labels, + ( + "candidate_all", + "candidate_c2_all", + "candidate_unfiltered", + "candidate_filtered", + ), + ) + self.assertEqual(decode.verify_extra_tokens, 0) + # Never enable the shortcut for a draft worker or another algorithm. + runner.is_draft_worker = True + self.assertIsNone( + create_dsv41_candidate_graph_variants( + runner, ForwardMode.TARGET_VERIFY, 6 + ) + ) + runner.is_draft_worker = False + runner.spec_algorithm.is_dspark = lambda: False + self.assertIsNone( + create_dsv41_candidate_graph_variants( + runner, ForwardMode.TARGET_VERIFY, 6 + ) + ) + runner.spec_algorithm.is_dspark = lambda: True + self.assertIsNone( + create_dsv41_candidate_graph_variants( + runner, ForwardMode.TARGET_VERIFY, 0 + ) + ) + + def test_decode_ignores_verify_request_budget(self): + policy = self.make_policy(width=0) + batch = SimpleNamespace( + seq_lens_cpu=None, + spec_info=SimpleNamespace(candidate_max_seq_len_upper_bound=5120), + ) + self.assertEqual(policy.select(batch), "candidate_filtered") + + def test_host_lengths_take_precedence_over_request_budget(self): + batch = SimpleNamespace( + seq_lens_cpu=torch.tensor([16379]), + spec_info=SimpleNamespace(candidate_max_seq_len_upper_bound=5120), + ) + self.assertEqual(self.make_policy().select(batch), "candidate_filtered") + + def test_request_bound_uses_full_output_budget(self): + req = SimpleNamespace( + origin_input_ids=[0] * 4096, + sampling_params=SimpleNamespace(max_new_tokens=1024), + ) + self.assertEqual(candidate_request_length_bound([req] * 64), 5120) + self.assertEqual(candidate_request_length_bound([req] * 64, 6), 5126) + req.output_ids = [] # Accepted tokens can still be in flight. + self.assertEqual(candidate_request_length_bound([req]), 5120) + for attr, value in ( + ("to_finish", object()), + ("input_embeds", object()), + ("multimodal_inputs", object()), + ): + with self.subTest(attr=attr): + setattr(req, attr, value) + self.assertIsNone(candidate_request_length_bound([req])) + setattr(req, attr, None) + for budget in (None, -1, 1.5): + with self.subTest(budget=budget): + req.sampling_params.max_new_tokens = budget + self.assertIsNone(candidate_request_length_bound([req])) + self.assertIsNone(candidate_request_length_bound([])) + + +if __name__ == "__main__": + unittest.main()