diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 7f45aea5c..c386e6d30 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -2436,6 +2436,12 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec `None` Type: int + + `--cuda-graph-max-seq-len-prefill` + Longest sequence a prefill CUDA graph replay admits; longer batches run eager prefill. Folds into `cuda_graph_config[prefill].max_seq_len`. + `None` + Type: int + `--cuda-graph-bs-decode` Explicit list of batch sizes to capture for the decode CUDA graph. diff --git a/docs/src/snippets/configs/deepseek-ai/deepseek-v4_1.jsx b/docs/src/snippets/configs/deepseek-ai/deepseek-v4_1.jsx index cf6b0d512..b1adb3d41 100644 --- a/docs/src/snippets/configs/deepseek-ai/deepseek-v4_1.jsx +++ b/docs/src/snippets/configs/deepseek-ai/deepseek-v4_1.jsx @@ -369,7 +369,7 @@ export const config = { "--mem-fraction-static 0.8", "--speculative-algorithm DSPARK", "--speculative-dspark-block-size 5", - "--cuda-graph-max-bs 64", + "--cuda-graph-max-bs-decode 64", "--cuda-graph-backend-prefill breakable", "--cuda-graph-max-bs-prefill 4096", "--reasoning-parser auto", diff --git a/python/sglang/kernels/ops/attention/dsv4/__init__.py b/python/sglang/kernels/ops/attention/dsv4/__init__.py index 067d1cfd6..a99daa313 100644 --- a/python/sglang/kernels/ops/attention/dsv4/__init__.py +++ b/python/sglang/kernels/ops/attention/dsv4/__init__.py @@ -32,7 +32,12 @@ from .moe import ( silu_and_mul_contig_post_quant, silu_and_mul_masked_post_quant, ) -from .topk import plan_topk_v2, topk_transform_paged, topk_transform_paged_v2 +from .topk import ( + plan_topk_v2, + topk_transform_paged, + topk_transform_paged_v2, + topk_transform_ragged_v2, +) from .utils import make_name __all__ = [ @@ -56,6 +61,7 @@ __all__ = [ "triton_create_paged_compress_data", "topk_transform_paged", "topk_transform_paged_v2", + "topk_transform_ragged_v2", "plan_topk_v2", "hash_topk", "mega_moe_pre_dispatch", diff --git a/python/sglang/kernels/ops/attention/dsv4/c2_decode_pool.py b/python/sglang/kernels/ops/attention/dsv4/c2_decode_pool.py new file mode 100644 index 000000000..5f9f9bcff --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/c2_decode_pool.py @@ -0,0 +1,151 @@ +"""Fused ratio-2 decode pair-pooling, bitwise identical to the torch pool_pairs +path; a rounding difference here can change the indexer's top-k selection.""" + +from typing import Tuple + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + + +@triton.jit +def _c2_decode_pool_kernel( + kv_ptr, # [n, D] fp32 + score_ptr, # [n, D] fp32 + pos_ptr, # [n] int64 + raw_out_loc_ptr, # [n] int32/int64 + out_loc_ptr, # [n] int32/int64 + req_ptr, # [n] int64 + state_kv_ptr, # [R, D] fp32, in/out + state_score_ptr, # [R, D] fp32, in/out + pooled_ptr, # [n, D] fp32, out + group_pos_ptr, # [n] int64, out + slots_ptr, # [n] int64, out + pad_row, + RING_SIZE: tl.constexpr, + STATE_KV_STRIDE: tl.constexpr, + STATE_SCORE_STRIDE: tl.constexpr, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + row = tl.program_id(0) + + pos = tl.load(pos_ptr + row) + raw_loc = tl.load(raw_out_loc_ptr + row) + out_loc = tl.load(out_loc_ptr + row) + req = tl.load(req_ptr + row) + + # Raw location 0 is the padded-graph-row sentinel, and its req_pool_idx 0 may + # be a live request, so such a row's pair state goes to the spare row. + odd = (pos % 2) == 1 + if RING_SIZE: + r = tl.where( + (raw_loc == 0) | (pos == 0), + pad_row, + req * RING_SIZE + (pos - 1) % RING_SIZE, + ) + else: + r = tl.where(raw_loc == 0, pad_row, req) + + offs = tl.arange(0, BLOCK_D) + mask = offs < D + + kv = tl.load(kv_ptr + row * D + offs, mask=mask, other=0.0) + score = tl.load(score_ptr + row * D + offs, mask=mask, other=0.0) + p_kv = tl.load(state_kv_ptr + r * STATE_KV_STRIDE + offs, mask=mask, other=0.0) + p_score = tl.load( + state_score_ptr + r * STATE_SCORE_STRIDE + offs, mask=mask, other=0.0 + ) + + if RING_SIZE: + # Each live request has one decode row. Pad rows never modify the ring. + write_row = req * RING_SIZE + pos % RING_SIZE + tl.store( + state_kv_ptr + write_row * STATE_KV_STRIDE + offs, + kv, + mask=mask & (raw_loc != 0), + ) + tl.store( + state_score_ptr + write_row * STATE_SCORE_STRIDE + offs, + score, + mask=mask & (raw_loc != 0), + ) + else: + tl.store( + state_kv_ptr + r * STATE_KV_STRIDE + offs, + tl.where(odd, p_kv, kv), + mask=mask, + ) + tl.store( + state_score_ptr + r * STATE_SCORE_STRIDE + offs, + tl.where(odd, p_score, score), + mask=mask, + ) + + # libdevice.exp, not tl.exp: the approximate exponential changes the latent. + m = tl.maximum(p_score, score) + e0 = libdevice.exp(p_score - m) + e1 = libdevice.exp(score - m) + denom = e0 + e1 + # The + 0.0 below prevents FMA contraction: torch rounds both products first. + # libdevice.div_rn matches torch division; Triton's / is an approximate reciprocal. + t0 = p_kv * libdevice.div_rn(e0, denom) + t1 = kv * libdevice.div_rn(e1, denom) + t0 = t0 + 0.0 + t1 = t1 + 0.0 + pooled = t0 + t1 + tl.store(pooled_ptr + row * D + offs, pooled, mask=mask) + + # One program per row, so these are written exactly once each; no guard. + tl.store(group_pos_ptr + row, tl.where(odd, pos - 1, pos)) + tl.store(slots_ptr + row, tl.where(out_loc >= 0, out_loc, 0)) + + +def c2_decode_pool( + kv: torch.Tensor, + score: torch.Tensor, + pos: torch.Tensor, + raw_out_loc: torch.Tensor, + out_loc: torch.Tensor, + req: torch.Tensor, + state_kv: torch.Tensor, + state_score: torch.Tensor, + pad_row: int, + *, + ring_size: int = 0, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Ratio-2 decode pair pooling; updates `state_kv` / `state_score` in place. + + With ring_size > 0 the state halves may be views of an interleaved + CompressStatePool ring, one row per request otherwise; pad_row is the + padded-graph-row sentinel and is never written. + """ + assert kv.is_contiguous() and score.is_contiguous() + assert state_kv.stride(1) == state_score.stride(1) == 1 + assert kv.dtype == torch.float32 and score.dtype == torch.float32 + n, D = kv.shape + pooled = torch.empty_like(kv) + group_pos = torch.empty_like(pos) + slots = torch.empty(n, dtype=out_loc.dtype, device=out_loc.device) + _c2_decode_pool_kernel[(n,)]( + kv, + score, + pos, + raw_out_loc, + out_loc, + req, + state_kv, + state_score, + pooled, + group_pos, + slots, + pad_row, + RING_SIZE=ring_size, + STATE_KV_STRIDE=state_kv.stride(0), + STATE_SCORE_STRIDE=state_score.stride(0), + D=D, + BLOCK_D=triton.next_power_of_2(D), + num_warps=4, + ) + return pooled, group_pos, slots diff --git a/python/sglang/kernels/ops/attention/dsv4/decode_attention_sm100.py b/python/sglang/kernels/ops/attention/dsv4/decode_attention_sm100.py new file mode 100644 index 000000000..a64d1e23b --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/decode_attention_sm100.py @@ -0,0 +1,160 @@ +"""SM100 small-batch paged attention with the heads on the MMA N dimension; the +caller applies the inverse RoPE to the result.""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from .kv_layout import KVLayout + +LAYOUT = KVLayout.V4 +MAX_BATCH = 8 +NUM_HEADS = 16 +HEAD_DIM = 512 +SOFTMAX_SCALE = HEAD_DIM**-0.5 + + +def can_use_swapab_attention( + q: torch.Tensor, + kv: torch.Tensor, + extra_kv: Optional[torch.Tensor], + num_heads: int, + head_dim_v: int, + softmax_scale: float, +) -> bool: + """The caller adds the SM100 and single-query forward-mode gates.""" + return ( + 0 < q.shape[0] <= MAX_BATCH + and num_heads == NUM_HEADS + and q.dtype == torch.bfloat16 + and q.shape[-1] == HEAD_DIM + and head_dim_v == HEAD_DIM + and softmax_scale == SOFTMAX_SCALE + and kv.shape[-1] == LAYOUT.bytes_per_token + and (extra_kv is None or extra_kv.shape[-1] == LAYOUT.bytes_per_token) + ) + + +@triton.jit +def _combine( + PART, + MAX, + SUM, + SINK, + OUT, + NT: tl.constexpr, + ST: tl.constexpr, + H: tl.constexpr, + BD: tl.constexpr, +): + b, h, tile = tl.program_id(0), tl.program_id(1), tl.program_id(2) + t, d = tl.arange(0, ST), tile * BD + tl.arange(0, BD) + den = tl.load(SUM + (b * NT + t) * H + h, t < NT, 0) + mx = tl.load(MAX + (b * NT + t) * H + h, t < NT, 0) + mx = tl.where(den > 0, mx, -float("inf")) + sink = tl.load(SINK + h) + m = tl.maximum(tl.max(mx, 0), sink) + m = tl.where(tl.abs(m) == float("inf"), 0.0, m) + factor = tl.exp(mx - m) + denominator = tl.sum(den * factor, 0) + tl.exp(sink - m) + vals = tl.load( + PART + ((b * NT + t[:, None]) * H + h) * 512 + d[None, :], + t[:, None] < NT, + 0, + ) + out = tl.sum(vals * factor[:, None], 0) / denominator + out = tl.where((denominator > 0) & (sink != float("inf")), out, 0.0) + tl.store(OUT + (b * H + h) * 512 + d, out) + + +def swapab_attention( + q, + kv, + indices, + lengths, + sink, + extra_kv=None, + extra_indices=None, + extra_lengths=None, +): + """V4-layout attention on 16 heads; `extra_*` is a second slot range appended + to each request's keys, and the attention sink is folded in exactly once.""" + from .decode_attention_sm100_gluon import partial_gluon + + block = 64 + b, h, d = q.shape[0], q.shape[-2], q.shape[-1] + assert q.ndim in (3, 4) and (q.ndim == 3 or q.shape[1] == 1) + assert 0 < b <= MAX_BATCH and h == NUM_HEADS and d == HEAD_DIM + assert q.dtype == torch.bfloat16 and q.stride(-1) == 1 + assert kv.shape[-1] == LAYOUT.bytes_per_token + assert kv.dtype in (torch.uint8, torch.float8_e4m3fn) + assert indices.stride(-1) == 1 and lengths.is_contiguous() + assert sink.stride(0) == 1 and sink.numel() >= h + nk = indices.shape[-1] + ne = 0 if extra_indices is None else extra_indices.shape[-1] + assert block in (32, 64, 128) and nk > 0 + if extra_kv is None: + assert extra_indices is None and extra_lengths is None + extra_kv, extra_indices, extra_lengths = kv, indices, lengths + else: + assert extra_kv.shape[-1] == LAYOUT.bytes_per_token + assert extra_kv.dtype in (torch.uint8, torch.float8_e4m3fn) + assert extra_indices.stride(-1) == 1 and extra_lengths.is_contiguous() + kv, extra_kv = kv.view(torch.uint8), extra_kv.view(torch.uint8) + kt = triton.cdiv(nk, block) + nt = kt + triton.cdiv(ne, block) + partial = torch.empty((b, nt, h, 512), dtype=torch.float32, device=q.device) + maximum = torch.empty((b, nt, h), dtype=torch.float32, device=q.device) + sums = torch.empty_like(maximum) + out = torch.empty((b, h, 512), dtype=q.dtype, device=q.device) + partial_gluon[(b, nt)]( + q, + kv, + extra_kv, + indices, + extra_indices, + lengths, + extra_lengths, + partial, + maximum, + sums, + QS=q.stride(0), + QH=q.stride(-2), + IS=indices.stride(0), + EIS=extra_indices.stride(0), + KP=kv.shape[1], + KS=kv.stride(0), + EP=extra_kv.shape[1], + ES=extra_kv.stride(0), + NK=nk, + NE=ne, + NT=nt, + KT=kt, + BT=block, + H=h, + SCALE=SOFTMAX_SCALE, + KTOKENS=kv.shape[0] * kv.shape[1], + ETOKENS=extra_kv.shape[0] * extra_kv.shape[1], + COMPENSATE=True, + SWAP_AB=True, + DATA_BYTES=LAYOUT.data_bytes, + SCALE_BYTES=LAYOUT.scale_bytes, + TILE=LAYOUT.tile_size, + num_warps=4, + ) + bd = 64 if ne else 512 + _combine[(b, h, triton.cdiv(512, bd))]( + partial, + maximum, + sums, + sink, + out, + NT=nt, + ST=triton.next_power_of_2(nt), + H=h, + BD=bd, + num_warps=4, + ) + return out diff --git a/python/sglang/kernels/ops/attention/dsv4/decode_attention_sm100_gluon.py b/python/sglang/kernels/ops/attention/dsv4/decode_attention_sm100_gluon.py new file mode 100644 index 000000000..bb45e4834 --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/decode_attention_sm100_gluon.py @@ -0,0 +1,190 @@ +"""Native 16-head Blackwell (tcgen05) MMA layouts for paged V4-layout attention.""" + +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.blackwell import ( + TensorMemoryLayout, + allocate_tensor_memory, + fence_async_shared, + get_tmem_reg_layout, + mbarrier, + tcgen05_commit, + tcgen05_mma, +) + + +@gluon.jit +def _load_v4( + CACHE, + ids, + valid, + PAGE: gl.constexpr, + STRIDE: gl.constexpr, + KV_LAYOUT: gl.constexpr, + DATA_BYTES: gl.constexpr, + SCALE_BYTES: gl.constexpr, + TILE: gl.constexpr, +): + # V4 row: 448 fp8 nope + 64 bf16 rope = DATA_BYTES, plus one ue8m0 scale per + # TILE values in the page's scale rows. + d = gl.arange(0, 512, gl.SliceLayout(0, KV_LAYOUT)) + base = (ids // PAGE).to(gl.int64)[:, None] * STRIDE + slot = (ids % PAGE)[:, None] + mask = valid[:, None] & (d[None, :] < 448) + bits = gl.load(CACHE + base + slot * DATA_BYTES + d[None, :], mask, 0) + fp8 = bits.to(gl.float8e4nv, bitcast=True).to(gl.float32) + exponent = gl.load( + CACHE + base + PAGE * DATA_BYTES + slot * SCALE_BYTES + d[None, :] // TILE, + mask, + 0, + ).to(gl.int32) + scale = gl.where(exponent == 0, 0x00400000, exponent << 23).to( + gl.float32, bitcast=True + ) + rope_ptr = (CACHE + base + slot * DATA_BYTES + 448 + (d[None, :] - 448) * 2).to( + gl.pointer_type(gl.bfloat16) + ) + rope = gl.load(rope_ptr, valid[:, None] & (d[None, :] >= 448), 0) + return gl.where(d[None, :] < 448, fp8 * scale, rope.to(gl.float32)).to(gl.bfloat16) + + +@gluon.jit +def partial_gluon( + Q, + K, + E, + IDX, + EI, + L, + EL, + PART, + MAX, + SUM, + QS: gl.constexpr, + QH: gl.constexpr, + IS: gl.constexpr, + EIS: gl.constexpr, + KP: gl.constexpr, + KS: gl.constexpr, + EP: gl.constexpr, + ES: gl.constexpr, + NK: gl.constexpr, + NE: gl.constexpr, + NT: gl.constexpr, + KT: gl.constexpr, + BT: gl.constexpr, + H: gl.constexpr, + SCALE: gl.constexpr, + KTOKENS: gl.constexpr, + ETOKENS: gl.constexpr, + COMPENSATE: gl.constexpr, + SWAP_AB: gl.constexpr, + DATA_BYTES: gl.constexpr, + SCALE_BYTES: gl.constexpr, + TILE: gl.constexpr, +): + gl.static_assert(SWAP_AB and H == 16 and (BT == 64 or BT == 128)) + b, t = gl.program_id(0), gl.program_id(1) + kv_layout: gl.constexpr = gl.BlockedLayout([1, 8], [4, 8], [4, 1], [1, 0]) + n = gl.arange(0, BT, gl.SliceLayout(1, kv_layout)) + if t < KT: + at = t * BT + n + length = gl.load(L + b) + ids = gl.load(IDX + b * IS + at, at < NK, -1) + valid = (at < NK) & (at < length) & (ids >= 0) & (ids < KTOKENS) + kv = _load_v4( + K, + gl.maximum(ids, 0), + valid, + KP, + KS, + kv_layout, + DATA_BYTES, + SCALE_BYTES, + TILE, + ) + else: + at = (t - KT) * BT + n + length = gl.load(EL + b) + ids = gl.load(EI + b * EIS + at, at < NE, -1) + valid = (at < NE) & (at < length) & (ids >= 0) & (ids < ETOKENS) + kv = _load_v4( + E, + gl.maximum(ids, 0), + valid, + EP, + ES, + kv_layout, + DATA_BYTES, + SCALE_BYTES, + TILE, + ) + qh = gl.arange(0, H, gl.SliceLayout(1, kv_layout)) + qd = gl.arange(0, 512, gl.SliceLayout(0, kv_layout)) + q = gl.load(Q + b * QS + qh[:, None] * QH + qd[None, :]) + q_smem = gl.allocate_shared_memory( + gl.bfloat16, + [H, 512], + gl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16), + value=q, + ) + kv_smem = gl.allocate_shared_memory( + gl.bfloat16, + [BT, 512], + gl.NVMMASharedLayout(swizzle_byte_width=128, element_bitwidth=16), + value=kv, + ) + score_tmem = allocate_tensor_memory( + gl.float32, [BT, H], TensorMemoryLayout(block=(BT, H), col_stride=1) + ) + bar = gl.allocate_shared_memory(gl.int64, [1], mbarrier.MBarrierLayout()) + mbarrier.init(bar, count=1) + fence_async_shared() + tcgen05_mma(kv_smem, q_smem.permute((1, 0)), score_tmem, use_acc=False) + tcgen05_commit(bar) + mbarrier.wait(bar, phase=0) + score_layout: gl.constexpr = get_tmem_reg_layout( + gl.float32, (BT, H), TensorMemoryLayout(block=(BT, H), col_stride=1), 4 + ) + scores = score_tmem.load(score_layout) * SCALE + valid = gl.convert_layout(valid, gl.SliceLayout(1, score_layout)) + scores = gl.where(valid[:, None], scores, -float("inf")) + mx = gl.max(scores, 0) + mx = gl.where(mx == -float("inf"), 0.0, mx) + prob = gl.exp(scores - mx[None, :]) + denom = gl.sum(prob, 0) + p_hi = prob.to(gl.bfloat16) + p_smem = gl.allocate_shared_memory( + gl.bfloat16, + [BT, H], + gl.NVMMASharedLayout(swizzle_byte_width=32, element_bitwidth=16), + value=p_hi, + ) + out_tmem = allocate_tensor_memory( + gl.float32, [512, H], TensorMemoryLayout(block=(128, H), col_stride=1) + ) + fence_async_shared() + tcgen05_mma(kv_smem.permute((1, 0)), p_smem, out_tmem, use_acc=False) + tcgen05_commit(bar) + mbarrier.wait(bar, phase=1) + if COMPENSATE: + p_lo = (prob - p_hi.to(gl.float32)).to(gl.bfloat16) + p_smem.store(p_lo) + fence_async_shared() + tcgen05_mma(kv_smem.permute((1, 0)), p_smem, out_tmem, use_acc=True) + tcgen05_commit(bar) + mbarrier.wait(bar, phase=0) + mbarrier.invalidate(bar) + out_layout: gl.constexpr = get_tmem_reg_layout( + gl.float32, (512, H), TensorMemoryLayout(block=(128, H), col_stride=1), 4 + ) + value = out_tmem.load(out_layout) + hd_layout: gl.constexpr = gl.BlockedLayout([1, 4], [4, 8], [4, 1], [1, 0]) + value_hd = gl.convert_layout(value.permute((1, 0)), hd_layout) + h = gl.arange(0, H, gl.SliceLayout(1, hd_layout)) + d = gl.arange(0, 512, gl.SliceLayout(0, hd_layout)) + gl.store(PART + ((b * NT + t) * H + h[:, None]) * 512 + d[None, :], value_hd) + stat_layout: gl.constexpr = gl.BlockedLayout([1], [32], [4], [0]) + hs = gl.arange(0, H, stat_layout) + gl.store(MAX + (b * NT + t) * H + hs, gl.convert_layout(mx, stat_layout)) + gl.store(SUM + (b * NT + t) * H + hs, gl.convert_layout(denom, stat_layout)) diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py index 256fc4438..9cac53e38 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py +++ b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py @@ -7,8 +7,11 @@ from triton.language.extra import libdevice from sglang.kernels.ops.attention.dsv4.torch_quant import FP4_AMAX_FLOOR +INDEX_HEAD_DIM = 128 # One index-K slot: 64 packed e2m1 bytes and four ue8m0 block exponents. -INDEX_K_SLOT_BYTES = 64 + 4 +INDEX_K_PAYLOAD_BYTES = tl.constexpr(64) +INDEX_K_SCALE_BYTES = tl.constexpr(4) +INDEX_K_SLOT_BYTES = INDEX_K_PAYLOAD_BYTES.value + INDEX_K_SCALE_BYTES.value @triton.jit @@ -326,3 +329,138 @@ def index_k_rope_pack( num_warps=4, ) return (payload, scale) if cache is None else None + + +@triton.jit +def _e2m1_decode(code): + # code: uint 0..15 -> e2m1 value. exp = bits 2..1, mantissa = bit 0, sign = bit 3. + e = (code >> 1) & 3 + m = (code & 1).to(tl.float32) + sub = m * 0.5 + nor = (1.0 + m * 0.5) * tl.exp2((e - 1).to(tl.float32)) + v = tl.where(e == 0, sub, nor) + return tl.where((code >> 3) == 1, -v, v) + + +@triton.jit +def _fp4_index_logits_kernel( + q_ptr, # [B, H, D] bf16, fq4 queries (already rope'd) + w_ptr, # [B, H] bf16 head weights (softmax scale folded in) + slots_ptr, # [B, L] int64 pool slots per (request, compressed position) + lens_ptr, # [B] int64 visible compressed positions per request + table_ptr, # [num_pages, page_size * 64 + page_size * 4] uint8 + out_ptr, # [B, L] fp32 logits, -inf beyond lens + L, + page_size, + row_stride, + stride_qb, + stride_qh, + stride_wb, + H: tl.constexpr, + HALF_D: tl.constexpr, # D // 2 == 64 nibble-pairs per row + BLOCK_L: tl.constexpr, +): + b = tl.program_id(0) + lb = tl.program_id(1) + offs_l = lb * BLOCK_L + tl.arange(0, BLOCK_L) + offs_h = tl.arange(0, H) + offs_i = tl.arange( + 0, HALF_D + ) # byte index i holds elements 2i (low nibble), 2i+1 (high nibble) + + n_vis = tl.load(lens_ptr + b) + valid = offs_l < tl.minimum(n_vis, L) + slot = tl.load(slots_ptr + b * L + offs_l, mask=offs_l < L, other=0).to(tl.int64) + page = slot // page_size + off = slot % page_size + row_base = page * row_stride + + # K payload: [BLOCK_L, HALF_D] uint8 + pay = tl.load( + table_ptr + + row_base[:, None] + + off[:, None] * INDEX_K_PAYLOAD_BYTES + + offs_i[None, :], + mask=valid[:, None], + other=0, + ) + low = _e2m1_decode(pay & 0x0F) + high = _e2m1_decode((pay >> 4) & 0x0F) + # e8m0 block scales: element j uses block j // 32 -> byte i uses block i // 16. + sc_idx = offs_i // 16 + exps = tl.load( + table_ptr + + row_base[:, None] + + page_size * INDEX_K_PAYLOAD_BYTES + + off[:, None] * INDEX_K_SCALE_BYTES + + sc_idx[None, :], + mask=valid[:, None], + other=127, + ) + scale = tl.exp2(exps.to(tl.float32) - 127.0) + k_low = (low * scale).to(tl.bfloat16) # [BLOCK_L, HALF_D] elements 2i + k_high = (high * scale).to(tl.bfloat16) # elements 2i+1 + + # queries: even / odd elements, [H, HALF_D] bf16 + q_even = tl.load( + q_ptr + b * stride_qb + offs_h[:, None] * stride_qh + 2 * offs_i[None, :] + ) + q_odd = tl.load( + q_ptr + b * stride_qb + offs_h[:, None] * stride_qh + 2 * offs_i[None, :] + 1 + ) + + acc = tl.dot(q_even, tl.trans(k_low)) # [H, BLOCK_L] fp32 + acc += tl.dot(q_odd, tl.trans(k_high)) + # reference rounding points: bf16 dot -> relu -> * bf16 weight -> bf16 -> sum -> bf16 + s = acc.to(tl.bfloat16).to(tl.float32) + s = tl.maximum(s, 0.0) + w = tl.load(w_ptr + b * stride_wb + offs_h).to(tl.float32) + s = (s * w[:, None]).to(tl.bfloat16).to(tl.float32) + logit = tl.sum(s, axis=0).to(tl.bfloat16).to(tl.float32) + logit = tl.where(valid, logit, float("-inf")) + tl.store(out_ptr + b * L + offs_l, logit, mask=offs_l < L) + + +def fp4_index_logits_decode( + q: torch.Tensor, + weights: torch.Tensor, + slots: torch.Tensor, + lens: torch.Tensor, + table: torch.Tensor, + page_size: int, +) -> torch.Tensor: + """Decode index logits from the fp4 index-K pool. q [B, H, 128] bf16, weights + [B, H], slots [B, L] int64, lens [B] int64, table = the layer's index-K page + buffer (uint8, 2D). Returns [B, L] fp32 logits, -inf at positions >= lens, + rounded as the torch reference does.""" + assert q.dtype == torch.bfloat16 and q.shape[-1] == INDEX_HEAD_DIM + B, H, _ = q.shape + L = slots.shape[1] + assert table.dtype == torch.uint8 and table.dim() == 2 + q = q.contiguous() + weights = weights.to(torch.bfloat16).contiguous() + slots = slots.contiguous() + out = torch.empty((B, L), dtype=torch.float32, device=q.device) + if L == 0: + return out + BLOCK_L = 64 + grid = (B, triton.cdiv(L, BLOCK_L)) + _fp4_index_logits_kernel[grid]( + q, + weights, + slots, + lens.to(torch.int64).contiguous(), + table, + out, + L, + page_size, + table.stride(0), + q.stride(0), + q.stride(1), + weights.stride(0), + H=H, + HALF_D=INDEX_HEAD_DIM // 2, + BLOCK_L=BLOCK_L, + num_warps=4, + ) + return out diff --git a/python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py b/python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py index 44beb470d..b34b98eca 100644 --- a/python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py +++ b/python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py @@ -400,6 +400,7 @@ class BuildCausalSwaPageIndices: seq_lens_casual: torch.Tensor, swa_window: int, page_index_aligned_size: int, + swa_replay_start: Optional[torch.Tensor] = None, ) -> torch.Tensor: return build_causal_swa_page_indices( req_to_token=req_to_token, @@ -408,6 +409,7 @@ class BuildCausalSwaPageIndices: seq_lens_casual=seq_lens_casual, swa_window=swa_window, page_index_aligned_size=page_index_aligned_size, + swa_replay_start=swa_replay_start, ) @classmethod @@ -420,6 +422,7 @@ class BuildCausalSwaPageIndices: seq_lens_casual: torch.Tensor, swa_window: int, page_index_aligned_size: int, + swa_replay_start: Optional[torch.Tensor] = None, ) -> torch.Tensor: return build_causal_swa_page_indices_triton( req_to_token=req_to_token, @@ -428,9 +431,41 @@ class BuildCausalSwaPageIndices: seq_lens_casual=seq_lens_casual, swa_window=swa_window, page_index_aligned_size=page_index_aligned_size, + swa_replay_start=swa_replay_start, ) +def late_layer_tail_layout( + *, + extend_lens_cpu: list[int], + seq_lens_cpu: list[int], + tail_len: int, + device: torch.device, +) -> tuple[torch.Tensor, list[int], torch.Tensor]: + """Tail rows of each prefill extend: its last min(tail_len, extend_len) tokens. + Returns (token indices into the extend, per-request tail lengths, per-row + absolute window floor).""" + tail_lens_cpu = [min(tail_len, n) for n in extend_lens_cpu] + if len(extend_lens_cpu) == 1: + n, t, s = extend_lens_cpu[0], tail_lens_cpu[0], seq_lens_cpu[0] + floor = torch.full((t,), s - t, dtype=torch.int32, device=device) + return torch.arange(n - t, n, device=device), tail_lens_cpu, floor + # One H2D copy for the three length vectors; launch count does not grow with bs. + lens = torch.tensor([extend_lens_cpu, tail_lens_cpu, seq_lens_cpu], device=device) + extend_lens, tail_lens, seq_lens = lens[0], lens[1], lens[2] + total = sum(tail_lens_cpu) + req = torch.repeat_interleave( + torch.arange(len(tail_lens_cpu), device=device), tail_lens, output_size=total + ) + offs = ( + torch.arange(total, device=device) + - (torch.cumsum(tail_lens, 0) - tail_lens)[req] + ) + token_indices = (torch.cumsum(extend_lens, 0) - tail_lens)[req] + offs + floor = (seq_lens - tail_lens)[req].to(torch.int32) + return token_indices, tail_lens_cpu, floor + + def build_causal_swa_page_indices( *, req_to_token: torch.Tensor, @@ -439,14 +474,20 @@ def build_causal_swa_page_indices( seq_lens_casual: torch.Tensor, swa_window: int, page_index_aligned_size: int, + swa_replay_start: Optional[torch.Tensor] = None, ) -> torch.Tensor: + """Window slots each query attends to, -1 where empty. swa_replay_start floors + each row's window at that absolute position; None is the plain causal window.""" device = seq_lens_casual.device pos_causal = seq_lens_casual - 1 num_qo_tokens = seq_lens_casual.size(0) offsets = pos_causal.unsqueeze(1) - torch.arange( swa_window, dtype=torch.int32, device=device ).unsqueeze(0) - invalid_offset_mask = offsets < 0 + if swa_replay_start is None: + invalid_offset_mask = offsets < 0 + else: + invalid_offset_mask = offsets < swa_replay_start.to(offsets.dtype).unsqueeze(1) offsets.masked_fill_(invalid_offset_mask, 0) raw_indices = req_to_token[req_pool_indices_repeated[:, None], offsets] assert raw_indices.shape == (num_qo_tokens, swa_window) @@ -470,10 +511,12 @@ def _causal_swa_page_indices_kernel( full_to_swa_ptr, req_pool_ptr, seq_lens_ptr, + swa_replay_start_ptr, out_ptr, rt_stride, swa_window, padded_width, + HAS_SWA_REPLAY_START: tl.constexpr, BLOCK_K: tl.constexpr, ): row = tl.program_id(0) @@ -481,12 +524,16 @@ def _causal_swa_page_indices_kernel( rp = tl.load(req_pool_ptr + row).to(tl.int64) base = req_to_token_ptr + rp * rt_stride out_base = out_ptr + row.to(tl.int64) * padded_width + if HAS_SWA_REPLAY_START: + floor = tl.load(swa_replay_start_ptr + row).to(tl.int64) + else: + floor = tl.zeros((), dtype=tl.int64) for k0 in range(0, padded_width, BLOCK_K): k = k0 + tl.arange(0, BLOCK_K) kmask = k < padded_width off = pos - k.to(tl.int64) - valid = (k < swa_window) & (off >= 0) & kmask + valid = (k < swa_window) & (off >= floor) & kmask full_loc = tl.load(base + tl.where(valid, off, 0), mask=valid, other=-1).to( tl.int64 ) @@ -502,6 +549,7 @@ def build_causal_swa_page_indices_triton( seq_lens_casual: torch.Tensor, swa_window: int, page_index_aligned_size: int, + swa_replay_start: Optional[torch.Tensor] = None, ) -> torch.Tensor: num_qo_tokens = seq_lens_casual.size(0) padded_width = ( @@ -513,15 +561,19 @@ def build_causal_swa_page_indices_triton( device=seq_lens_casual.device, ) BLOCK_K = 256 + has_swa_replay_start = swa_replay_start is not None _causal_swa_page_indices_kernel[(num_qo_tokens,)]( req_to_token, full_to_swa_mapping, req_pool_indices_repeated, seq_lens_casual, + # Unused when HAS_SWA_REPLAY_START is False; any valid pointer will do. + swa_replay_start if has_swa_replay_start else seq_lens_casual, out, req_to_token.stride(0), swa_window, padded_width, + HAS_SWA_REPLAY_START=has_swa_replay_start, BLOCK_K=BLOCK_K, ) return out diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index faea161e2..b7c9fde82 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -1,11 +1,12 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple import torch import triton import triton.language as tl +from triton.language.extra import libdevice from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl, load_jit from sglang.kernels.kernel_api_logging import debug_kernel_api @@ -90,8 +91,12 @@ def moe_fused_gate_jit( def _router_triton_kernel( scores_ptr, # [M, N] raw logits, fp32/fp16/bf16 (upcast to fp32 on load) bias_ptr, # [N] fp32/fp16/bf16 (upcast to fp32 on load) + bias_alt_ptr, + input_ids_ptr, + num_token_non_padded_ptr, out_weights_ptr, # [M, K] fp32 out_indices_ptr, # [M, K] int32 + out_packed_ptr, # [M, K] int32 (HAS_PACKED) M, routed_scaling_factor, moe_softcapping, @@ -106,17 +111,28 @@ def _router_triton_kernel( EXPERTS_PER_GROUP: tl.constexpr, # N // N_GROUP BLOCK_G: tl.constexpr, # >= N_GROUP, power of 2 SCORING_FUNC: tl.constexpr, # 0 = sigmoid, 1 = sqrtsoftplus, 2 = softmax + SQRTSOFTPLUS_LOG1P: tl.constexpr, # sqrtsoftplus via log1p (V4.1 numerics) HAS_SOFTCAP: tl.constexpr, # tanh softcapping (softmax only) RENORMALIZE: tl.constexpr, APPLY_SCALE: tl.constexpr, # apply_routed_scaling_factor_on_output HAS_BIAS: tl.constexpr, + HAS_TOKEN_BIAS: tl.constexpr, + BIAS_ALT_TOKEN_ID: tl.constexpr, + HAS_PADDING: tl.constexpr, + HAS_PACKED: tl.constexpr, + RENORMALIZE_EPSILON: tl.constexpr, USE_PDL: tl.constexpr, + stride_bias, + stride_bias_alt, + stride_input_ids, stride_sm, stride_sn, stride_wm, stride_wk, stride_im, stride_ik, + stride_pm, + stride_pk, ) -> None: # Row-tiled: each program handles BLOCK_M rows; all reductions run along the # expert (N) axis. Tiling rows keeps CTAs large enough to stay occupancy-bound @@ -136,12 +152,30 @@ def _router_triton_kernel( # Plain softmax routing has no bias, so keep the zero value in registers # rather than materializing and clearing a device tensor per call. if HAS_BIAS: - bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32) + bias = tl.load(bias_ptr + offs_n * stride_bias, mask=mask_n, other=0.0).to( + tl.float32 + ) else: bias = tl.zeros([BLOCK_N], dtype=tl.float32) + if HAS_TOKEN_BIAS: + bias_alt = tl.load( + bias_alt_ptr + offs_n * stride_bias_alt, mask=mask_n, other=0.0 + ).to(tl.float32) + + live_m = mask_m + if HAS_PADDING: + live_m = live_m & (offs_m < tl.load(num_token_non_padded_ptr)) + row_bias = bias[None, :] + if HAS_TOKEN_BIAS: + input_ids = tl.load( + input_ids_ptr + offs_m * stride_input_ids, mask=live_m, other=0 + ) + row_bias = tl.where( + (input_ids == BIAS_ALT_TOKEN_ID)[:, None], bias_alt[None, :], row_bias + ) row_ptr = scores_ptr + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn - mask2d = mask_m[:, None] & mask_n[None, :] + mask2d = live_m[:, None] & mask_n[None, :] scores = tl.load(row_ptr, mask=mask2d, other=0.0).to( tl.float32 ) # [BLOCK_M, BLOCK_N] @@ -149,17 +183,21 @@ def _router_triton_kernel( if SCORING_FUNC == 0: # sigmoid(x) = 1 / (1 + exp(-x)); bias is for ranking only, weight is bias-free. activated = tl.sigmoid(scores) - biased = activated + bias[None, :] + biased = activated + row_bias elif SCORING_FUNC == 1: - # sqrt(softplus(x)). log(1.0 + exp(x)) rounds to 0 below -16.64 and overflows - # above 88.7; Triton has no log1p, so recover it from log via z*log(u)/(u-1). - z = tl.exp(-tl.abs(scores)) - u = 1.0 + z - exact = u == 1.0 - log1p_z = tl.where(exact, z, z * tl.log(u) / tl.where(exact, 1.0, u - 1.0)) - sp = tl.maximum(scores, 0.0) + log1p_z - activated = tl.sqrt(sp) - biased = activated + bias[None, :] + if SQRTSOFTPLUS_LOG1P: + # log1p preserves small positive scores for negative logits. + sp = tl.where(scores > 20.0, scores, libdevice.log1p(libdevice.exp(scores))) + activated = libdevice.sqrt(sp) + else: + # Open-coded log1p; reproduces the DeepSeek-V4 sqrtsoftplus numerics. + z = tl.exp(-tl.abs(scores)) + u = 1.0 + z + exact = u == 1.0 + log1p_z = tl.where(exact, z, z * tl.log(u) / tl.where(exact, 1.0, u - 1.0)) + sp = tl.maximum(scores, 0.0) + log1p_z + activated = tl.sqrt(sp) + biased = activated + row_bias else: # softmax over the row: weight is the softmax probability (bias kept), with # optional tanh softcapping. Ranking by the (softcapped, biased) logit is @@ -169,7 +207,7 @@ def _router_triton_kernel( # tanh(z) = 2*sigmoid(2z) - 1 (avoids relying on tl.math.tanh availability). z = logit / moe_softcapping logit = moe_softcapping * (2.0 * tl.sigmoid(2.0 * z) - 1.0) - biased = logit + bias[None, :] + biased = logit + row_bias biased = tl.where(mask_n[None, :], biased, -float("inf")) row_max = tl.max(biased, axis=1)[:, None] # [BLOCK_M, 1] exp_row = tl.where(mask_n[None, :], tl.exp(biased - row_max), 0.0) @@ -178,8 +216,11 @@ def _router_triton_kernel( biased = tl.where(mask_n[None, :], biased, -float("inf")) # [BLOCK_M, BLOCK_N] - # Map NaN -> a finite floor - biased = tl.where(biased == biased, biased, -1e30) # [BLOCK_M, BLOCK_N] + if SCORING_FUNC == 1 and SQRTSOFTPLUS_LOG1P: + # Rank NaNs above finite scores, matching torch.topk. + biased = tl.where(biased == biased, biased, float("inf")) + else: + biased = tl.where(biased == biased, biased, -1e30) # Grouped routing (DeepSeek-V3 noaux_tc): per-group score = sum of the top-2 # biased values; keep TOPK_GROUP groups (lowest group id wins ties); mask the @@ -214,9 +255,10 @@ def _router_triton_kernel( selected_idx = tl.zeros([BLOCK_M, BLOCK_K], dtype=tl.int32) cur = biased # [BLOCK_M, BLOCK_N] + remaining = tl.broadcast_to(mask_n[None, :], (BLOCK_M, BLOCK_N)) for k in tl.static_range(K_ROUTED): max_val = tl.max(cur, axis=1)[:, None] # [BLOCK_M, 1] - is_max = cur == max_val + is_max = remaining & (cur == max_val) lane_id = tl.where(is_max, offs_n[None, :], N + 1) # lowest expert id wins ties win_lane = tl.min(lane_id, axis=1)[:, None].to(tl.int32) # [BLOCK_M, 1] win_activated = tl.sum( @@ -225,7 +267,8 @@ def _router_triton_kernel( slot = offs_k[None, :] == k # [1, BLOCK_K] selected_vals = tl.where(slot, win_activated, selected_vals) selected_idx = tl.where(slot, win_lane, selected_idx) - cur = tl.where(offs_n[None, :] == win_lane, -float("inf"), cur) + remaining = remaining & (offs_n[None, :] != win_lane) + cur = tl.where(remaining, cur, -float("inf")) routed_sum = tl.sum(tl.where(mask_k_routed[None, :], selected_vals, 0.0), axis=1)[ :, None @@ -244,10 +287,16 @@ def _router_triton_kernel( tl.extra.cuda.gdc_launch_dependents() if RENORMALIZE: - norm = tl.where(routed_sum > 0.0, routed_sum, 1.0) # [BLOCK_M, 1] + if RENORMALIZE_EPSILON > 0.0: + norm = routed_sum + RENORMALIZE_EPSILON + else: + norm = tl.where(routed_sum > 0.0, routed_sum, 1.0) # [BLOCK_M, 1] selected_vals = selected_vals / norm if APPLY_SCALE: selected_vals = selected_vals * routed_scaling_factor + if HAS_PADDING: + selected_vals = tl.where(live_m[:, None], selected_vals, 0.0) + selected_idx = tl.where(live_m[:, None], selected_idx, -1) out_w_ptr = ( out_weights_ptr + offs_m[:, None] * stride_wm + offs_k[None, :] * stride_wk @@ -258,6 +307,25 @@ def _router_triton_kernel( store_mask = mask_m[:, None] & mask_k_total[None, :] tl.store(out_w_ptr, selected_vals, mask=store_mask) tl.store(out_i_ptr, selected_idx, mask=store_mask) + if HAS_PACKED: + # Must stay bitwise identical to fused_pack_topk. + w_bits = selected_vals.to(tl.bfloat16).to(tl.int16, bitcast=True).to(tl.int32) + packed = (selected_idx << 16) | (w_bits & 0xFFFF) + out_p_ptr = ( + out_packed_ptr + offs_m[:, None] * stride_pm + offs_k[None, :] * stride_pk + ) + tl.store(out_p_ptr, packed, mask=store_mask) + + +_DUMMY_I32: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_i32(device: torch.device) -> torch.Tensor: + # Placeholder pointer for kernel args whose constexpr flag is off. + t = _DUMMY_I32.get(device) + if t is None: + t = _DUMMY_I32[device] = torch.empty(1, dtype=torch.int32, device=device) + return t @debug_kernel_api @@ -273,14 +341,29 @@ def moe_fused_gate( moe_softcapping: float = 0.0, num_expert_group: int = 1, topk_group: int = 1, + *, + bias_alt: Optional[torch.Tensor] = None, + input_ids: Optional[torch.Tensor] = None, + bias_alt_token_id: Optional[int] = None, + num_token_non_padded: Optional[torch.Tensor] = None, + renormalize_epsilon: float = 0.0, + packed_out: Optional[torch.Tensor] = None, + sqrtsoftplus_log1p: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """Triton fused router: scoring + bias + topk + (optional) renorm/scale. - Mirrors the semantics of :func:`moe_fused_gate_jit` (the CUDA JIT kernel). + Mirrors :func:`moe_fused_gate_jit` (the CUDA JIT kernel) for the shared + parameters; the keyword-only extras are Triton-only. With ``num_expert_group > 1`` it performs DeepSeek-V3 grouped routing (per-group top-2-sum group scores, keep ``topk_group`` groups, then top-k - within). The first argument is named ``scores`` (raw GEMM logits) to match - the existing call sites. + within). ``scores`` contains raw GEMM logits. + + Rows past the device scalar ``num_token_non_padded`` return zero weights and -1 ids. + Positive ``renormalize_epsilon`` uses ``sum + epsilon`` instead of the zero-sum guard. + ``sqrtsoftplus_log1p`` evaluates sqrtsoftplus through ``log1p`` and ranks NaNs first + (DeepSeek-V4.1); off, the DeepSeek-V4 formula and NaN order are kept. + ``packed_out`` ([M, topk] int32, optional) receives the FlashInfer routed-MoE form + ``(id << 16) | bf16_bits(weight)``, bitwise identical to ``fused_pack_topk``. """ scoring_func_int = _SCORING_FUNC_MAP.get(scoring_func.lower()) assert scoring_func_int is not None, ( @@ -310,6 +393,15 @@ def moe_fused_gate( "scores and bias must have same num_experts" ) assert topk > num_fused_shared_experts, "topk must be > num_fused_shared_experts" + if input_ids is not None: + assert bias_alt is not None and bias_alt_token_id is not None + assert bias is not None and bias_alt.shape == bias.shape + assert input_ids.shape == (scores.size(0),) + if packed_out is not None: + assert packed_out.dtype == torch.int32, "packed_out must be int32" + assert packed_out.shape == (scores.size(0), topk), ( + "packed_out must be [M, topk]" + ) if routed_scaling_factor is None: routed_scaling_factor = 1.0 @@ -325,6 +417,11 @@ def moe_fused_gate( and num_fused_shared_experts == 0 and num_expert_group <= 1 and moe_softcapping == 0.0 + and input_ids is None + and num_token_non_padded is None + and renormalize_epsilon == 0.0 + and packed_out is None + and bias.stride(0) == 1 ): radix_args = ( scores, @@ -348,6 +445,8 @@ def moe_fused_gate( weights = torch.empty((M, K), dtype=torch.float32, device=scores.device) indices = torch.empty((M, K), dtype=torch.int32, device=scores.device) + if M == 0: + return weights, indices BLOCK_N = triton.next_power_of_2(N) # 256 -> 256, 384 -> 512 BLOCK_K = triton.next_power_of_2(K) # 6 -> 8, 8 -> 8 @@ -363,11 +462,18 @@ def moe_fused_gate( grid = (triton.cdiv(M, BLOCK_M),) use_pdl = is_arch_support_pdl() extra = {"launch_pdl": True} if use_pdl else {} + # Dynamo cannot analyze the kernel (PDL inline asm), so it writes back every + # pointer arg; aliasing an output as an unused arg's fallback clobbers it. + _unused_i32 = _dummy_i32(scores.device) _router_triton_kernel[grid]( scores, bias if bias is not None else scores, + bias_alt, + input_ids, + num_token_non_padded, weights, indices, + packed_out if packed_out is not None else _unused_i32, M, float(routed_scaling_factor), float(moe_softcapping), @@ -382,17 +488,28 @@ def moe_fused_gate( EXPERTS_PER_GROUP=experts_per_group, BLOCK_G=BLOCK_G, SCORING_FUNC=scoring_func_int, + SQRTSOFTPLUS_LOG1P=bool(sqrtsoftplus_log1p), HAS_SOFTCAP=bool(moe_softcapping != 0.0), RENORMALIZE=bool(renormalize), APPLY_SCALE=bool(apply_routed_scaling_factor_on_output), HAS_BIAS=bias is not None, + HAS_TOKEN_BIAS=input_ids is not None, + BIAS_ALT_TOKEN_ID=bias_alt_token_id, + HAS_PADDING=num_token_non_padded is not None, + HAS_PACKED=packed_out is not None, + RENORMALIZE_EPSILON=renormalize_epsilon, USE_PDL=use_pdl, + stride_bias=bias.stride(0) if bias is not None else 0, + stride_bias_alt=bias_alt.stride(0) if bias_alt is not None else 0, + stride_input_ids=input_ids.stride(0) if input_ids is not None else 0, stride_sm=scores.stride(0), stride_sn=scores.stride(1), stride_wm=weights.stride(0), stride_wk=weights.stride(1), stride_im=indices.stride(0), stride_ik=indices.stride(1), + stride_pm=packed_out.stride(0) if packed_out is not None else 0, + stride_pk=packed_out.stride(1) if packed_out is not None else 0, num_warps=num_warps, **extra, ) diff --git a/python/sglang/kernels/ops/speculative/dspark/commit_swa.py b/python/sglang/kernels/ops/speculative/dspark/commit_swa.py new file mode 100644 index 000000000..e94f61fa9 --- /dev/null +++ b/python/sglang/kernels/ops/speculative/dspark/commit_swa.py @@ -0,0 +1,43 @@ +"""Translate committed cache locations to the draft SWA pool.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _committed_swa_locations( + LOC, + MAP, + LENS, + OUT, + N: tl.constexpr, + WIDTH: tl.constexpr, + MAP_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + length = tl.load(LENS + i // WIDTH, i < N, other=0) + committed = (i < N) & (i % WIDTH < length) + loc = tl.load(LOC + i, committed, other=0).to(tl.int64) + # Preserve torch indexing for an unused negative location. + loc = tl.where(loc < 0, loc + MAP_SIZE, loc) + swa = tl.load(MAP + loc, committed, other=-1).to(tl.int32) + tl.store(OUT + i, swa, i < N) + + +def committed_swa_locations(cache_loc, full_to_swa_mapping, commit_lens, width): + assert cache_loc.numel() == commit_lens.numel() * width + out = torch.empty_like(cache_loc, dtype=torch.int32) + if cache_loc.numel(): + _committed_swa_locations[(triton.cdiv(cache_loc.numel(), 256),)]( + cache_loc, + full_to_swa_mapping, + commit_lens, + out, + cache_loc.numel(), + width, + full_to_swa_mapping.numel(), + 256, + ) + return out diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index b9f64cbed..785a9ddd3 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -584,12 +584,14 @@ class AcceptGreedy: target_logits: torch.Tensor, verify_num_draft_tokens: int, cutoff_verify_lens: Optional[torch.Tensor] = None, + fused_argmax: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return accept_greedy( candidates=candidates, target_logits=target_logits, verify_num_draft_tokens=verify_num_draft_tokens, cutoff_verify_lens=cutoff_verify_lens, + fused_argmax=fused_argmax, ) @classmethod @@ -600,12 +602,14 @@ class AcceptGreedy: target_logits: torch.Tensor, verify_num_draft_tokens: int, cutoff_verify_lens: Optional[torch.Tensor] = None, + fused_argmax: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return accept_greedy_triton( candidates=candidates, target_logits=target_logits, verify_num_draft_tokens=verify_num_draft_tokens, cutoff_verify_lens=cutoff_verify_lens, + fused_argmax=fused_argmax, ) @@ -615,9 +619,10 @@ def accept_greedy( target_logits: torch.Tensor, verify_num_draft_tokens: int, cutoff_verify_lens: Optional[torch.Tensor] = None, + fused_argmax: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bs = candidates.shape[0] - target_predict = torch.argmax(target_logits, dim=-1).view( + target_predict = _row_argmax(target_logits, fused=fused_argmax).view( bs, verify_num_draft_tokens ) correct_len, bonus = compute_dflash_correct_drafts_and_bonus( @@ -661,15 +666,35 @@ def gather_row_bonus_triton(*, table: torch.Tensor, idx: torch.Tensor) -> torch. return out +def _row_argmax(logits: torch.Tensor, fused: bool = False) -> torch.Tensor: + # torch.argmax uses one block per row; at few rows x wide vocab that is ~7x + # off the memory the reduction touches. The fused kernel does not reproduce + # torch.argmax's NaN selection, hence the opt-in. + if ( + fused + and logits.is_cuda + and logits.dim() == 2 + and logits.dtype == torch.float32 + and logits.stride(1) == 1 + and logits.shape[0] <= 64 + and logits.shape[1] >= 4096 + ): + from sglang.kernels.ops.speculative.row_argmax import row_argmax + + return row_argmax(logits) + return torch.argmax(logits, dim=-1) + + def accept_greedy_triton( *, candidates: torch.Tensor, target_logits: torch.Tensor, verify_num_draft_tokens: int, cutoff_verify_lens: Optional[torch.Tensor] = None, + fused_argmax: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bs = candidates.shape[0] - target_predict = torch.argmax(target_logits, dim=-1).view( + target_predict = _row_argmax(target_logits, fused=fused_argmax).view( bs, verify_num_draft_tokens ) correct_len, bonus = compute_dflash_correct_drafts_and_bonus( diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py index f8b26e8d0..1422b8ff6 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py @@ -384,9 +384,14 @@ class CommitKvProj: *, main_x: torch.Tensor, wkv_linears: list[torch.nn.Module], + allow_strided_output: bool = False, ) -> list[torch.Tensor]: if main_x.is_cuda and _fused_commit_kv_proj_supported(wkv_linears=wkv_linears): - return cls.triton(main_x=main_x, wkv_linears=wkv_linears) + return cls.triton( + main_x=main_x, + wkv_linears=wkv_linears, + allow_strided_output=allow_strided_output, + ) return cls.torch(main_x=main_x, wkv_linears=wkv_linears) @classmethod @@ -404,8 +409,13 @@ class CommitKvProj: *, main_x: torch.Tensor, wkv_linears: list[torch.nn.Module], + allow_strided_output: bool = False, ) -> list[torch.Tensor]: - return commit_kv_proj_fused(main_x=main_x, wkv_linears=wkv_linears) + return commit_kv_proj_fused( + main_x=main_x, + wkv_linears=wkv_linears, + allow_strided_output=allow_strided_output, + ) def commit_kv_proj( @@ -420,11 +430,20 @@ def commit_kv_proj_fused( *, main_x: torch.Tensor, wkv_linears: list[torch.nn.Module], + allow_strided_output: bool = False, ) -> list[torch.Tensor]: num_stages = len(wkv_linears) stacked = _stacked_wkv_weight(wkv_linears=wkv_linears) - if stacked.fp8_scale is not None: + if stacked.mxfp8_scale is not None: + kv_all = wkv_linears[0].quant_method.w8a8_mxfp8_linear( + input=main_x, + weight=stacked.weight, + weight_scale=stacked.mxfp8_scale, + input_scale=None, + bias=None, + ) + elif stacked.fp8_scale is not None: quant_method = wkv_linears[0].quant_method kv_all = quant_method.w8a8_block_fp8_linear( input=main_x, @@ -438,15 +457,14 @@ def commit_kv_proj_fused( kv_all = torch.nn.functional.linear(main_x, stacked.weight) head_dim = kv_all.shape[-1] // num_stages - return [ - kv_all[:, i * head_dim : (i + 1) * head_dim].contiguous() - for i in range(num_stages) - ] + slices = list(kv_all.split(head_dim, dim=-1)) + return slices if allow_strided_output else [kv.contiguous() for kv in slices] class _StackedWkvWeight(msgspec.Struct): weight: torch.Tensor fp8_scale: Optional[torch.Tensor] + mxfp8_scale: Optional[torch.Tensor] = None def _stacked_wkv_weight(*, wkv_linears: list[torch.nn.Module]) -> _StackedWkvWeight: @@ -500,6 +518,21 @@ def _build_stacked_wkv_weight( ) -> _StackedWkvWeight: if _block_quant_stack_applies(wkv_linears=wkv_linears): weight = torch.cat([linear.weight for linear in wkv_linears], dim=0) + backend = getattr(wkv_linears[0].quant_method, "mxfp8_dense_backend", None) + if ( + backend is not None + and (backend.is_flashinfer_cutlass() or backend.is_flashinfer_cutedsl()) + and all( + getattr(linear, "block_fp8_mxfp8_ready", False) + and linear.weight.shape[0] % 128 == 0 + for linear in wkv_linears + ) + ): + # 128-row-aligned scale tiles concatenate without breaking the swizzle. + scale = torch.cat( + [linear.weight_scale_inv_swizzled.reshape(-1) for linear in wkv_linears] + ) + return _StackedWkvWeight(weight=weight, fp8_scale=None, mxfp8_scale=scale) if wkv_linears[0].weight_scale_inv.dtype == torch.int32: from sglang.srt.layers.quantization.fp8_utils import ( inverse_transform_scale_ue8m0, diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py index 0fdd88ecb..1aeee2839 100644 --- a/python/sglang/srt/arg_groups/cuda_graph_hook.py +++ b/python/sglang/srt/arg_groups/cuda_graph_hook.py @@ -78,6 +78,8 @@ def parse_cuda_graph_config(server_args: Any): _set(Phase.DECODE, "max_bs", cfg.cuda_graph_max_bs_decode) if cfg.cuda_graph_max_bs_prefill is not None: _set(Phase.PREFILL, "max_bs", cfg.cuda_graph_max_bs_prefill) + if cfg.cuda_graph_max_seq_len_prefill is not None: + _set(Phase.PREFILL, "max_seq_len", cfg.cuda_graph_max_seq_len_prefill) if cfg.cuda_graph_bs_decode is not None: _set(Phase.DECODE, "bs", cfg.cuda_graph_bs_decode) if cfg.cuda_graph_bs_prefill is not None: diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 516e4a40b..22e098c26 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING from sglang.srt.arg_groups.overrides import ( _deepseek_v4_kv_cache_dtype, declare_resolution, + model_config_of, resolving_view, run_post_process_pass, ) @@ -244,14 +245,139 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: f"DeepSeekV4 CP supports moe_a2a_backend in {supported_a2a_backends}, " f"got {cfg.moe_a2a_backend!r}." ) - logger.warning( - "Disabling SGLANG_OPT_FLASHMLA_SPARSE_PREFILL because DeepSeekV4 " - "context parallelism is enabled." - ) - envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.set(False) + if model_config_of(server_args).hf_config.model_type != "deepseek_v41": + # The CP-aware sparse prefill chunk cache is validated on V4.1 only. + logger.warning( + "Disabling SGLANG_OPT_FLASHMLA_SPARSE_PREFILL because DeepSeekV4 " + "context parallelism is enabled." + ) + envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.set(False) logger.warning( f"Enable Context Parallel for DeepSeekV4, " f"strategy={cfg.cp_strategy}, " f"dp_size={cfg.dp_size}, moe_dense_tp_size={cfg.moe_dense_tp_size}, " f"attn_cp_size={cfg.attn_cp_size}, ep_size={cfg.ep_size}, tp_size={cfg.tp_size}" ) + + +def validate_deepseek_v41_features(server_args: ServerArgs) -> None: + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_triton, + ) + + cfg = resolving_view(server_args) + if model_config_of(server_args).hf_config.model_type != "deepseek_v41": + if cfg.enable_encoder_swa_bounded_replay: + raise ValueError( + "--enable-encoder-swa-bounded-replay requires DeepSeek-V4.1" + ) + return + if cfg.enable_encoder_swa_bounded_replay: + from sglang.srt.model_executor.cuda_graph_config import Backend + + incompatible = ( + ("non-CUDA hardware", not get_platform().is_cuda), + ( + "prefill CUDA graphs", + cfg.cuda_graph_config.prefill.backend != Backend.DISABLED, + ), + ("DP attention", cfg.enable_dp_attention), + ("context parallelism", cfg.attn_cp_size > 1), + ("external cache linker", cfg.enable_unified_cache_external_linker), + ("unified memory", cfg.enable_unified_memory), + ("PD disaggregation", cfg.disaggregation_mode != "null"), + ("mixed prefill/decode", cfg.enable_mixed_chunk), + ("LoRA", cfg.enable_lora), + ("radix sessions", cfg.enable_session_radix_cache), + ) + for feature, enabled in incompatible: + if enabled: + raise ValueError( + f"--enable-encoder-swa-bounded-replay does not support {feature} yet" + ) + if ( + cfg.max_running_requests is None + or cfg.max_running_requests <= 0 + or not cfg.chunked_prefill_size + or cfg.chunked_prefill_size < 128 + ): + raise ValueError( + "encoder SWA replay requires explicit --max-running-requests and --chunked-prefill-size >= 128" + ) + + unsupported = ( + ( + "speculative decoding other than DSpark", + cfg.speculative_algorithm is not None + and str(cfg.speculative_algorithm).upper() != "DSPARK", + ), + ("HiSparse", cfg.enable_hisparse), + ("the unified KV layout", is_unified_kv_triton()), + # The trtllm-gen path has no uniform-FP8 pool for V4.1's ratio-1/2 layers. + ("the trtllm DSv4 attention backend", cfg.dsv4_attn_backend == "trtllm"), + ("two-batch overlap", cfg.enable_two_batch_overlap), + ("pipeline parallelism", cfg.pp_size > 1), + ) + for feature, enabled in unsupported: + if enabled: + raise ValueError( + f"DeepSeek-V4.1 does not support {feature} yet; disable it to " + "serve this model." + ) + + if cfg.disaggregation_mode != "null" and cfg.speculative_algorithm is not None: + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) + + if ( + read_ragged_verify_mode() is not RaggedVerifyMode.STATIC + or cfg.disaggregation_transfer_backend != "mooncake" + or cfg.dp_size != 1 + or cfg.enable_dp_attention + or cfg.attn_cp_size != 1 + or cfg.dcp_size != 1 + ): + raise ValueError( + "DeepSeek-V4.1 DSpark PD requires static verify, Mooncake, " + "DP=1 and CP=1. Both servers must enable DSpark with the same " + "block size and TP size." + ) + + from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase + + prefill_graph = cfg.cuda_graph_config.prefill + if prefill_graph.backend != Backend.DISABLED and prefill_graph.max_seq_len is None: + # The captured low-ratio indexer scores a static context width; 16k + # keeps it inside the candidate window at under 1 ms per layer. + declare_resolution( + server_args, + "validate_deepseek_v41_features", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, max_seq_len=16 * 1024 + ), + ) + logger.warning( + "Setting cuda_graph_config[prefill].max_seq_len to 16384 for " + "DeepSeek-V4.1; longer contexts run eager prefill." + ) + + if cfg.enable_decoder_swa_bounded_replay: + from sglang.srt.model_executor.cuda_graph_config import Backend + + # Late layers see a per-request tail slice, not the captured prefill shape. + incompatible = ( + ( + "the prefill CUDA graph", + cfg.cuda_graph_config.prefill.backend != Backend.DISABLED, + ), + # input_ids_global is a DP-wide gather, so the tail slice cannot apply. + ("DP attention", cfg.enable_dp_attention), + ) + for feature, enabled in incompatible: + if enabled: + raise ValueError( + "--enable-decoder-swa-bounded-replay cannot be combined with " + f"{feature} yet; disable one of them." + ) diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index 05c80805c..314281680 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -96,6 +96,15 @@ class ExecFeatures(msgspec.Struct): bool, "Enable returning indexer topk indices of layers with indexer with responses.", ] = False + enable_encoder_swa_bounded_replay: A[ + bool, + "DeepSeek-V4.1 encoder SWA bounded replay: cache Main KV and Indexer keys only, " + "rebuild request-owned SWA windows on prefix hits. Experimental; CUDA only.", + ] = False + enable_decoder_swa_bounded_replay: A[ + bool, + "DeepSeek-V4.1 decoder SWA bounded replay: after the last kv_source layer, run the remaining layers over only the last window_size tokens of a prefill. Main and indexer KV stay exact; nothing is replayed. Deterministic for a fixed prompt and chunk size.", + ] = False sampling_mask_max_tokens: A[ int, "The maximum number of token IDs in a returned sampling mask. Requests " @@ -486,6 +495,11 @@ class ExecGraph(msgspec.Struct): cuda_graph_max_bs_prefill: A[ Optional[int], "Maximum batch size captured for the prefill cuda graph." ] = None + cuda_graph_max_seq_len_prefill: A[ + Optional[int], + "Longest sequence a prefill cuda graph replay admits; longer batches " + "run eager prefill. Folds into cuda_graph_config[prefill].max_seq_len.", + ] = None cuda_graph_bs_decode: A[ Optional[List[int]], "Explicit list of batch sizes to capture for the decode cuda graph.", diff --git a/python/sglang/srt/arg_groups/fields/schedule.py b/python/sglang/srt/arg_groups/fields/schedule.py index 2ed71b8a3..0d835e859 100644 --- a/python/sglang/srt/arg_groups/fields/schedule.py +++ b/python/sglang/srt/arg_groups/fields/schedule.py @@ -163,6 +163,23 @@ class Schedule(msgspec.Struct): fallback=0.8, ), ] = None + # Recorded by the cache hook; the effective field answers the fallback when unset. + _swa_full_tokens_ratio_explicitly_set: A[ + Optional[bool], + Arg(no_cli=True), + ] = None + swa_prefix_tails: A[ + Optional[int], + Arg( + help=( + "When the SWA KV pool is sized from the request cap (DeepSeek-V4 " + "family), how many radix-cached prefix tails it keeps room for. " + "Each tail is one sliding window plus one page. Default: 4 x " + "max_running_requests per attention-DP rank, 0 when the radix " + "cache is disabled." + ), + ), + ] = None disable_hybrid_swa_memory: A[ bool, Arg(help="Disable the hybrid SWA memory pool.", resolvable=True) ] = False diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 8c246858b..910ca2925 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -203,6 +203,13 @@ def handle_cache_compatibility(server_args: Any) -> None: "both build a decode host pool." ) + if cfg._swa_full_tokens_ratio_explicitly_set is None: + declare_resolution( + server_args, + "_handle_cache_compatibility", + _swa_full_tokens_ratio_explicitly_set=cfg.swa_full_tokens_ratio is not None, + ) + # Validate the effective ratio: model branches may declare a reset # (e.g. Step3p forces 1.0 under hierarchical cache) that supersedes # the user input before it ever takes effect. @@ -210,6 +217,9 @@ def handle_cache_compatibility(server_args: Any) -> None: # claimed the field, and the value to range-check is the effective one. if not (0 < resolution_result(server_args, "swa_full_tokens_ratio") <= 1.0): raise ValueError("--swa-full-tokens-ratio should be in range (0, 1.0].") + prefix_tails = resolved_view(server_args).swa_prefix_tails + if prefix_tails is not None and prefix_tails < 0: + raise ValueError("--swa-prefix-tails should be a non-negative integer.") def handle_unified_memory_pool(server_args: Any) -> None: diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index 86e3ff97c..e8a184a8f 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -416,8 +416,11 @@ def handle_model_specific_adjustments(server_args: Any): from sglang.srt.arg_groups.deepseek_v4_hook import ( validate_deepseek_v4_cp, validate_deepseek_v4_mega_moe_token_budget, + validate_deepseek_v41_features, ) + # Before the CP validation: V4.1 rejects CP outright, the actionable message. + validate_deepseek_v41_features(server_args) validate_deepseek_v4_cp(server_args) validate_deepseek_v4_mega_moe_token_budget(server_args) diff --git a/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py b/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py index a289c992e..601c2911b 100644 --- a/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py +++ b/python/sglang/srt/arg_groups/model_overrides/deepseek_v4.py @@ -13,21 +13,44 @@ from sglang.srt.arg_groups.model_override_base import ( ) from sglang.srt.environ import envs from sglang.srt.runtime_context import get_platform +from sglang.srt.utils import is_flashinfer_available logger = logging.getLogger(__name__) @_register_for("DeepseekV4ForCausalLM") def _deepseek_v4_overrides(server_args: Any, hf_config: Any) -> dict: - """DeepSeek V4 attention/page/window/MoE-runner defaults (from - arg_groups/deepseek_v4_hook.py). The kv-cache dtype and NPU split-backend - writes, the max_running_requests fill and the validations stay in the - hook at its legacy slot.""" + """Attention, page and MoE defaults; the rest lives in deepseek_v4_hook.""" cfg = resolving_view(server_args) model_arch = hf_config.architectures[0] overrides: Dict[str, Any] = {"attention_backend": "dsv4"} + # MXFP8 serves this checkpoint's 32-wide ue8m0 blocks on SM100/SM103; + # explicit backend choices, including Triton, take precedence. + quant = getattr(hf_config, "quantization_config", None) or {} + if ( + getattr(hf_config, "model_type", None) == "deepseek_v41" + and cfg.device == "cuda" + and not get_platform().is_hip + and get_platform().is_sm100 + and cfg.fp8_gemm_runner_backend == "auto" + and quant.get("quant_method") == "fp8" + and quant.get("weight_block_size") == [32, 32] + and quant.get("scale_fmt") == "ue8m0" + and is_flashinfer_available() + ): + overrides["fp8_gemm_runner_backend"] = "flashinfer_cutedsl" + logger.info("Use flashinfer_cutedsl for DeepSeek-V4.1 MXFP8 dense GEMMs.") + + # Left unset, the pool configurator sizes the SWA pool from the request cap. + if ( + cfg.swa_full_tokens_ratio is None + and getattr(hf_config, "model_type", None) != "deepseek_v41" + ): + overrides["swa_full_tokens_ratio"] = 0.1 + logger.info(f"Setting swa_full_tokens_ratio to 0.1 for {model_arch}.") + page_size = 256 if cfg.device == "npu": # NPU keeps the device-aware "dsv4" backend (the registry routes it to @@ -43,10 +66,6 @@ def _deepseek_v4_overrides(server_args: Any, hf_config: Any) -> dict: f"Use dsv4 attention backend for {model_arch}, setting page_size to {page_size}." ) - if cfg.swa_full_tokens_ratio is None: - overrides["swa_full_tokens_ratio"] = 0.1 - logger.info(f"Setting swa_full_tokens_ratio to 0.1 for {model_arch}.") - if cfg.moe_runner_backend == "auto": model_config = model_config_of(server_args) # nvidia/DeepSeek-V4-Pro-NVFP4 uses the routed TRT-LLM runner. diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index de96b45e6..144a2ee69 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1014,7 +1014,17 @@ def _flashinfer_allreduce_fusion_auto_enable(view: Any) -> dict: single-node systems. Reads the mid-resolution enable_dp_attention / moe_a2a_backend (after the DeepSeek CP and a2a declarations), exactly like the legacy tail block.""" - model_arch = model_config_of(view).hf_config.architectures[0] + hf_config = model_config_of(view).hf_config + model_arch = hf_config.architectures[0] + # V4.1 TP4 uses the custom push plane for decode and fused MoE finalize. + prefer_custom_dsv41 = ( + getattr(hf_config, "model_type", None) == "deepseek_v41" + and getattr(hf_config, "hidden_size", None) == 5120 + and get_platform().is_blackwell + and view.tp_size == 4 + and view.nnodes == 1 + and not view.disable_custom_all_reduce + ) if envs.SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION.get() and model_arch in { "Qwen3_5MoeForCausalLM", "Qwen3_5MoeForConditionalGeneration", @@ -1032,6 +1042,7 @@ def _flashinfer_allreduce_fusion_auto_enable(view: Any) -> dict: if ( view.flashinfer_allreduce_fusion_backend is None and model_arch in _FLASHINFER_ALLREDUCE_FUSION_ARCHS + and not prefer_custom_dsv41 and (get_platform().is_sm90 or get_platform().is_sm100) and view.tp_size > 1 and not view.enable_dp_attention diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index f1856a4ff..1cf69daa3 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -810,6 +810,7 @@ class TboForwardBatchPreparer: # The child runs the same forward, so it keeps the parent's # sharding verdict; its counts above are already per-child. attn_tp_sequence_sharded=batch.attn_tp_sequence_sharded, + encoder_swa_replay=batch.encoder_swa_replay, tbo_split_seq_index=None, tbo_parent_token_range=(start_token_index, end_token_index), tbo_children=None, diff --git a/python/sglang/srt/configs/deepseek_v4.py b/python/sglang/srt/configs/deepseek_v4.py index 54c83b7dd..087e7971a 100644 --- a/python/sglang/srt/configs/deepseek_v4.py +++ b/python/sglang/srt/configs/deepseek_v4.py @@ -103,8 +103,36 @@ class DeepSeekV4Config(PretrainedConfig): compress_rope_theta: int = 40000 compress_ratios: List[int] = field(default_factory=list) + kv_source_layer_ids: List[int] = field(default_factory=list) + index_source_layer_ids: List[int] = field(default_factory=list) + candidate_source_layer_id: int = -1 + candidate_topk_blocks: int = 0 + candidate_block_size: int = 0 + + engram_layer_ids: List[int] = field(default_factory=list) + engram_num_embeddings: List[int] = field(default_factory=list) + engram_max_ngram_size: int = 1 + engram_vocab_size: int = 0 + engram_n_heads: int = 0 + engram_head_dim: int = 0 + engram_pad_token_id: int = 2 + engram_compressed_vocab_size: int = 0 + + vision_n_layers: int = 0 + vision_dim: int = 1024 + vision_n_heads: int = 16 + vision_inter_dim: int = 2816 + vision_patch_size: int = 14 + vision_rope_theta: float = 10000.0 + vision_downsample_ratio: int = 3 + vision_max_n_token: int = 1024 + vision_min_pixels: int = 295936 + vision_max_wh_ratio: Optional[int] = None + image_token_id: int = 129264 n_hash_layers: int = 3 hc_mult: int = 4 + hc_pre_from_prev_sublayer: bool = False + q_head_norm: bool = True hc_sinkhorn_iters: int = 20 hc_eps: float = 1e-6 diff --git a/python/sglang/srt/configs/deepseek_v41.py b/python/sglang/srt/configs/deepseek_v41.py new file mode 100644 index 000000000..f23951dbb --- /dev/null +++ b/python/sglang/srt/configs/deepseek_v41.py @@ -0,0 +1,88 @@ +"""Translate DeepSeek V4.1 HF configs to the runtime's flat config schema.""" + +from transformers import DeepseekV3Config, PretrainedConfig + +_VISION_FIELDS = { + "num_hidden_layers": "vision_n_layers", + "hidden_size": "vision_dim", + "num_attention_heads": "vision_n_heads", + "intermediate_size": "vision_inter_dim", + "patch_size": "vision_patch_size", + "rope_theta": "vision_rope_theta", + "downsample_ratio": "vision_downsample_ratio", + "max_image_tokens": "vision_max_n_token", + "min_pixels": "vision_min_pixels", + "max_wh_ratio": "vision_max_wh_ratio", +} + + +def _config_dict(config): + return config.to_dict() if isinstance(config, PretrainedConfig) else dict(config) + + +def normalize_deepseek_v41_config(values): + values = dict(values) + text = values.pop("text_config", None) + vision = values.pop("vision_config", None) + if text is not None: + text = _config_dict(text) + text.pop("model_type", None) + values = {**text, **values} + if vision is not None: + vision = _config_dict(vision) + for source, target in _VISION_FIELDS.items(): + if source in vision: + values.setdefault(target, vision[source]) + if "model_type" in values: + values["model_type"] = "deepseek_v41" + if values.get("architectures") == ["DeepseekV41ForCausalLM"]: + values["architectures"] = ["DeepseekV4ForCausalLM"] + return values + + +class DeepseekV41Config(DeepseekV3Config): + # V3 accepts the V4.1 compression ratios; the native V4 config rejects 1/2. + model_type = "deepseek_v41" + vision_n_layers = 0 + hc_pre_from_prev_sublayer = True + q_head_norm = False + kv_source_layer_ids = () + index_source_layer_ids = () + candidate_source_layer_id = -1 + candidate_topk_blocks = 0 + candidate_block_size = 0 + engram_layer_ids = () + engram_num_embeddings = () + engram_max_ngram_size = 1 + engram_vocab_size = 0 + engram_n_heads = 0 + engram_head_dim = 0 + engram_pad_token_id = 2 + engram_compressed_vocab_size = 0 + + def __init__(self, **kwargs): + kwargs = normalize_deepseek_v41_config(kwargs) + kwargs["model_type"] = "deepseek_v41" + super().__init__(**kwargs) + + def to_dict(self): + values = super().to_dict() + values["model_type"] = "deepseek_v41" + return values + + +class DeepseekV41TextConfig(DeepseekV41Config): + model_type = "deepseek_v41_text" + # Transformers regenerates a dataclass initializer unless it is explicit. + __init__ = DeepseekV41Config.__init__ + + +class DeepseekV41VisionConfig(PretrainedConfig): + model_type = "deepseek_v41_vision" + + +DEEPSEEK_V41_CONFIG_CLASSES = ( + DeepseekV41Config, + DeepseekV41TextConfig, + DeepseekV41VisionConfig, +) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index d37dce6b4..417491c56 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -29,6 +29,7 @@ from sglang.srt.disaggregation.base.conn import ( from sglang.srt.disaggregation.utils import ( DisaggregationMode, filter_kv_indices_for_cp_rank, + get_dsv41_spec_layout, ) from sglang.srt.distributed import get_pp_group, get_world_group from sglang.srt.environ import envs @@ -102,6 +103,7 @@ class PrefillServerInfo: kv_cache_dtype: Optional[str] follow_bootstrap_room: bool enable_dsa_cache_layer_split: bool = False + dsv41_spec_layout: Optional[dict] = None # PD true-retraction rebootstrap: the prefill's HTTP API port. The decode # already knows the prefill host (the bootstrap_addr host), so it can POST @@ -152,6 +154,8 @@ class CommonKVManager(BaseKVManager): kv_status_msg_tag: Optional[bytes] = None kv_status_msg_carries_reason: bool = False + dsv41_spec_layout: Optional[dict] = None + # Used by decode when the prefill reported Failed without a reason frame. DEFAULT_PREFILL_FAILURE_REASON = ( "Failed to get kvcache from prefill instance, it might be dead" @@ -166,6 +170,7 @@ class CommonKVManager(BaseKVManager): ): self.kv_args = args self.kv_cache_dtype_str = args.kv_cache_dtype_str + self.dsv41_spec_layout = get_dsv41_spec_layout(args) self.kv_item_lens_sum = sum(args.kv_item_lens) self.state_item_lens_sum = sum(x for comp in args.state_item_lens for x in comp) self.is_mla_backend = is_mla_backend @@ -918,6 +923,27 @@ class CommonKVManager(BaseKVManager): f"Both servers must use the same --kv-cache-dtype value." ) + local_layout = self.dsv41_spec_layout + if local_layout is not None or info.dsv41_spec_layout is not None: + if local_layout != info.dsv41_spec_layout: + mismatched_fields = sorted( + key + for key in (local_layout or {}).keys() + | (info.dsv41_spec_layout or {}).keys() + if (local_layout or {}).get(key) + != (info.dsv41_spec_layout or {}).get(key) + ) + raise RuntimeError( + "DeepSeek-V4.1 DSpark PD layout mismatch " + f"({', '.join(mismatched_fields)}): both servers must " + "enable DSpark with the same block size and target/draft KV " + "layout. Upgrade both servers together." + ) + if info.attn_tp_size != self.attn_tp_size: + raise RuntimeError( + "DeepSeek-V4.1 DSpark PD requires the same TP size on both servers" + ) + if self.dcp_size > 1: if not (self.is_mla_backend or self.is_hybrid_mla_backend): raise RuntimeError( @@ -1079,6 +1105,7 @@ class CommonKVManager(BaseKVManager): "rank_port": self.rank_port, "page_size": self.kv_args.page_size, "kv_cache_dtype": self.kv_cache_dtype_str, + "dsv41_spec_layout": self.dsv41_spec_layout, "load_balance_method": get_parallel().load_balance_method, "enable_dsa_cache_layer_split": get_parallel().enable_dsa_cache_layer_split, # Self-register the HTTP API port so the decode can derive the PD @@ -1969,6 +1996,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): self.dp_size = None self.page_size = None self.kv_cache_dtype: Optional[str] = None + self.dsv41_spec_layout: Optional[dict] = None self.follow_bootstrap_room: Optional[bool] = None self.enable_dsa_cache_layer_split: Optional[bool] = None self.prefill_http_port: Optional[int] = None @@ -2039,6 +2067,14 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): page_size = int(data["page_size"]) kv_cache_dtype = data["kv_cache_dtype"] prefill_http_port = data.get("prefill_http_port") + dsv41_spec_layout = data.get("dsv41_spec_layout") + + if self._registered_count and self.dsv41_spec_layout != dsv41_spec_layout: + return web.Response( + text="DeepSeek-V4.1 DSpark PD layout differs across prefill ranks", + status=400, + ) + self.dsv41_spec_layout = dsv41_spec_layout if self.attn_tp_size is None: self.attn_tp_size = attn_tp_size @@ -2130,6 +2166,7 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): pp_size=self.pp_size, page_size=self.page_size, kv_cache_dtype=self.kv_cache_dtype, + dsv41_spec_layout=self.dsv41_spec_layout, follow_bootstrap_room=( self.follow_bootstrap_room if self.follow_bootstrap_room is not None @@ -2138,7 +2175,10 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): enable_dsa_cache_layer_split=bool(self.enable_dsa_cache_layer_split), prefill_http_port=self.prefill_http_port, ) - return web.json_response(dataclasses.asdict(info), status=200) + payload = dataclasses.asdict(info) + if info.dsv41_spec_layout is None: + payload.pop("dsv41_spec_layout") + return web.json_response(payload, status=200) if not self._is_ready(): return web.Response( diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 69a65ec4e..5548adf0c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -87,6 +87,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( EvictParams, ) from sglang.srt.mem_cache.common import ( + dsv41_dspark_needs_rebootstrap, kv_to_page_indices, page_align_floor, release_kv_cache, @@ -674,6 +675,16 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): if not is_retracted and not is_rebootstrap and is_unadmitted_reject(req): self.scheduler.retire_unadmitted_request(req) return + if is_retracted and dsv41_dspark_needs_rebootstrap( + self.token_to_kv_pool_allocator + ): + if req.output_ids: + req.pd_rebootstrap_forced_output_id = req.output_ids.pop() + req.pd_rebootstrap_in_progress = True + req.time_stats.set_retract_time() + is_retracted = False + is_rebootstrap = True + if self._check_if_req_exceed_kv_capacity(req): return @@ -2837,6 +2848,10 @@ class SchedulerDisaggregationDecodeMixin: # A finished request can still have one redundant forward in flight. # Drain it before a prebuilt request seeds a potentially reused row. self.schedule_stream.wait_stream(self.forward_stream) + # The prebuilt batch never reaches the forward loop's prepare call. + self.ngram_embedding_manager.prepare_for_forward( + new_batch, chunked_req=self.chunked_req + ) new_batch.process_prebuilt(self.future_map) return new_batch diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 06808d23b..d23f45282 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -24,6 +24,7 @@ from sglang.srt.disaggregation.base import KVPoll from sglang.srt.environ import envs from sglang.srt.runtime_context import ( get_disagg, + get_spec, ) from sglang.srt.utils import is_npu @@ -1674,6 +1675,29 @@ def setup_state_kv_args( ) +def get_dsv41_spec_layout(kv_args: KVArgs) -> Optional[dict]: + """Describe the positional transfer layout without pool capacities or pointers.""" + ratios = getattr(kv_args, "mla_compression_ratios", None) or [] + if 2 not in ratios or str(get_spec().speculative_algorithm).upper() != "DSPARK": + return None + + from sglang.srt.disaggregation.base.conn import StateType + + if kv_args.state_types.count(StateType.SWA) != 2: + raise RuntimeError( + "DeepSeek-V4.1 DSpark PD requires target and draft SWA state" + ) + + return { + "num_draft_tokens": get_spec().speculative_num_draft_tokens, + "compression_ratios": list(ratios), + "kv_layer_ids": list(kv_args.kv_layer_ids), + "kv_item_lens": list(kv_args.kv_item_lens), + "state_types": [state_type.value for state_type in kv_args.state_types], + "state_item_lens": [list(items) for items in kv_args.state_item_lens], + } + + def prepare_abort(req: Req, error_message: str, status_code=None): from sglang.srt.managers.schedule_batch import FINISH_ABORT diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 16b5f4b5c..b5d388416 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -2269,6 +2269,7 @@ def _execute_server_warmup(server_args: ServerArgs): bool(model_info.get("has_image_understanding", False)) and not get_disagg().language_only and not get_disagg().language_model_only + and not get_exec().features.enable_encoder_swa_bounded_replay and not is_mps() ) if model_info["is_generation"]: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 50dbc48c3..3ac57ca59 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -549,6 +549,11 @@ class Envs: SGLANG_DSPARK_OPT_MARKOV_W2_BF16 = EnvBool(True) SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD = EnvBool(True) SGLANG_DSPARK_OPT_FUSED_GREEDY_MARKOV = EnvBool(False) + # With the TP-sharded markov_w2, gather each step's vocab-parallel logits over + # the NVLink push collective (CustomAllReduceV2's multicast plane) instead of + # the NCCL ring. Only taken when the group's communicator has a multicast + # plane; off, or no such plane, keeps the NCCL all-gather. + SGLANG_DSPARK_NVLINK_VOCAB_GATHER = EnvBool(True) SGLANG_DSPARK_ENABLE_MULTI_STREAM = EnvBool(True) SGLANG_DSPARK_CONFIDENCE_RELAY_LAG_STEPS = EnvInt(2) @@ -1479,6 +1484,13 @@ class Envs: # Quantize the SWA fp8 KV cache from bf16-rounded values (matches # trainer-side QAT and the DSA-CP path) instead of fp32 registers. SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE = EnvBool(False) + # Paged KV layout of the DeepSeek-V4 family pools: "v4" (584 B/token, every + # GPU), "v41" (the SM100 FlashMLA V4.1 formats: 528 B fp8 SWA cache, fp8 or + # fp4 compressed caches) or "auto" (v41 on SM100 when FlashMLA supports it). + SGLANG_DSV4_KV_LAYOUT = EnvStr("v4") + # Compressed-cache layout under "v41": "auto" (fp4 for the fp4-rounded + # ratio-1 / ratio-2 latents, fp8 for ratios 4 / 128), "fp8" or "fp4" for all. + SGLANG_DSV4_COMPRESSED_KV_LAYOUT = EnvStr("auto") # unified_kv only: split the pool into an fp8 nope pool plus a parallel # bf16 rope pool, 640 B/token instead of 1024. The unified pool takes no # dtype, so --kv-cache-dtype has no effect there and this switch is the @@ -1510,6 +1522,9 @@ class Envs: SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False) SGLANG_EXPERIMENTAL_ONLINE_C128_MTP = EnvBool(False) SGLANG_DSV4_COMPRESS_STATE_DTYPE = EnvStr("float32") + # Run the DeepSeek-V4.1 ratio-1/2 prefill indexer on the torch path instead + # of the DeepGEMM dense fp4 logits kernel (test oracle / fallback). + SGLANG_DSV41_TORCH_PREFILL_INDEXER = EnvBool(False) SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False) SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(True) diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py index 8dbf9c005..fb0609a57 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py @@ -21,6 +21,7 @@ from typing import List, Optional, Tuple import torch import torch_npu +from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.hardware_backend.npu.utils import is_npu_arch35 from sglang.srt.mem_cache.deepseek_v4_compress_state import CompressStatePool @@ -291,6 +292,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): enable_memory_saver: bool, global_page_size: int, cls: type = DeepSeekV4SingleKVPool, + kv_layout: KVLayout = KVLayout.V4, ) -> NPUDeepSeekV4SingleKVPool: # NPU does not use the HiSparse c4 device pool; fail loud if someone # enables it so the silent layout mismatch surfaces at init. @@ -298,6 +300,8 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): "enable_hisparse is not supported on the NPU DSV4 KV pool " f"(got c4 pool class {cls.__name__})." ) + # The V4.1 fp8 / fp4 page layouts are CUDA FlashMLA formats. + assert kv_layout is KVLayout.V4, f"NPU pools do not support {kv_layout}" # Full/SWA use the global page size, C4 uses its native compressed page, # and C128 has an independent physical page size. is_c4_pool = page_size * 4 == global_page_size @@ -382,6 +386,9 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]: """Main PD buffers addressed by the full KV page id.""" + if self.c4_kv_pool is None: + # A draft pool whose layers are all uncompressed has no c4 buffers. + return [], [], [] indexer_pool = self._indexer_pool(4) buffers = ( self.c4_kv_pool.kv_buffer diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 309d43485..96a1cc223 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -15,9 +15,14 @@ from typing import ( Union, ) +import msgspec import torch import torch.nn.functional as F +from sglang.kernels.ops.attention.dsv4 import topk_transform_ragged_v2 +from sglang.kernels.ops.attention.dsv4.decode_attention_sm100 import ( + can_use_swapab_attention, +) from sglang.kernels.ops.attention.dsv4.dequant_k_cache import ( cast_q_fp8_for_q8kv8_prefill, dequantize_k_cache_paged, @@ -25,7 +30,11 @@ from sglang.kernels.ops.attention.dsv4.dequant_k_cache import ( gather_dequant_requant_fp8_paged, q8kv8_padded_num_heads, ) +from sglang.kernels.ops.attention.dsv4.fp4_indexer import fp4_index_logits_decode from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout +from sglang.kernels.ops.attention.dsv4.metadata_kernel import ( + fill_all_compressed_indices, +) from sglang.kernels.ops.attention.dsv4.metadata_kernel import ( init_compression_metadata as _init_compression_metadata_triton, ) @@ -34,6 +43,7 @@ from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import ( BuildCausalSwaPageIndices, BuildPageTablePositions, ExpandPrefillCausally, + late_layer_tail_layout, ) from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import ( BuildBlockSeqLensCausal, @@ -46,12 +56,30 @@ from sglang.srt.layers.attention.base_attn_backend import ( SharedReadEnds, ) from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend +from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp +from sglang.srt.layers.attention.dsv4.candidate_indexer import ( + CandidateMasks, + CandidateMetadata, + IndexerInputs, + make_candidate_indexer, + mask_topk_scores, + published_masks, + select_candidate_blocks, +) from sglang.srt.layers.attention.dsv4.compressor_v2 import ( CompressorBackendMixin, FusedCompressMetadata, create_paged_compressor_data, ) -from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin +from sglang.srt.layers.attention.dsv4.dsv41_sparse import ( + _rope_fq4, + token_req_indices, +) +from sglang.srt.layers.attention.dsv4.indexer import ( + C4IndexerBackendMixin, + deep_gemm_fp4_paged_mqa_logits, + topk_transform_paged_from_metadata, +) from sglang.srt.layers.attention.dsv4.metadata import ( _LARGE_INDEXER_QUERY_THRESHOLD, PagedIndexerMetadata, @@ -67,7 +95,19 @@ from sglang.srt.layers.attention.verify_mask import ( VerifyMask, maybe_create_verify_mask, ) -from sglang.srt.layers.cp.utils import is_cp_active +from sglang.srt.layers.cp.interleave import ( + InterleaveContextParallelMetadata, + interleave_rows_per_request, +) +from sglang.srt.layers.cp.utils import ( + cp_materialize_global_token_order, + is_cp_active, +) +from sglang.srt.layers.dp_attention import ( + get_local_dp_buffer_len, + set_local_dp_buffer_len, +) +from sglang.srt.mem_cache.deepseek_v4_compress_state import KVAndScore from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import ( @@ -103,6 +143,12 @@ DEFAULT_INDEX_TOPK = 512 PAGE_INDEX_ALIGNED_SIZE = 64 +@functools.lru_cache(maxsize=None) +def _is_sm100_or_newer() -> bool: + # DeepGEMM's fp8_fp4 mqa-logits kernels need SM100+; Hopper takes the torch indexer. + return torch.cuda.get_device_capability()[0] >= 10 + + def _get_logical_forward_mode(forward_batch: ForwardBatch) -> ForwardMode: # IDLE is a real per-DP-rank mode. Do not let a stale _original_forward_mode # from a reused/padded ForwardBatch turn an empty rank into TARGET_VERIFY. @@ -155,6 +201,224 @@ def _create_flashmla_metadata(): return flash_mla.get_mla_metadata()[0] +# FlashMLA's head64 sm100 decode scheduling constants; not exported, so they hold +# only for that shape and go stale silently if FlashMLA retunes it. +_FLASHMLA_SCHED_BLOCK_SIZE_N = 64 +_FLASHMLA_SCHED_FIXED_OVERHEAD = 5 + + +@functools.lru_cache(maxsize=None) +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def _fast_flashmla_sched_shape(q: torch.Tensor) -> bool: + return q.is_cuda and get_platform().is_blackwell and q.shape[-2] == 64 + + +def _maybe_precompute_flashmla_sched_meta( + flashmla_metadata, + *, + q: torch.Tensor, + indices: torch.Tensor, + topk_length: Optional[torch.Tensor], + extra_indices: Optional[torch.Tensor], + extra_topk_length: Optional[torch.Tensor], +) -> None: + """Fill the split-KV schedule buffers so `sparse_decode_fwd` skips its own + `<<<1, 32>>>` scheduling kernel on the decode critical path; + `flashmla_sched_meta` produces the same schedule bit for bit.""" + if flashmla_metadata is None: + return + if getattr(flashmla_metadata, "tile_scheduler_metadata", None) is not None: + return + if not _fast_flashmla_sched_shape(q): + return + from sglang.kernels.ops.attention.dsv4.flashmla_sched_meta import ( + META_INTS, + flashmla_sched_meta, + ) + + b, s_q = q.shape[0], q.shape[1] + num_sm_parts = max(_num_sms(q.device.index) // s_q, 1) + meta = torch.empty((num_sm_parts, META_INTS), dtype=torch.int32, device=q.device) + num_splits = torch.empty((b + 1,), dtype=torch.int32, device=q.device) + flashmla_sched_meta( + meta, + num_splits, + topk_length=topk_length, + extra_topk_length=extra_topk_length, + block_size_n=_FLASHMLA_SCHED_BLOCK_SIZE_N, + fixed_overhead_num_blocks=_FLASHMLA_SCHED_FIXED_OVERHEAD, + topk=indices.shape[-1], + extra_topk=0 if extra_indices is None else extra_indices.shape[-1], + ) + flashmla_metadata.tile_scheduler_metadata = meta + flashmla_metadata.num_splits = num_splits + + +def _expand_index_page_table( + page_table: torch.Tensor, + *, + full_page_size: int, + compress_ratio: int, + index_page_size: int, +) -> torch.Tensor: + """Block table of a low-ratio indexer-K pool, which pages at `index_page_size` + slots: [bs, n] -> [bs, n * blocks_per_page] int32. The kernel reads compressed + slot j at page_table[b, j // index_page_size] * index_page_size + j % + index_page_size, which after this expansion is the c1/c2 pool slot of the same + position.""" + slots_per_page = full_page_size // compress_ratio + assert slots_per_page % index_page_size == 0, ( + f"{full_page_size = } / {compress_ratio = } must be a multiple of " + f"{index_page_size = }" + ) + blocks_per_page = slots_per_page // index_page_size + if blocks_per_page == 1: + return page_table + bs, n = page_table.shape + base = page_table.to(torch.int64) * blocks_per_page + offsets = torch.arange(blocks_per_page, device=page_table.device, dtype=torch.int64) + expanded = base.unsqueeze(-1) + offsets # [bs, n, blocks_per_page] + return expanded.reshape(bs, n * blocks_per_page).to(torch.int32) + + +# Arbitrary cap on one bf16 [rows, heads, lc] score chunk; transients run ~3x this. +_TORCH_INDEXER_SCORE_BUDGET_BYTES = 1 << 30 + + +def _every_request_fits() -> bool: + from sglang.srt.model_executor.runner_utils.capture_mode import ( + get_capture_attention_variant, + ) + + # Captured only for batches where every request fits the candidate budget, so + # the plain top-k is the whole selection. + return get_capture_attention_variant() in ( + "candidate_all", + "candidate_c2_all", + "candidate_unfiltered", + ) + + +@functools.cache +def _has_dense_fp4_indexer() -> bool: + if not torch.cuda.is_available() or torch.version.cuda is None: + return False + try: + import deep_gemm + except ImportError: + return False + return hasattr(deep_gemm, "fp8_fp4_mqa_logits") + + +def _dense_fp4_mqa_logits( + q_fp4: Tuple[torch.Tensor, torch.Tensor], + kv_fp4: Tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + ks: torch.Tensor, + ke: torch.Tensor, + max_seqlen_k: int, +) -> torch.Tensor: + from deep_gemm import fp8_fp4_mqa_logits as fn + + # q (int8 [T, H, 64], int32 [T, H]) x kv (int8 [L, 64], int32 [L]) -> fp32 + # [T, max_seqlen_k]; row t column j is k[ks_t + j], garbage past ke_t - ks_t. + return fn(q_fp4, kv_fp4, weights, ks, ke, False, max_seqlen_k) + + +def _low_ratio_source_projections(layer, x, q_lora, positions, bufs): + from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + get_tc_piecewise_forward_context, + ) + + real = ( + get_tc_piecewise_forward_context().forward_batch.global_num_token_non_padded_cpu + ) + if real is None: + real = x.shape[0] + + # These GEMMs pick their algorithm by M, so at the bucket size the live rows + # differ from eager; everything downstream is row-independent. + def put(name, value): + buf = bufs[name] + buf[:real].copy_(value) + buf[real:].zero_() + + if real == 0: + # An idle DP-attention rank replays on fabricated rows with no live + # token; a zero-row GEMM is a launch error, so only zero the buffers. + for buf in bufs.values(): + buf.zero_() + return + + if layer.compressor is not None: + kv, score = layer.compressor.project(x[:real]) + put("kv", kv) + if score is not None: + put("score", score) + if layer.indexer is not None: + indexer = layer.indexer + put("q", indexer.queries(q_lora[:real], layer.freqs_cis[positions[:real]])) + put("w", indexer.head_weights(x[:real])) + + +def _bcg_low_ratio_source_projections(*args): + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import ( + eager_on_graph, + ) + + global _bcg_low_ratio_source_projections_fn + if _bcg_low_ratio_source_projections_fn is None: + _bcg_low_ratio_source_projections_fn = eager_on_graph(True)( + _low_ratio_source_projections + ) + return _bcg_low_ratio_source_projections_fn(*args) + + +_bcg_low_ratio_source_projections_fn = None + + +def _as_int_list(values) -> Optional[List[int]]: + if values is None: + return None + if isinstance(values, torch.Tensor): + if values.device.type != "cpu": + return None + values = values.tolist() + return [int(v) for v in values] + + +def _low_ratio_compression_metadata( + compress_ratio: int, seq_lens_casual: torch.Tensor, raw_out_loc: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + num_write_tokens = raw_out_loc.shape[0] + completes_group = seq_lens_casual[:num_write_tokens] % compress_ratio == 0 + out_loc = torch.where( + completes_group, raw_out_loc.to(torch.int64) // compress_ratio, -1 + ) + topk_lengths_clamp1 = (seq_lens_casual // compress_ratio).clamp_min(1) + return out_loc, topk_lengths_clamp1.to(torch.int32) + + +def _low_ratio_sparse_buffers( + topk_lengths_clamp1: torch.Tensor, topk: int, is_prefill: bool +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + # Not the extra_topk_length the kernel reads: that one comes from positions. + sparse_topk_lengths = torch.clamp(topk_lengths_clamp1, max=topk) + page_indices = _pad_last_dim( + torch.full( + (topk_lengths_clamp1.size(0), topk), + -1, + dtype=torch.int32, + device=topk_lengths_clamp1.device, + ) + ) + raw_indices = torch.empty_like(page_indices) if is_prefill else None + return sparse_topk_lengths, page_indices, raw_indices + + def _create_dummy_paged_compress_data(compress_ratio: int): return None @@ -183,6 +447,7 @@ class DSV4AttnMetadata: # Sorted compress ratios present in this stage; absent ratios keep no # buffers or schedules. present_ratios: Tuple[int, ...] + request_window_layout: Optional[object] = None # Shared by all layer stores; locations are in SWA space. swa_out_cache_loc: Optional[torch.Tensor] = None c4_out_loc: Optional[torch.Tensor] = None @@ -196,6 +461,20 @@ class DSV4AttnMetadata: c128_page_indices: Optional[torch.Tensor] = None c128_topk_lengths_clamp1: Optional[torch.Tensor] = None + # The (1, 2) subset of present_ratios: one latent per ratio tokens, at slot + # raw_out_loc // ratio of the c1 / c2 pool, attended through the extra cache. + low_ratios: Tuple[int, ...] = () + c1_out_loc: Optional[torch.Tensor] = None + c1_topk_lengths_clamp1: Optional[torch.Tensor] = None + c1_sparse_topk_lengths: Optional[torch.Tensor] = field(init=False, default=None) + c1_sparse_page_indices: Optional[torch.Tensor] = field(init=False, default=None) + c1_sparse_raw_indices: Optional[torch.Tensor] = field(init=False, default=None) + c2_out_loc: Optional[torch.Tensor] = None + c2_topk_lengths_clamp1: Optional[torch.Tensor] = None + c2_sparse_topk_lengths: Optional[torch.Tensor] = field(init=False, default=None) + c2_sparse_page_indices: Optional[torch.Tensor] = field(init=False, default=None) + c2_sparse_raw_indices: Optional[torch.Tensor] = field(init=False, default=None) + # Combined decode tables. Only the c4 tail and lens vary by layer. trtllm_swa_lens: Optional[torch.Tensor] = None trtllm_c4_indices: Optional[torch.Tensor] = None @@ -209,6 +488,12 @@ class DSV4AttnMetadata: trtllm_prefill_c128: Optional[tuple] = None c0_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False) + c1_flashmla_metadata: Optional[FlashMLASchedMeta] = field( + init=False, default=None, repr=False + ) + c2_flashmla_metadata: Optional[FlashMLASchedMeta] = field( + init=False, default=None, repr=False + ) c4_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False) c128_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False) @@ -224,9 +509,13 @@ class DSV4AttnMetadata: def has_c128(self) -> bool: return 128 in self.present_ratios - def get_flashmla_metadata(self, compress_ratio: Literal[0, 4, 128]): + def get_flashmla_metadata(self, compress_ratio: Literal[0, 1, 2, 4, 128]): if compress_ratio == 0: return self.c0_flashmla_metadata + elif compress_ratio == 1: + return self.c1_flashmla_metadata + elif compress_ratio == 2: + return self.c2_flashmla_metadata elif compress_ratio == 4: return self.c4_flashmla_metadata elif compress_ratio == 128: @@ -239,17 +528,25 @@ class DSV4AttnMetadata: def sparse_page_indices(self, compress_ratio: int) -> torch.Tensor: """Slots into the ratio's extra cache, -1 padded: the indexer's top-k for - c4, every compressed block up to the position for c128.""" - if compress_ratio == 4: + c1 / c2 / c4, every compressed block up to the position for c128.""" + if compress_ratio == 1: + return self.c1_sparse_page_indices + elif compress_ratio == 2: + return self.c2_sparse_page_indices + elif compress_ratio == 4: return self.c4_sparse_page_indices - if compress_ratio == 128: + elif compress_ratio == 128: return self.c128_page_indices raise ValueError(f"invalid {compress_ratio=}") def sparse_topk_lengths(self, compress_ratio: int) -> torch.Tensor: - if compress_ratio == 4: + if compress_ratio == 1: + return self.c1_sparse_topk_lengths + elif compress_ratio == 2: + return self.c2_sparse_topk_lengths + elif compress_ratio == 4: return self.c4_sparse_topk_lengths - if compress_ratio == 128: + elif compress_ratio == 128: return self.c128_topk_lengths_clamp1 raise ValueError(f"invalid {compress_ratio=}") @@ -257,7 +554,11 @@ class DSV4AttnMetadata: """The top-k as request-local compressed positions, for the sparse prefill workspace; allocated for prefill metadata only. Only the indexer ratios have one (c128 remaps its page indices instead).""" - if compress_ratio == 4: + if compress_ratio == 1: + return self.c1_sparse_raw_indices + elif compress_ratio == 2: + return self.c2_sparse_raw_indices + elif compress_ratio == 4: return self.c4_sparse_raw_indices raise ValueError(f"invalid {compress_ratio=}") @@ -270,7 +571,17 @@ class DSV4AttnMetadata: raw_indices: Optional[torch.Tensor] = None, ) -> None: """Writer counterpart of the accessors above.""" - if compress_ratio == 4: + if compress_ratio == 1: + self.c1_sparse_page_indices = page_indices + self.c1_sparse_topk_lengths = topk_lengths + if raw_indices is not None: + self.c1_sparse_raw_indices = raw_indices + elif compress_ratio == 2: + self.c2_sparse_page_indices = page_indices + self.c2_sparse_topk_lengths = topk_lengths + if raw_indices is not None: + self.c2_sparse_raw_indices = raw_indices + elif compress_ratio == 4: self.c4_sparse_page_indices = page_indices self.c4_sparse_topk_lengths = topk_lengths if raw_indices is not None: @@ -282,6 +593,48 @@ class DSV4AttnMetadata: else: raise ValueError(f"invalid {compress_ratio=}") + def init_trtllm_sparse_buffers(self) -> None: + """Decode tables of 128 SWA columns then compressed KV, -1 for an invalid + index, lens counting all 128 SWA slots; only the c4 tail is per layer.""" + + num_tokens = self.seq_lens_casual.shape[0] + assert self.swa_page_indices.shape == (num_tokens, SWA_WINDOW) + + # VarSeq reads rows to the 64-token tile boundary. Back every live view + # with an aligned parent whose extra rows contain inert values. + n_pad = (num_tokens + 63) // 64 * 64 + + def _tile_padded(fill, src=None, width=None): + shape = (n_pad,) if width is None else (n_pad, width) + buf = torch.full(shape, fill, **self.cuda_int32_kwargs) + if src is not None: + buf[:num_tokens].copy_(src) + return buf[:num_tokens] + + if n_pad != num_tokens: + self.seq_lens_casual = _tile_padded(1, self.seq_lens_casual) + self.swa_page_indices = _tile_padded( + -1, self.swa_page_indices, width=SWA_WINDOW + ) + self.trtllm_swa_lens = _tile_padded(SWA_WINDOW) + if self.c4_sparse_page_indices is not None: + w4 = self.c4_sparse_page_indices.shape[-1] + assert w4 % 4 == 0, f"{w4=}" + # Unwritten c4 rows must remain inert until the per-layer fill. + self.trtllm_c4_indices = _tile_padded(-1, width=SWA_WINDOW + w4) + self.trtllm_c4_indices[:, :SWA_WINDOW].copy_(self.swa_page_indices) + self.trtllm_c4_lens = _tile_padded(SWA_WINDOW) + if self.c128_page_indices is not None: + w128 = self.c128_page_indices.shape[-1] + assert w128 % 4 == 0, f"{w128=}" + self.trtllm_c128_indices = _tile_padded(-1, width=SWA_WINDOW + w128) + self.trtllm_c128_indices[:, :SWA_WINDOW].copy_(self.swa_page_indices) + self.trtllm_c128_indices[:, SWA_WINDOW:].copy_(self.c128_page_indices) + self.trtllm_c128_lens = _tile_padded( + SWA_WINDOW, + (self.c128_topk_lengths_clamp1 + SWA_WINDOW).to(torch.int32), + ) + def copy_(self, other: DSV4AttnMetadata) -> None: copy_metadata( src=other, @@ -291,6 +644,7 @@ class DSV4AttnMetadata: "page_size", "cuda_int32_kwargs", "present_ratios", + "low_ratios", ], copy_fields=[ "raw_out_loc", @@ -308,6 +662,17 @@ class DSV4AttnMetadata: "c4_sparse_topk_lengths", "c4_sparse_page_indices", "c4_sparse_raw_indices", + "c1_out_loc", + "c1_topk_lengths_clamp1", + "c1_sparse_topk_lengths", + "c1_sparse_page_indices", + "c1_sparse_raw_indices", + "c2_out_loc", + "c2_topk_lengths_clamp1", + "c2_sparse_topk_lengths", + "c2_sparse_page_indices", + "c2_sparse_raw_indices", + "request_window_layout", "trtllm_swa_lens", "trtllm_c4_indices", "trtllm_c4_lens", @@ -319,6 +684,8 @@ class DSV4AttnMetadata: # each forward; not copied across replays. "swa_out_cache_loc", "c0_flashmla_metadata", + "c1_flashmla_metadata", + "c2_flashmla_metadata", "c4_flashmla_metadata", "c128_flashmla_metadata", # Eager-only lazy caches are assigned, not content-copied. @@ -334,6 +701,7 @@ class DSV4AttnMetadata: assert self.page_size == other.page_size assert self.cuda_int32_kwargs == other.cuda_int32_kwargs assert self.present_ratios == other.present_ratios + assert self.low_ratios == other.low_ratios tensor_copy_fields = [ "raw_out_loc", @@ -344,6 +712,12 @@ class DSV4AttnMetadata: "c4_topk_lengths_raw", "c4_topk_lengths_clamp1", "c4_sparse_topk_lengths", + "c1_out_loc", + "c1_topk_lengths_clamp1", + "c1_sparse_topk_lengths", + "c2_out_loc", + "c2_topk_lengths_clamp1", + "c2_sparse_topk_lengths", # Preserve graph-captured table addresses; refill c4 per layer. "trtllm_swa_lens", "trtllm_c4_indices", @@ -358,6 +732,8 @@ class DSV4AttnMetadata: "c128_page_indices", "c128_topk_lengths_clamp1", "c0_flashmla_metadata", + "c1_flashmla_metadata", + "c2_flashmla_metadata", "c4_flashmla_metadata", "c128_flashmla_metadata", # Reset eager-only caches so a replay cannot reuse another shape. @@ -382,7 +758,9 @@ class DSV4AttnMetadata: for field_name in reference_assign_fields: setattr(self, field_name, getattr(other, field_name)) - def init_compression_metadata(self, num_tokens: Optional[int] = None) -> None: + def init_compression_metadata( + self, num_tokens: Optional[int] = None, low_ratio_buffers=None + ) -> None: assert self.page_table.dim() == 2 # CP pads causal metadata for per-rank partitioning, while cache-write # locations remain one-per-logical-token. num_tokens tracks that unpadded @@ -425,6 +803,23 @@ class DSV4AttnMetadata: self.swa_page_indices = _pad_last_dim(self.swa_page_indices) + if low_ratio_buffers is not None: + self.c1_out_loc, self.c1_topk_lengths_clamp1 = low_ratio_buffers[:2] + self.c2_out_loc, self.c2_topk_lengths_clamp1 = low_ratio_buffers[4:6] + return + if 1 in self.low_ratios: + self.c1_out_loc, self.c1_topk_lengths_clamp1 = ( + _low_ratio_compression_metadata( + 1, self.seq_lens_casual, self.raw_out_loc + ) + ) + if 2 in self.low_ratios: + self.c2_out_loc, self.c2_topk_lengths_clamp1 = ( + _low_ratio_compression_metadata( + 2, self.seq_lens_casual, self.raw_out_loc + ) + ) + # Cache-write locations stay in global logical order and are intentionally # excluded from CP reindexing. _CP_REINDEX_FIELDS = [ @@ -434,30 +829,42 @@ class DSV4AttnMetadata: "swa_topk_lengths", "page_table", ] - # Same treatment, None for stages without that compress ratio. + # Same treatment, None for models without that compress ratio. _CP_REINDEX_OPTIONAL_FIELDS = [ "c4_topk_lengths_raw", "c4_topk_lengths_clamp1", "c128_page_indices", "c128_topk_lengths_clamp1", + "c1_topk_lengths_clamp1", + "c2_topk_lengths_clamp1", ] _CP_GLOBAL_FIELDS = [ "raw_out_loc", "swa_out_cache_loc", "c4_out_loc", "c128_out_loc", + "c1_out_loc", + "c2_out_loc", ] - def apply_cp_reindex(self, num_tokens: Optional[int] = None) -> None: + def apply_cp_reindex( + self, + num_tokens: Optional[int] = None, + local_index: Optional[torch.Tensor] = None, + ) -> None: cp_rank = get_parallel().attn_cp_rank cp_size = get_parallel().attn_cp_size - idx = slice(cp_rank, None, cp_size) pre_global_len = self.seq_lens_casual.shape[0] - assert pre_global_len % cp_size == 0, ( - f"apply_cp_reindex: global token count {pre_global_len} is not divisible by cp_size={cp_size}. " - "CP round-robin requires padding to ensure divisibility." - ) - expected_local_len = pre_global_len // cp_size + if local_index is not None: + idx = local_index + expected_local_len = local_index.shape[0] + else: + idx = slice(cp_rank, None, cp_size) + assert pre_global_len % cp_size == 0, ( + f"apply_cp_reindex: global token count {pre_global_len} is not divisible by cp_size={cp_size}. " + "CP round-robin requires padding to ensure divisibility." + ) + expected_local_len = pre_global_len // cp_size if num_tokens is None: num_tokens = pre_global_len for field_name in self._CP_REINDEX_FIELDS + self._CP_REINDEX_OPTIONAL_FIELDS: @@ -485,9 +892,7 @@ class DSV4AttnMetadata: f"!= num_tokens={num_tokens} (must remain global for compressor write path)" ) - def init_flashmla_related(self, is_prefill: bool = False): - # index_topk is set from model_config.index_topk per-model - # (small model: 512, large model: 1024). + def init_flashmla_related(self, is_prefill: bool = False, low_ratio_buffers=None): assert self.index_topk in (512, 1024), ( f"unexpected index_topk={self.index_topk}; " "supported: 512 (small) or 1024 (large)" @@ -517,51 +922,84 @@ class DSV4AttnMetadata: self.c128_flashmla_metadata = ( _create_flashmla_metadata() if self.has_c128 else None ) - - def init_trtllm_sparse_buffers(self) -> None: - """Build decode tables with 128 SWA columns followed by compressed KV. - - Indices use -1 for invalid entries; lens include all 128 SWA slots. - Only the c4 tail and lens are filled per layer. - """ - - num_tokens = self.seq_lens_casual.shape[0] - assert self.swa_page_indices.shape == (num_tokens, SWA_WINDOW) - - # VarSeq reads rows to the 64-token tile boundary. Back every live view - # with an aligned parent whose extra rows contain inert values. - n_pad = (num_tokens + 63) // 64 * 64 - - def _tile_padded(fill, src=None, width=None): - shape = (n_pad,) if width is None else (n_pad, width) - buf = torch.full(shape, fill, **self.cuda_int32_kwargs) - if src is not None: - buf[:num_tokens].copy_(src) - return buf[:num_tokens] - - if n_pad != num_tokens: - self.seq_lens_casual = _tile_padded(1, self.seq_lens_casual) - self.swa_page_indices = _tile_padded( - -1, self.swa_page_indices, width=SWA_WINDOW + if low_ratio_buffers is not None: + assert not is_prefill and self.low_ratios == (1, 2) + self.c1_sparse_topk_lengths, self.c1_sparse_page_indices = ( + low_ratio_buffers[2:4] ) - self.trtllm_swa_lens = _tile_padded(SWA_WINDOW) - if self.c4_sparse_page_indices is not None: - w4 = self.c4_sparse_page_indices.shape[-1] - assert w4 % 4 == 0, f"{w4=}" - # Unwritten c4 rows must remain inert until the per-layer fill. - self.trtllm_c4_indices = _tile_padded(-1, width=SWA_WINDOW + w4) - self.trtllm_c4_indices[:, :SWA_WINDOW].copy_(self.swa_page_indices) - self.trtllm_c4_lens = _tile_padded(SWA_WINDOW) - if self.c128_page_indices is not None: - w128 = self.c128_page_indices.shape[-1] - assert w128 % 4 == 0, f"{w128=}" - self.trtllm_c128_indices = _tile_padded(-1, width=SWA_WINDOW + w128) - self.trtllm_c128_indices[:, :SWA_WINDOW].copy_(self.swa_page_indices) - self.trtllm_c128_indices[:, SWA_WINDOW:].copy_(self.c128_page_indices) - self.trtllm_c128_lens = _tile_padded( - SWA_WINDOW, - (self.c128_topk_lengths_clamp1 + SWA_WINDOW).to(torch.int32), + self.c2_sparse_topk_lengths, self.c2_sparse_page_indices = ( + low_ratio_buffers[6:8] ) + self.c1_flashmla_metadata = _create_flashmla_metadata() + self.c2_flashmla_metadata = _create_flashmla_metadata() + return + if 1 in self.low_ratios: + ( + self.c1_sparse_topk_lengths, + self.c1_sparse_page_indices, + self.c1_sparse_raw_indices, + ) = _low_ratio_sparse_buffers( + self.c1_topk_lengths_clamp1, self.index_topk, is_prefill + ) + self.c1_flashmla_metadata = _create_flashmla_metadata() + if 2 in self.low_ratios: + ( + self.c2_sparse_topk_lengths, + self.c2_sparse_page_indices, + self.c2_sparse_raw_indices, + ) = _low_ratio_sparse_buffers( + self.c2_topk_lengths_clamp1, self.index_topk, is_prefill + ) + self.c2_flashmla_metadata = _create_flashmla_metadata() + + +class LateLayerTail(msgspec.Struct, frozen=True): + """Rows the layers after the last kv_source layer run over under decoder SWA + bounded replay: the last tail tokens of each request in the extend.""" + + token_indices: torch.Tensor + positions: torch.Tensor + extend_seq_lens: torch.Tensor + extend_seq_lens_cpu: List[int] + swa_out_cache_loc: torch.Tensor + # Set when the tail is the extend's last rows (one request): a view, not a gather. + contiguous_start: Optional[int] = None + # prefill CP: this rank's tail rows padded to the largest share; cp_metadata is that layout + pad_rows: int = 0 + cp_metadata: Optional[InterleaveContextParallelMetadata] = None + local_lens_cpu: Optional[List[int]] = None + req_global: Optional[torch.Tensor] = None + pos_global: Optional[torch.Tensor] = None + + def rows(self, t: torch.Tensor) -> torch.Tensor: + rows = self.real_rows(t) + if self.pad_rows: + rows = torch.cat([rows, rows.new_zeros((self.pad_rows, *rows.shape[1:]))]) + return rows + + def real_rows(self, t: torch.Tensor) -> torch.Tensor: + return _tail_rows( + t, token_indices=self.token_indices, contiguous_start=self.contiguous_start + ) + + +def _tail_rows( + t: torch.Tensor, *, token_indices: torch.Tensor, contiguous_start: Optional[int] +) -> torch.Tensor: + if contiguous_start is not None: + return t[contiguous_start:] + return t[token_indices] + + +# Rows per logits chunk for the ratio-1/2 indexer inside the prefill CUDA graph; +# its width is the graph's max_seq_len, and longer contexts replay eagerly. +_PREFILL_GRAPH_INDEXER_ROW_CHUNK = 2048 + + +def _prefill_graph_max_seq_len() -> Optional[int]: + from sglang.srt.runtime_context import get_exec + + return get_exec().graph.cuda_graph_config.prefill.max_seq_len @dataclass @@ -569,14 +1007,33 @@ class DSV4Metadata: core_attn_metadata: DSV4AttnMetadata indexer_metadata: Optional[PagedIndexerMetadata] + # Low-ratio paged indexer metadata; ratio 4 uses indexer_metadata above. + c1_indexer_metadata: Optional[PagedIndexerMetadata] = None + c2_indexer_metadata: Optional[PagedIndexerMetadata] = None + c4_compress_metadata: Optional[FusedCompressMetadata] = None c128_compress_metadata: Optional[FusedCompressMetadata] = None + # Shared by all low-ratio source layers; graph replay refreshes them live. + low_ratio_req_indices: Optional[torch.Tensor] = None + low_ratio_pos_i64: Optional[torch.Tensor] = None + + # Per-step scratch for TP-padded query heads, zeroed by the first user. + # Later layers overwrite real heads and preserve the zero padding. + q_pad_buffer: Optional[torch.Tensor] = None + + # What the candidate-source layer published for the index-source layers after + # it, in the implementation's own type; never copied from the host. + candidate_metadata: Optional[CandidateMetadata] = None + # Built at the runner's prefill WAR boundary when the fast path is on, # otherwise lazily by ``_forward_prefill_sparse``. sparse_prefill_cache: Optional[SparsePrefillChunkCache] = None prefill_shared_reads_snapshotted: bool = False + # Set only on the metadata built for the late layers under bounded SWA replay. + late_layer_tail: Optional[LateLayerTail] = None + @property def core_metadata(self) -> DSV4AttnMetadata: return self.core_attn_metadata @@ -584,6 +1041,8 @@ class DSV4Metadata: def copy_(self, other: DSV4Metadata): self.core_attn_metadata.copy_(other.core_attn_metadata) maybe_copy_inplace(self.indexer_metadata, src=other.indexer_metadata) + maybe_copy_inplace(self.c1_indexer_metadata, src=other.c1_indexer_metadata) + maybe_copy_inplace(self.c2_indexer_metadata, src=other.c2_indexer_metadata) maybe_copy_inplace(self.c4_compress_metadata, src=other.c4_compress_metadata) maybe_copy_inplace( self.c128_compress_metadata, src=other.c128_compress_metadata @@ -596,6 +1055,18 @@ class DSV4Metadata: static_metadata.core_attn_metadata ) maybe_copy_inplace(self.indexer_metadata, src=static_metadata.indexer_metadata) + maybe_copy_inplace( + self.c1_indexer_metadata, src=static_metadata.c1_indexer_metadata + ) + maybe_copy_inplace( + self.c2_indexer_metadata, src=static_metadata.c2_indexer_metadata + ) + maybe_copy_inplace( + self.low_ratio_req_indices, src=static_metadata.low_ratio_req_indices + ) + maybe_copy_inplace( + self.low_ratio_pos_i64, src=static_metadata.low_ratio_pos_i64 + ) maybe_copy_inplace( self.c4_compress_metadata, src=static_metadata.c4_compress_metadata ) @@ -706,6 +1177,7 @@ class DeepseekV4AttnBackend( ): super().__init__() self.model_runner = model_runner + self.encoder_replay = False self.device = torch.device(model_runner.device) self.max_context_len = model_runner.model_config.context_len head_dim = model_runner.model_config.head_dim @@ -727,8 +1199,18 @@ class DeepseekV4AttnBackend( self.req_to_token = model_runner.req_to_token_pool.req_to_token # Nothing is built for a compress ratio outside the pool's set. self.present_ratios: Tuple[int, ...] = self.token_to_kv_pool.present_ratios + self.low_ratios: Tuple[int, ...] = tuple( + ratio for ratio in (1, 2) if ratio in self.present_ratios + ) self.has_c4: bool = 4 in self.present_ratios self.has_c128: bool = 128 in self.present_ratios + # Two-level low-ratio indexer (dsv4/candidate_indexer.py). + cfg = model_runner.model_config.hf_text_config + self.is_dsv41: bool = getattr(cfg, "model_type", None) == "deepseek_v41" + self.candidate_indexer = make_candidate_indexer( + getattr(cfg, "candidate_topk_blocks", 0), + getattr(cfg, "candidate_block_size", 0), + ) self.MAX_SEQ_LEN_FOR_CAPTURE = self.req_to_token.shape[1] assert isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool) @@ -738,6 +1220,11 @@ class DeepseekV4AttnBackend( kernel = get_exec().kernel self.enable_deepseek_v4_fp4_indexer = kernel.enable_deepseek_v4_fp4_indexer + self.enable_decoder_swa_bounded_replay: bool = ( + get_exec().features.enable_decoder_swa_bounded_replay + ) + # The model switches onto this metadata in enter_late_layer_tail. + self.tail_forward_metadata: Optional[DSV4Metadata] = None self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner) self.dsv4_prefill_backend = getattr(kernel, "dsv4_prefill_backend", "auto") if use_dsv4_q8kv8_sparse_prefill(self.dsv4_prefill_backend): @@ -866,13 +1353,37 @@ class DeepseekV4AttnBackend( self, core_attn_metadata: DSV4AttnMetadata, *, + compress_ratio: int = 4, use_prefill_cuda_graph: bool = False, ): + page_table = core_attn_metadata.page_table + index_page_size = 0 + if compress_ratio == 4: + c_seq_lens = core_attn_metadata.c4_topk_lengths_raw + elif compress_ratio in (1, 2): + c_seq_lens = ( + core_attn_metadata.c1_topk_lengths_clamp1 + if compress_ratio == 1 + else core_attn_metadata.c2_topk_lengths_clamp1 + ) + # The low-ratio indexer-K pool pages at 64 slots, not page_size // + # ratio, so the kernel needs a block table at that granularity. + index_page_size = self.token_to_kv_pool.get_index_k_page_size( + compress_ratio + ) + page_table = _expand_index_page_table( + page_table, + full_page_size=self.page_size, + compress_ratio=compress_ratio, + index_page_size=index_page_size, + ) + else: + raise ValueError(f"Unsupported indexer {compress_ratio = }") return PagedIndexerMetadata( page_size=self.page_size, - compressed_page_size=self.token_to_kv_pool.get_index_k_page_size(), - page_table=core_attn_metadata.page_table, - compressed_seq_lens=core_attn_metadata.c4_topk_lengths_raw, + compressed_page_size=index_page_size or self.page_size // compress_ratio, + page_table=page_table, + compressed_seq_lens=c_seq_lens, use_topk_v2=self.dsa_topk_backend.should_use_topk_v2() and not _is_xpu, # The SM120 FP4 kernel schedules split_kv=128, while the generic # JIT metadata planner encodes split_kv=256. @@ -880,6 +1391,7 @@ class DeepseekV4AttnBackend( self.enable_deepseek_v4_fp4_indexer and get_platform().is_sm120 ), use_prefill_cuda_graph=use_prefill_cuda_graph, + compress_ratio=compress_ratio, ) def init_forward_metadata_decode( @@ -915,13 +1427,25 @@ class DeepseekV4AttnBackend( online_c128_state_slot_offset: int = 0, dspark_block_size: Optional[int] = None, forward_batch: Optional[ForwardBatch] = None, + swa_replay_start: Optional[torch.Tensor] = None, + cp_metadata: Optional[InterleaveContextParallelMetadata] = None, + dspark_swa_buffers: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> DSV4Metadata: padded_num_tokens = out_cache_loc.shape[0] cp_active = forward_batch is not None and is_cp_active(forward_batch) if cp_active: - cp_metadata = forward_batch.attn_cp_metadata + if cp_metadata is None: + cp_metadata = forward_batch.attn_cp_metadata assert cp_metadata is not None padded_num_tokens = sum(cp_metadata.per_rank_actual_token) + if ( + swa_replay_start is not None + and swa_replay_start.shape[0] < padded_num_tokens + ): + swa_replay_start = torch.nn.functional.pad( + swa_replay_start, + (0, padded_num_tokens - swa_replay_start.shape[0]), + ) seq_lens_casual, req_pool_indices_repeated = self.expand_prefill_casually( num_tokens=num_tokens, @@ -942,10 +1466,15 @@ class DeepseekV4AttnBackend( need_compress=need_compress, is_prefill=True, dspark_block_size=dspark_block_size, + dspark_swa_buffers=dspark_swa_buffers, num_tokens=num_tokens if cp_active else None, + swa_replay_start=swa_replay_start, + num_groups=len(extend_seq_lens_cpu), ) if cp_active: - core_attn_metadata.apply_cp_reindex(num_tokens=num_tokens) + core_attn_metadata.apply_cp_reindex( + num_tokens=num_tokens, local_index=cp_metadata.local_index + ) core_attn_metadata.init_flashmla_related(is_prefill=True) indexer_metadata = ( self.init_forward_metadata_indexer( @@ -995,7 +1524,7 @@ class DeepseekV4AttnBackend( online_state_slot_offset=online_c128_state_slot_offset, ) - return DSV4Metadata( + metadata = DSV4Metadata( core_attn_metadata, indexer_metadata, c4_compress_metadata=create(compress_ratio=4) if self.has_c4 else None, @@ -1003,6 +1532,283 @@ class DeepseekV4AttnBackend( create(compress_ratio=128) if self.has_c128 else None ), ) + if use_prefill_cuda_graph and self.low_ratio_prefill_graph: + low = core_attn_metadata.low_ratios + metadata.c1_indexer_metadata = ( + self._low_ratio_prefill_indexer_metadata(core_attn_metadata, 1) + if 1 in low + else None + ) + metadata.c2_indexer_metadata = ( + self._low_ratio_prefill_indexer_metadata(core_attn_metadata, 2) + if 2 in low + else None + ) + metadata.low_ratio_req_indices = req_pool_indices_repeated.to(torch.int64) + metadata.low_ratio_pos_i64 = core_attn_metadata.positions_casual.to( + torch.int64 + ) + return metadata + + def _low_ratio_prefill_indexer_metadata( + self, core: DSV4AttnMetadata, compress_ratio: int + ) -> PagedIndexerMetadata: + num_pages = core.page_table.shape[1] + max_seq_len = _prefill_graph_max_seq_len() + if max_seq_len is not None: + num_pages = min(max_seq_len // self.page_size, num_pages) + index_page_size = self.token_to_kv_pool.get_index_k_page_size(compress_ratio) + page_table = _expand_index_page_table( + core.page_table[:, :num_pages], + full_page_size=self.page_size, + compress_ratio=compress_ratio, + index_page_size=index_page_size, + ) + # Unclamped: a token with no completed group scores nothing, as in eager. + c_seq_lens = (core.seq_lens_casual // compress_ratio).to(torch.int32) + row_chunk = _PREFILL_GRAPH_INDEXER_ROW_CHUNK + return PagedIndexerMetadata( + page_size=self.page_size, + compressed_page_size=index_page_size, + page_table=page_table, + compressed_seq_lens=c_seq_lens, + use_topk_v2=False, + use_prefill_cuda_graph=True, + compress_ratio=compress_ratio, + row_chunk=row_chunk if row_chunk < c_seq_lens.shape[0] else 0, + ) + + @property + def low_ratio_prefill_graph(self) -> bool: + return ( + bool(self.low_ratios) and _has_dense_fp4_indexer() and _is_sm100_or_newer() + ) + + def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool: + max_seq_len = _prefill_graph_max_seq_len() + seq_lens_cpu = forward_batch.seq_lens_cpu + if max_seq_len is None or seq_lens_cpu is None or seq_lens_cpu.numel() == 0: + return True + return int(seq_lens_cpu.max().item()) <= max_seq_len + + def _build_late_layer_tail_metadata( + self, forward_batch: ForwardBatch + ) -> DSV4Metadata: + # Each request contributes only its last SWA_WINDOW extend tokens, with the + # window floored at the tail start: window KV before it is never written here. + extend_lens_cpu = forward_batch.extend_seq_lens_cpu + seq_lens_cpu = forward_batch.seq_lens_cpu + assert extend_lens_cpu is not None and seq_lens_cpu is not None + device = forward_batch.out_cache_loc.device + token_indices, tail_lens_cpu, swa_replay_start = late_layer_tail_layout( + extend_lens_cpu=extend_lens_cpu, + seq_lens_cpu=seq_lens_cpu.tolist(), + tail_len=SWA_WINDOW, + device=device, + ) + contiguous_start = ( + extend_lens_cpu[0] - tail_lens_cpu[0] if len(extend_lens_cpu) == 1 else None + ) + out_cache_loc = _tail_rows( + forward_batch.out_cache_loc, + token_indices=token_indices, + contiguous_start=contiguous_start, + ) + tail_lens = torch.tensor(tail_lens_cpu, dtype=torch.int32, device=device) + cp_tail = ( + self._late_layer_tail_cp_layout(forward_batch, token_indices, tail_lens) + if is_cp_active(forward_batch) + else None + ) + + metadata = self.init_forward_metadata_prefill( + max_seq_len=int(seq_lens_cpu.max().item()), + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens.to(torch.int32), + seq_lens_cpu=seq_lens_cpu.tolist(), + out_cache_loc=out_cache_loc, + num_tokens=sum(tail_lens_cpu), + extend_seq_lens=tail_lens, + extend_seq_lens_cpu=tail_lens_cpu, + extend_start_loc=torch.cumsum(tail_lens, dim=0) - tail_lens, + swa_replay_start=swa_replay_start, + forward_batch=forward_batch if cp_tail is not None else None, + cp_metadata=cp_tail["cp_metadata"] if cp_tail is not None else None, + ) + swa_out_cache_loc = ( + metadata.core_attn_metadata.request_window_layout.write_loc + if self.token_to_kv_pool.request_window is not None + else self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( + torch.int32 + ) + ) + metadata.core_attn_metadata.swa_out_cache_loc = swa_out_cache_loc + metadata.low_ratio_req_indices = torch.repeat_interleave( + forward_batch.req_pool_indices.to(torch.int64), tail_lens.to(torch.int64) + ) + positions = _tail_rows( + forward_batch.positions, + token_indices=token_indices, + contiguous_start=contiguous_start, + ) + metadata.low_ratio_pos_i64 = positions.to(torch.int64) + if cp_tail is None: + # Without CP, tail rows index the full extend on this rank. + metadata.late_layer_tail = LateLayerTail( + token_indices=token_indices, + positions=positions, + extend_seq_lens=tail_lens, + extend_seq_lens_cpu=tail_lens_cpu, + swa_out_cache_loc=swa_out_cache_loc, + contiguous_start=contiguous_start, + ) + else: + # With CP, select from this rank's extend and pad for collectives. + metadata.late_layer_tail = LateLayerTail( + token_indices=cp_tail["local_token_indices"], + positions=cp_tail["local_positions"], + extend_seq_lens=tail_lens, + extend_seq_lens_cpu=tail_lens_cpu, + swa_out_cache_loc=swa_out_cache_loc, + pad_rows=cp_tail["pad_rows"], + cp_metadata=cp_tail["cp_metadata"], + local_lens_cpu=cp_tail["local_lens_cpu"], + req_global=metadata.low_ratio_req_indices, + pos_global=metadata.low_ratio_pos_i64, + ) + return metadata + + def _late_layer_tail_cp_layout( + self, + forward_batch: ForwardBatch, + token_indices: torch.Tensor, + tail_lens: torch.Tensor, + ) -> dict: + cp_rank = get_parallel().attn_cp_rank + cp_size = get_parallel().attn_cp_size + device = token_indices.device + total = token_indices.shape[0] + owner_rank = token_indices % cp_size + counts = torch.bincount(owner_rank, minlength=cp_size).tolist() + max_local = max(counts) + order = torch.argsort(owner_rank, stable=True) + rank_starts = torch.tensor( + [sum(counts[:r]) for r in range(cp_size)], + dtype=torch.int64, + device=device, + ) + slot = torch.empty_like(owner_rank) + slot[order] = ( + torch.arange(total, device=device) - rank_starts[owner_rank[order]] + ) + gather_index = owner_rank * max_local + slot + + local_tail_rows = (owner_rank == cp_rank).nonzero().squeeze(1) + pad_rows = max_local - counts[cp_rank] + # Give each rank distinct padding rows in the compact tail metadata. + pad_start = total + sum(max_local - c for c in counts[:cp_rank]) + local_metadata_rows = torch.cat( + [ + local_tail_rows, + torch.arange(pad_start, pad_start + pad_rows, device=device), + ] + ) + tail_request_ids = torch.repeat_interleave( + torch.arange(forward_batch.batch_size, device=device), + tail_lens.to(torch.int64), + output_size=total, + ) + local_positions = torch.cat( + [ + forward_batch.positions[token_indices[local_tail_rows]], + forward_batch.positions.new_zeros(pad_rows), + ] + ) + cp_metadata = InterleaveContextParallelMetadata( + per_rank_actual_token=[max_local] * cp_size, + max_rank_len=[max_local] * cp_size, + total_seq_lens=total, + bs=forward_batch.batch_size, + per_rank_logical_token=counts, + gather_index=gather_index, + local_index=local_metadata_rows, + ) + return dict( + cp_metadata=cp_metadata, + local_token_indices=(token_indices[local_tail_rows] - cp_rank) // cp_size, + local_positions=local_positions, + local_lens_cpu=torch.bincount( + tail_request_ids[local_tail_rows], minlength=forward_batch.batch_size + ).tolist(), + pad_rows=pad_rows, + ) + + def enter_late_layer_tail(self, forward_batch: ForwardBatch) -> tuple: + """Switch the late layers onto the tail; the return value goes back to + exit_late_layer_tail.""" + tail_metadata = self.tail_forward_metadata + assert tail_metadata is not None, "no tail metadata for this forward" + saved = ( + self.forward_metadata, + forward_batch.attn_cp_metadata, + get_local_dp_buffer_len(), + ) + tail = tail_metadata.late_layer_tail + tail_lens_cpu = ( + tail.local_lens_cpu + if tail.cp_metadata is not None + else tail.extend_seq_lens_cpu + ) + # TODO(candidate): goes away once the source publishes its tail rows straight + # onto the tail metadata (publish_prefill); until then cut the full masks. + full_masks = self.forward_metadata.candidate_metadata + if isinstance(full_masks, CandidateMasks) and full_masks.request_masks: + tail_metadata.candidate_metadata = CandidateMasks( + request_masks=[ + mask[mask.shape[0] - t :] + for mask, t in zip(full_masks.request_masks, tail_lens_cpu) + ] + ) + # The layers before the switch published top-k into the full metadata's + # buffers; carry the tail rows into the tail metadata's (padding stays -1). + full_core = saved[0].core_attn_metadata + tail_core = tail_metadata.core_attn_metadata + for ratio in tail_core.low_ratios: + for full_buf, tail_buf in ( + ( + full_core.sparse_page_indices(ratio), + tail_core.sparse_page_indices(ratio), + ), + ( + full_core.sparse_topk_lengths(ratio), + tail_core.sparse_topk_lengths(ratio), + ), + ( + full_core.sparse_raw_indices(ratio), + tail_core.sparse_raw_indices(ratio), + ), + ): + if full_buf is None or tail_buf is None: + continue + rows = tail.real_rows(full_buf) + tail_buf[: rows.shape[0]].copy_(rows) + self.forward_metadata = tail_metadata + if self.token_to_kv_pool.request_window is not None: + self.token_to_kv_pool.request_window.activate( + tail_core.request_window_layout + ) + if tail.cp_metadata is not None: + forward_batch.attn_cp_metadata = tail.cp_metadata + set_local_dp_buffer_len(sum(tail.cp_metadata.per_rank_actual_token)) + return saved + + def exit_late_layer_tail(self, saved: tuple, forward_batch: ForwardBatch) -> None: + ( + self.forward_metadata, + forward_batch.attn_cp_metadata, + local_dp_buffer_len, + ) = saved + set_local_dp_buffer_len(local_dp_buffer_len) def init_forward_metadata_target_verify( self, @@ -1063,6 +1869,7 @@ class DeepseekV4AttnBackend( seq_lens_cpu: Optional[torch.Tensor], out_cache_loc: torch.Tensor, block_size: int, + dspark_swa_buffers: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> DSV4Metadata: if seq_lens_cpu is None: seq_lens_cpu_list = seq_lens.tolist() @@ -1075,7 +1882,9 @@ class DeepseekV4AttnBackend( ) extend_seq_lens = self._move_to_device(lengths.extend_seq_lens_cpu) return self.init_forward_metadata_prefill( - max_seq_len=max_seq_len, + # DSpark draft blocks are SWA-only, like draft extend. Their full + # context page table is unused; retain only its 2-D placeholder. + max_seq_len=self.page_size, req_pool_indices=req_pool_indices, seq_lens=lengths.seq_lens_extended, seq_lens_cpu=lengths.seq_lens_cpu_extended, @@ -1087,6 +1896,7 @@ class DeepseekV4AttnBackend( need_compress=False, use_prefill_cuda_graph=False, dspark_block_size=block_size, + dspark_swa_buffers=dspark_swa_buffers, ) def make_forward_metadata_from_raw_verify( @@ -1135,6 +1945,7 @@ class DeepseekV4AttnBackend( max_seq_len=self.MAX_SEQ_LEN_FOR_CAPTURE, out_loc=out_cache_loc, need_compress=True, + num_groups=bs, ) indexer_metadata = ( self.init_forward_metadata_indexer(core_attn_metadata) @@ -1158,9 +1969,20 @@ class DeepseekV4AttnBackend( c128_compress_metadata = raw_metadata.c128_compress_metadata if c128_compress_metadata is None and self.has_c128: c128_compress_metadata = create(compress_ratio=128) + low = core_attn_metadata.low_ratios return DSV4Metadata( core_attn_metadata, indexer_metadata, + c1_indexer_metadata=( + self.init_forward_metadata_indexer(core_attn_metadata, compress_ratio=1) + if 1 in low + else None + ), + c2_indexer_metadata=( + self.init_forward_metadata_indexer(core_attn_metadata, compress_ratio=2) + if 2 in low + else None + ), c4_compress_metadata=create(compress_ratio=4) if self.has_c4 else None, c128_compress_metadata=c128_compress_metadata, ) @@ -1187,6 +2009,18 @@ class DeepseekV4AttnBackend( else None ) + low = core_attn_metadata.low_ratios + c1_indexer_metadata = ( + self.init_forward_metadata_indexer(core_attn_metadata, compress_ratio=1) + if 1 in low + else None + ) + c2_indexer_metadata = ( + self.init_forward_metadata_indexer(core_attn_metadata, compress_ratio=2) + if 2 in low + else None + ) + create = functools.partial( create_paged_compressor_data, is_prefill=False, @@ -1199,6 +2033,8 @@ class DeepseekV4AttnBackend( return DSV4Metadata( core_attn_metadata, indexer_metadata, + c1_indexer_metadata=c1_indexer_metadata, + c2_indexer_metadata=c2_indexer_metadata, c4_compress_metadata=create(compress_ratio=4) if self.has_c4 else None, c128_compress_metadata=( create(compress_ratio=128) if self.has_c128 else None @@ -1218,9 +2054,12 @@ class DeepseekV4AttnBackend( if swa_out_cache_loc is None and out_cache_loc is not None: # Eager-only miss (no graph state / oversized batch): translate once # per step instead of per layer at store time. - swa_out_cache_loc = self.token_to_kv_pool.translate_loc_from_full_to_swa( - out_cache_loc - ).to(torch.int32) + if self.token_to_kv_pool.request_window is None: + swa_out_cache_loc = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + out_cache_loc + ).to(torch.int32) + ) if out_cache_loc is None: out_cache_loc = seq_lens.new_zeros(num_tokens) @@ -1246,6 +2085,7 @@ class DeepseekV4AttnBackend( out_loc=out_cache_loc, need_compress=False, is_prefill=True, + num_groups=batch_size, ) if swa_out_cache_loc is not None: # Captures store_cache's cached path instead of a per-layer @@ -1278,7 +2118,12 @@ class DeepseekV4AttnBackend( return buf[:n] def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: - # Raw metadata must be materialized inside the graph to refresh on replay. + from sglang.srt.model_executor.runner_utils.capture_mode import ( + skip_low_ratio_indexer, + ) + + # Upgrade Raw->Full so compress + core_attn + indexer materialization is + # recorded inside the cuda graph; already Full when PREP_IN_CUDA_GRAPH=0. if isinstance(self.forward_metadata, DSV4RawVerifyMetadata): self.forward_metadata = self.make_forward_metadata_from_raw_verify( raw_metadata=self.forward_metadata, @@ -1289,10 +2134,23 @@ class DeepseekV4AttnBackend( raw_metadata=self.forward_metadata, ) - # Spec-v2 and DP padding can rebind out_cache_loc after out-graph prep; - # capture the translation here so replay reads live locations. - # FlashMLA requires int32 indices. metadata = self.forward_metadata + if isinstance(metadata, DSV4Metadata): + core = metadata.core_metadata + for ratio in core.low_ratios: + if skip_low_ratio_indexer(ratio): + # Share the full-position indices across layers of this ratio. + fill_all_compressed_indices( + core.page_table, + core.sparse_topk_lengths(ratio), + core.sparse_page_indices(ratio), + compress_ratio=ratio, + page_size=core.page_size, + raw_indices=core.sparse_raw_indices(ratio), + ) + + # Recorded inside the cuda graph, so replay re-reads the live out_cache_loc + # buffer (spec-v2 and DP padding rebind it). flash_mla needs int32 indices. if ( isinstance(metadata, DSV4Metadata) and forward_batch.out_cache_loc is not None @@ -1311,13 +2169,26 @@ class DeepseekV4AttnBackend( self.topk, self.speculative_num_steps, )[self.speculative_step_id] - metadata.core_attn_metadata.swa_out_cache_loc = ( - self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to( - torch.int32 + if self.token_to_kv_pool.request_window is None: + metadata.core_attn_metadata.swa_out_cache_loc = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + out_cache_loc + ).to(torch.int32) ) - ) - if self.is_dspark_draft and forward_batch.forward_mode.is_target_verify(): + # Refresh low-ratio source metadata from the live decode inputs. + if ( + metadata.core_metadata.low_ratios + and forward_batch.forward_mode.is_decode() + ): + metadata.low_ratio_req_indices = token_req_indices(forward_batch) + metadata.low_ratio_pos_i64 = forward_batch.positions.to(torch.int64) + + if ( + self.is_dspark_draft + and forward_batch.forward_mode.is_target_verify() + and self.token_to_kv_pool.request_window is None + ): block_size = int(forward_batch.spec_info.draft_token_num) seq_lens_casual = self._dspark_seq_lens_casual( seq_lens=forward_batch.seq_lens, block_size=block_size @@ -1440,6 +2311,15 @@ class DeepseekV4AttnBackend( req_pool_indices, seq_lens, ) + dspark_swa_buffers = None + captured_metadata = self.cuda_graph_metadata_of_bucket_and_bs[bucket].get( + bs + ) + if not in_capture and captured_metadata is not None: + # Reuse only storage: the draft graph rebuilds both tensors from + # live inputs before attention. copy_ onto itself is a no-op. + core = captured_metadata.core_attn_metadata + dspark_swa_buffers = (core.swa_page_indices, core.swa_topk_lengths) temp_metadata = self.init_forward_metadata_dspark_draft_block( max_seq_len=chosen_max_seq_len, req_pool_indices=req_pool_indices, @@ -1447,6 +2327,7 @@ class DeepseekV4AttnBackend( seq_lens_cpu=seq_lens_cpu, out_cache_loc=out_cache_loc_padded, block_size=block_size, + dspark_swa_buffers=dspark_swa_buffers, ) elif bucket == _GraphBucket.TARGET_VERIFY: verify_bs = _get_target_verify_bs(forward_batch) @@ -1537,8 +2418,20 @@ class DeepseekV4AttnBackend( self.online_c128_mtp.clear() return + self.encoder_replay = forward_batch.encoder_swa_replay self.forward_metadata = self._build_forward_metadata(forward_batch) self.init_forward_metadata_in_graph(forward_batch) + self.tail_forward_metadata = ( + self._build_late_layer_tail_metadata(forward_batch) + if self.enable_decoder_swa_bounded_replay + and forward_batch.forward_mode.is_extend_without_speculative() + else None + ) + + if self.token_to_kv_pool.request_window is not None: + self.token_to_kv_pool.request_window.activate( + self.forward_metadata.core_attn_metadata.request_window_layout + ) def prepare_prefill_shared_read_snapshot( self, forward_batch: ForwardBatch, *, num_qo_tokens: int @@ -1547,6 +2440,8 @@ class DeepseekV4AttnBackend( # first layer. DFLASH/DSPARK have no later prefill draft-extend reader; # CP shards the query layout that this global snapshot assumes. metadata = self.forward_metadata + if self.token_to_kv_pool.request_window is not None: + return if isinstance(metadata, DSV4Metadata): metadata.prefill_shared_reads_snapshotted = False snapshot_shared_prefill_reads = ( @@ -1559,9 +2454,14 @@ class DeepseekV4AttnBackend( return assert isinstance(metadata, DSV4Metadata) - use_sparse_prefill = not get_platform().is_sm120 and ( - num_qo_tokens > _LARGE_INDEXER_QUERY_THRESHOLD - or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get() + # The tail never takes the sparse path, so it carries no chunk cache. + use_sparse_prefill = ( + not get_platform().is_sm120 + and metadata.late_layer_tail is None + and ( + num_qo_tokens > _LARGE_INDEXER_QUERY_THRESHOLD + or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get() + ) ) if use_sparse_prefill: metadata.sparse_prefill_cache = self._build_sparse_prefill_chunk_cache( @@ -1580,6 +2480,9 @@ class DeepseekV4AttnBackend( ) -> SparsePrefillChunkCache: seq_lens_cpu = forward_batch.seq_lens_cpu assert seq_lens_cpu is not None + # The chunk cache gathers the W-1 positions before the chunk; under the + # tail those are late-layer window slots this prefill never wrote. + assert self.forward_metadata.late_layer_tail is None extend_seq_lens = forward_batch.extend_seq_lens extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu assert extend_seq_lens_cpu is not None @@ -1590,12 +2493,26 @@ class DeepseekV4AttnBackend( seq_lens_cpu_list, extend_seq_lens_cpu, strict=True ) ) - # The rows this forward runs are the extend, one per causal position. + if is_cp_active(forward_batch): + query_lens = torch.tensor( + interleave_rows_per_request( + _as_int_list(extend_seq_lens_cpu), + get_parallel().attn_cp_rank, + get_parallel().attn_cp_size, + ), + dtype=torch.int32, + device=extend_seq_lens.device, + ) + else: + query_lens = extend_seq_lens.to(torch.int32) + # padding rows are never combined query_pos = core_attn_metadata.seq_lens_casual[:num_qo_tokens] - 1 + if query_pos.shape[0] < num_qo_tokens: + query_pos = _pad_tensor_to_size(query_pos, num_qo_tokens, value=0) return SparsePrefillChunkCache.build( seq_lens=forward_batch.seq_lens.to(torch.int32), extend_seq_lens=extend_seq_lens.to(torch.int32), - query_lens=extend_seq_lens.to(torch.int32), + query_lens=query_lens, query_pos=query_pos, req_pool_indices=forward_batch.req_pool_indices.to(torch.int32), req_to_token=self.req_to_token, @@ -1731,8 +2648,40 @@ class DeepseekV4AttnBackend( max_seq_len_override=max_seq_len, use_prefill_cuda_graph=True, ) + if self.low_ratio_prefill_graph and forward_batch.forward_mode.is_extend(): + for ratio in self.low_ratios: + self._source_projection_buffers( + forward_batch.out_cache_loc.shape[0], ratio + ) return self.forward_metadata + def _source_projection_buffers(self, num_tokens: int, ratio: int) -> dict: + cfg = self.model_runner.model_config.hf_text_config + heads, dim = int(cfg.index_n_heads), int(cfg.index_head_dim) + sets = getattr(self, "_source_proj_bufs", None) or {} + have = sets.get(ratio) + if have is None or have[0]["q"].shape[0] < num_tokens: + latent = self.model_runner.model_config.head_dim + zeros = lambda *shape, dtype: torch.zeros( + *shape, dtype=dtype, device=self.device + ) + bufs = { + "q": zeros(num_tokens, heads, dim, dtype=torch.bfloat16), + "w": zeros(num_tokens, heads, dtype=torch.bfloat16), + "kv": zeros( + num_tokens, + latent, + dtype=torch.bfloat16 if ratio == 1 else torch.float32, + ), + } + if ratio == 2: + bufs["score"] = zeros(num_tokens, latent, dtype=torch.float32) + # Keep the previous allocations alive: captured graphs still hold them. + sets[ratio] = [bufs] + (have or []) + self._source_proj_bufs = sets + bufs = sets[ratio][0] + return {name: buf[:num_tokens] for name, buf in bufs.items()} + def prepare_forward_metadata_for_breakable_cuda_graph_replay( self, capture_metadata, @@ -1827,6 +2776,10 @@ class DeepseekV4AttnBackend( core.c4_flashmla_metadata = _create_flashmla_metadata() if core.has_c128: core.c128_flashmla_metadata = _create_flashmla_metadata() + if 1 in core.low_ratios: + core.c1_flashmla_metadata = _create_flashmla_metadata() + if 2 in core.low_ratios: + core.c2_flashmla_metadata = _create_flashmla_metadata() # PREP_IN_CUDA_GRAPH=True: warmup upgraded raw->full on the host; # restore raw so capture re-runs the upgrade inside the graph. @@ -1834,8 +2787,897 @@ class DeepseekV4AttnBackend( if current_raw is not None: self.forward_metadata = current_raw + # ---- DeepSeek V4.1 ratio 1/2 compressor and indexer, torch bring-up path ---- + + def forward_low_ratio_sources( + self, + *, + layer, + x, + q_lora, + positions, + forward_batch: ForwardBatch, + run_compressor: bool = True, + run_indexer: bool = True, + ) -> None: + """Runs on every ratio 1/2 layer before its attention.""" + if forward_batch.forward_mode.is_idle(): + return + if forward_batch.encoder_swa_replay: + run_compressor = False + if dsa_use_prefill_cp(forward_batch) and forward_batch.forward_mode.is_extend(): + self._forward_low_ratio_sources_cp( + layer=layer, + x=x, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + run_compressor=run_compressor, + run_indexer=run_indexer, + ) + return + meta = self.forward_metadata + hoisted_req = getattr(meta, "low_ratio_req_indices", None) + hoisted_pos = getattr(meta, "low_ratio_pos_i64", None) + if ( + hoisted_req is not None + and hoisted_pos is not None + and hoisted_pos.shape[0] == positions.shape[0] + ): + # Bucket-sized under the prefill graph; an eager break sees the + # live rows only and falls through. + req, pos = hoisted_req, hoisted_pos + else: + req = token_req_indices(forward_batch, num_tokens=positions.shape[0]) + # Every consumer takes int32 or int64 positions; keep the caller's. + pos = positions + if ( + forward_batch.forward_mode.is_extend() + and self._low_ratio_in_prefill_graph() + ): + bufs = self._source_projection_buffers(x.shape[0], layer.compress_ratio) + _bcg_low_ratio_source_projections(layer, x, q_lora, pos, bufs) + if run_compressor and layer.compressor is not None: + self._low_ratio_compress_torch( + layer, x, req, pos, projected=(bufs["kv"], bufs.get("score")) + ) + if run_indexer and layer.indexer is not None: + self._low_ratio_index_topk_prefill_graph( + layer, pos, bufs["q"], bufs["w"] + ) + return + if run_compressor and layer.compressor is not None: + self._low_ratio_compress(layer, x, req, pos, forward_batch) + if run_indexer and layer.indexer is not None: + self._low_ratio_index_topk(layer, x, q_lora, req, pos, forward_batch) + + def _forward_low_ratio_sources_cp( + self, *, layer, x, q_lora, positions, forward_batch, run_compressor, run_indexer + ) -> None: + # Every rank writes the whole prompt's compressed state, scoring its own rows. + cp_meta = forward_batch.attn_cp_metadata + total = int(cp_meta.total_seq_lens) + tail = self.forward_metadata.late_layer_tail + if tail is not None: + q_lens_cpu = tail.local_lens_cpu + req_global, pos_global = tail.req_global, tail.pos_global + else: + q_lens_cpu = interleave_rows_per_request( + _as_int_list(forward_batch.extend_seq_lens_cpu), + get_parallel().attn_cp_rank, + get_parallel().attn_cp_size, + ) + req_global = token_req_indices(forward_batch, num_tokens=total) + pos_global = forward_batch.positions[:total].to(torch.int64) + num_local = sum(q_lens_cpu) + if run_compressor and layer.compressor is not None: + x_global = cp_materialize_global_token_order( + x.contiguous(), forward_batch, torch.cuda.current_stream() + )[:total] + self._low_ratio_compress_torch(layer, x_global, req_global, pos_global) + if run_indexer and layer.indexer is not None: + self._low_ratio_index_topk_dense( + layer, + x[:num_local], + q_lora[:num_local], + positions[:num_local].to(torch.int64), + forward_batch, + torch.tensor(q_lens_cpu, dtype=torch.int32, device=x.device), + q_lens_cpu, + ) + + def _low_ratio_compress(self, layer, x, req, pos, forward_batch) -> None: + if forward_batch.forward_mode.is_decode(): + self._low_ratio_compress_decode(layer, x, req, pos) + elif ( + forward_batch.forward_mode.is_target_verify() + and not self.is_dspark_draft + and layer.compress_ratio in (1, 2) + and layer.compressor.use_fused_compress + and read_ragged_verify_mode() is not RaggedVerifyMode.COMPACT + and self.speculative_num_draft_tokens is not None + and self.speculative_num_draft_tokens > 1 + and x.shape[0] + == forward_batch.batch_size * self.speculative_num_draft_tokens + ): + # Static verify is request-major with consecutive positions; compact + # verify has variable block lengths and keeps the general path. + self._low_ratio_compress_fused( + layer, x, req, pos, draft_len=self.speculative_num_draft_tokens + ) + else: + self._low_ratio_compress_torch( + layer, + x, + req, + pos, + fuse_index_store=( + forward_batch.forward_mode.is_target_verify() + and layer.compressor.use_fused_compress + ), + ) + + def _low_ratio_in_prefill_graph(self) -> bool: + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, + ) + + return self.low_ratio_prefill_graph and is_in_breakable_cuda_graph() + + def _low_ratio_compress_decode(self, layer, x, req, pos) -> None: + # Projection layout and fused-write support are fixed together at load time. + if layer.compressor.use_fused_compress: + self._low_ratio_compress_fused(layer, x, req, pos) + return + if layer.compress_ratio == 1: + core = self.forward_metadata.core_metadata + kv, _ = layer.compressor.project(x) + slots = torch.where( + core.c1_out_loc >= 0, core.c1_out_loc, torch.zeros_like(core.c1_out_loc) + ) + self._low_ratio_write_group( + layer, + kv, + slots, + pos, + fuse_index_store=( + x.is_cuda + and torch.version.cuda is not None + and _is_sm100_or_newer() + ), + ) + return + if not (x.is_cuda and torch.version.cuda): + self._low_ratio_compress_torch(layer, x, req, pos) + return + + from sglang.kernels.ops.attention.dsv4.c2_decode_pool import c2_decode_pool + + core = self.forward_metadata.core_metadata + state = self.token_to_kv_pool.get_attention_compress_states(layer.layer_id) + kv, score = layer.compressor.project(x) + pooled, group_pos, slots = c2_decode_pool( + kv, + score, + pos, + core.raw_out_loc, + core.c2_out_loc, + req, + state.kv_score_buffer.kv, + state.kv_score_buffer.score, + state.kv_score_buffer.shape[0] - 1, + ring_size=state.ring_size, + ) + self._low_ratio_write_group( + layer, + pooled, + slots, + group_pos, + fuse_index_store=_is_sm100_or_newer(), + ) + + def _low_ratio_compress_fused(self, layer, x, req, pos, *, draft_len=1) -> None: + from sglang.kernels.ops.attention.dsv4.fp4_indexer_rope import ( + index_k_norm_rope_pack_store, + ) + from sglang.kernels.ops.attention.dsv4.low_ratio_compress import ( + c1_decode_norm_rope_store, + c2_decode_norm_rope_store, + ) + + pool = self.token_to_kv_pool + core = self.forward_metadata.core_metadata + compressor = layer.compressor + layer_id = layer.layer_id + # Contiguous complex64 freqs_cis gives a real/imag-interleaved view without copying. + freqs_cis = torch.view_as_real(layer.freqs_cis).flatten(-2) + kv_cache = pool.get_extra_key_buffer(layer_id) + page_size = pool.get_extra_key_page_size(layer_id) + # The pool's page format: V4, or the V4.1 fp8 / fp4 layouts. + kv_layout = pool.get_extra_key_layout(layer_id) + assert kv_cache is not None + + if layer.compress_ratio == 1: + # At ratio 1, c1_out_loc equals the int64 raw_out_loc supplied by the scheduler. + latent = c1_decode_norm_rope_store( + compressor.wkv(x), + compressor.norm.weight.data, + pos, + core.raw_out_loc, + compressor.norm.eps, + freqs_cis, + kv_cache, + page_size=page_size, + layout=kv_layout, + ) + out_loc = core.c1_out_loc + else: + # CompressStatePool stores each request's pending pairs in a position ring. + # KVAndScore rows use | kv | score |, addressed as req * ring_size + pos % ring_size. + state = pool.get_attention_compress_states(layer_id) + latent = c2_decode_norm_rope_store( + compressor.project_fused(x), + state.kv_score_buffer.kv_score, + compressor.norm.weight.data, + pos, + req, + core.raw_out_loc, + compressor.norm.eps, + freqs_cis, + kv_cache, + page_size=page_size, + ring_size=state.ring_size, + draft_len=draft_len, + layout=kv_layout, + ) + out_loc = core.c2_out_loc + + indexer = layer.indexer + if indexer is not None and indexer.owns_k: + # out_loc is -1 for an incomplete group and 0 for padding; + # the kernel suppresses both stores. + assert out_loc is not None + index_k_norm_rope_pack_store( + indexer.forward_wk(latent), + indexer.k_norm.weight.data, + indexer.k_norm.eps, + freqs_cis, + pos, + out_loc, + pool.get_index_k_with_scale_buffer(layer_id), + ratio=layer.compress_ratio, + ) + + def _low_ratio_compress_torch( + self, layer, x, req, pos, projected=None, *, fuse_index_store=False + ) -> None: + core = self.forward_metadata.core_metadata + num_tokens = pos.shape[0] + kv, score = projected if projected is not None else layer.compressor.project(x) + if not num_tokens: + return + if layer.compress_ratio == 1: + self._low_ratio_write_group( + layer, + kv, + core.c1_out_loc[:num_tokens], + pos, + fuse_index_store=fuse_index_store, + ) + return + + partner_kv, partner_score = self._low_ratio_pair_partners( + layer_id=layer.layer_id, + kv=kv, + score=score, + req=req, + pos=pos, + pad=core.raw_out_loc[:num_tokens] == 0, + ) + pooled = layer.compressor.pool_pairs( + torch.stack([partner_kv, kv], dim=1), + torch.stack([partner_score, score], dim=1), + ) + group_pos = torch.where(pos % 2 == 1, pos - 1, pos) + out_loc = core.c2_out_loc[:num_tokens] + slots = torch.where(out_loc >= 0, out_loc, torch.zeros_like(out_loc)) + self._low_ratio_write_group( + layer, pooled, slots, group_pos, fuse_index_store=fuse_index_store + ) + + def _low_ratio_pair_partners( + self, *, layer_id, kv, score, req, pos, pad + ) -> Tuple[torch.Tensor, torch.Tensor]: + state = self.token_to_kv_pool.get_attention_compress_states(layer_id) + ring = state.ring_size + num_tokens = pos.shape[0] + if not num_tokens: + return kv, score + + read_pos = (pos - 1).masked_fill(pad, -1) + carried = state.get_state_by_state_loc( + state.translate_from_req_position_to_state_loc(req, read_pos) + ) + in_batch = torch.zeros_like(pad) + in_batch[1:] = ( + (req[1:] == req[:-1]) & (pos[1:] == pos[:-1] + 1) & ~pad[1:] & ~pad[:-1] + ) + partner_kv = torch.where(in_batch[:, None], torch.roll(kv, 1, 0), carried.kv) + partner_score = torch.where( + in_batch[:, None], torch.roll(score, 1, 0), carried.score + ) + + # All reads above precede every write, and keeping only a request's last + # ring_size rows leaves each write a distinct live slot, even when a + # prefill chunk is longer than the ring. + keep = ~pad + if num_tokens > ring: + keep[:-ring] &= (req[:-ring] != req[ring:]) | pad[ring:] + write_pos = pos.masked_fill(~keep, -1) + state.set_state_by_state_loc( + state.translate_from_req_position_to_state_loc(req, write_pos), + KVAndScore.from_kv_score(kv=kv, score=score), + ) + return partner_kv, partner_score + + def _low_ratio_write_group( + self, + layer, + pooled, + slots, + group_pos, + *, + fuse_index_store=False, + ) -> None: + pool = self.token_to_kv_pool + latent = layer.compressor.finish(pooled) + freqs = layer.freqs_cis[group_pos] + # Index keys come from the pre-RoPE latent, so publish them first. Stored + # as fp4 (per-32 ue8m0, no hadamard), matching the reference indexer. + if layer.indexer is not None and layer.indexer.owns_k: + if ( + fuse_index_store + and latent.is_cuda + and torch.version.cuda is not None + and latent.dtype == torch.bfloat16 + and layer.indexer.index_head_dim == 128 + ): + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + index_k_rope_pack, + ) + + indexer = layer.indexer + k = indexer.k_norm(indexer.forward_wk(latent)) + index_k_rope_pack( + k, + freqs, + indexer.rope_head_dim, + cache=pool.get_index_k_with_scale_buffer(layer.layer_id), + loc=slots, + ) + else: + pool.set_index_k_fp4( + layer_id=layer.layer_id, + loc=slots, + cache_k=layer.indexer.index_keys(latent, freqs), + ) + if pool.get_extra_key_layout(layer.layer_id) is KVLayout.V41_FP4: + # The fp4 cache stores e2m1 codes: the kernel rotates the tail and + # rounds once, with no fake quantization in between. + pool.set_extra_key_buffer_fused( + layer_id=layer.layer_id, loc=slots, cache_k=latent, freqs_cis=freqs + ) + return + # The fp8 FlashMLA caches requantize the FP4/E4M3 latent into their layout. + latent = _rope_fq4(latent, freqs, layer.rope_head_dim, compressed_kv=True) + pool.set_extra_key_buffer_fused( + layer_id=layer.layer_id, loc=slots, cache_k=latent + ) + + def _low_ratio_index_topk(self, layer, x, q_lora, req, pos, forward_batch) -> None: + is_decode_or_verify = ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ) + if is_decode_or_verify: + if _is_sm100_or_newer(): + # DeepGEMM pairs verify rows by request id; decode has one row each. + req_ids = None if forward_batch.forward_mode.is_decode() else req + self._low_ratio_index_topk_decode(layer, x, q_lora, pos, req_ids) + else: + self._low_ratio_index_topk_sm90_decode(layer, x, q_lora, req, pos) + elif ( + self._use_dense_fp4_prefill_indexer(forward_batch) and _is_sm100_or_newer() + ): + self._low_ratio_index_topk_extend(layer, x, q_lora, pos, forward_batch) + else: + self._low_ratio_index_topk_torch(layer, x, q_lora, req, pos) + + @staticmethod + def _use_dense_fp4_prefill_indexer(forward_batch) -> bool: + return ( + not envs.SGLANG_DSV41_TORCH_PREFILL_INDEXER.get() + and _has_dense_fp4_indexer() + and forward_batch.forward_mode.is_extend() + and forward_batch.seq_lens_cpu is not None + and forward_batch.extend_seq_lens_cpu is not None + ) + + def _low_ratio_index_topk_extend( + self, layer, x, q_lora, pos, forward_batch + ) -> None: + tail = self.forward_metadata.late_layer_tail + if tail is not None: + q_lens, q_lens_cpu = tail.extend_seq_lens, tail.extend_seq_lens_cpu + else: + q_lens = forward_batch.extend_seq_lens + q_lens_cpu = _as_int_list(forward_batch.extend_seq_lens_cpu) + assert q_lens_cpu is not None + self._low_ratio_index_topk_dense( + layer, x, q_lora, pos, forward_batch, q_lens, q_lens_cpu + ) + + def _low_ratio_index_topk_dense( + self, layer, x, q_lora, pos, forward_batch, q_lens, q_lens_cpu + ) -> None: + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + quantize_fp4_indexer_tensor, + ) + + pool = self.token_to_kv_pool + core = self.forward_metadata.core_metadata + ratio = layer.compress_ratio + indexer = layer.indexer + page_indices = core.sparse_page_indices(ratio) + raw_indices = core.sparse_raw_indices(ratio) + page_indices.fill_(-1) + if raw_indices is not None: + raw_indices.fill_(-1) + + seq_lens_cpu = _as_int_list(forward_batch.seq_lens_cpu) + assert seq_lens_cpu is not None + device = pos.device + # Visible compressed positions per request at its newest token; the + # per-token count (pos + 1) // ratio bounds each row below. + lc_per_req = [s // ratio for s in seq_lens_cpu] + req_pool_indices = forward_batch.req_pool_indices.to(torch.int64) + slot_chunks, starts, start = [], [], 0 + for r, lc in enumerate(lc_per_req): + starts.append(start) + if lc == 0: + continue + j = torch.arange(lc, device=device) + slot_chunks.append( + self.req_to_token[req_pool_indices[r], j * ratio].to(torch.int64) + // ratio + ) + start += lc + empty_mask = torch.zeros(0, 0, dtype=torch.bool, device=device) + num_tokens = pos.shape[0] + # TODO(candidate): move this to candidate indexer + if not slot_chunks or num_tokens == 0: + if indexer.is_candidate_source: + self.forward_metadata.candidate_metadata = CandidateMasks( + request_masks=[empty_mask for _ in lc_per_req] + ) + return + k_slots = torch.cat(slot_chunks) + k_fp4, k_sf = pool.get_low_ratio_index_k_fp4(layer.layer_id, k_slots) + + q = indexer.queries(q_lora, layer.freqs_cis[pos]) # [T, H, 128] fp4 grid + num_heads = q.shape[1] + q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) + q_fp4 = q_fp4.view(num_tokens, num_heads, 64) + q_sf = q_sf.view(num_tokens, num_heads) + weights = indexer.head_weights(x).float() + compress_lens = ((pos + 1) // ratio).to(torch.int32) + ks = torch.repeat_interleave( + torch.tensor(starts, dtype=torch.int32, device=device), + q_lens.to(torch.int64), + output_size=num_tokens, + ) + logits = _dense_fp4_mqa_logits( + (q_fp4, q_sf), + (k_fp4, k_sf), + weights, + ks, + ks + compress_lens, + # the fused top-k reads score rows through 16-byte vectors + ceil_align(max(lc_per_req), 4), + ) + if indexer.is_candidate_source or indexer.uses_candidates: + self._publish_or_consume_candidates( + indexer, logits, compress_lens, lc_per_req, q_lens_cpu, empty_mask + ) + topk = indexer.index_topk + selected = torch.empty((num_tokens, topk), dtype=torch.int32, device=device) + topk_transform_ragged_v2( + logits, compress_lens, out_offsets=ks, out_indices=selected + ) + if indexer.uses_candidates and not indexer.is_candidate_source: + selected = mask_topk_scores(logits, selected, ks) + # ascending positions, padding last: the layout the consumers expect + unselected = torch.iinfo(torch.int32).max + selected = selected.masked_fill(selected < 0, unselected).sort(dim=-1).values + chosen = selected != unselected + page_indices[:num_tokens, :topk] = torch.where( + chosen, k_slots[selected.clamp_max(k_slots.shape[0] - 1)], -1 + ).to(torch.int32) + if raw_indices is not None: + raw_indices[:num_tokens, :topk] = torch.where( + chosen, selected - ks[:, None], -1 + ) + + # TODO(candidate): dense-prefill level one / level two inline with masks; move + # into the candidate indexer as publish_prefill / select_prefill. + def _publish_or_consume_candidates( + self, indexer, logits, compress_lens, lc_per_req, q_lens_cpu, empty_mask + ) -> None: + publish = [] if indexer.is_candidate_source else None + consume = ( + None + if publish is not None + else published_masks(self.forward_metadata.candidate_metadata) + ) + j = torch.arange(logits.shape[1], device=logits.device) + tok_start = 0 + for b, (lc, t_len) in enumerate(zip(lc_per_req, q_lens_cpu)): + rows = slice(tok_start, tok_start + t_len) + tok_start += t_len + if lc == 0 or t_len == 0: + if publish is not None: + publish.append(empty_mask) + continue + scores = logits[rows, :lc] + if publish is None: + scores.masked_fill_(~consume.request_masks[b], -torch.inf) + continue + lens = compress_lens[rows, None] + # the block selection tells unreachable positions apart by -inf + scores.masked_fill_(j[None, :lc] >= lens, -torch.inf) + # the block selection pads and pools a copy of its rows; bound that copy + step = max(1, _TORCH_INDEXER_SCORE_BUDGET_BYTES // (lc * 4)) + masks = [ + select_candidate_blocks( + scores[start : start + step], + lens[start : start + step], + topk_blocks=indexer.candidate_topk_blocks, + block_size=indexer.candidate_block_size, + ) + for start in range(0, t_len, step) + ] + publish.append(masks[0] if len(masks) == 1 else torch.cat(masks)) + if publish is not None: + self.forward_metadata.candidate_metadata = CandidateMasks( + request_masks=publish + ) + + def _low_ratio_index_topk_prefill_graph(self, layer, pos, q, w) -> None: + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + quantize_fp4_indexer_tensor, + ) + + pool = self.token_to_kv_pool + core = self.forward_metadata.core_metadata + ratio = layer.compress_ratio + indexer = layer.indexer + metadata = ( + self.forward_metadata.c1_indexer_metadata + if ratio == 1 + else self.forward_metadata.c2_indexer_metadata + ) + assert metadata is not None, f"no prefill graph indexer metadata for {ratio = }" + assert indexer.n_local_heads == indexer.n_heads + width = metadata.max_compressed_seq_len + if indexer.uses_candidates or indexer.is_candidate_source: + # Every reachable block is a candidate inside the window, so the + # two-level selection collapses to the plain top-k below. + assert ( + width <= indexer.candidate_topk_blocks * indexer.candidate_block_size + ), f"prefill graph indexer width {width} exceeds the candidate window" + + num_tokens, num_heads = q.shape[0], q.shape[1] + q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) + q_fp4 = q_fp4.view(num_tokens, 1, num_heads, 64) + q_sf = q_sf.view(num_tokens, 1, num_heads) + weights = w.float() + + k_cache = pool.get_index_k_with_scale_buffer(layer.layer_id) + assert k_cache.dim() == 2 + page_size = metadata.compressed_page_size + k_cache = k_cache.view(k_cache.shape[0], page_size, 1, 68) + + lens = metadata.compressed_seq_lens + page_table = metadata.page_table + page_indices = core.sparse_page_indices(ratio) + raw_indices = core.sparse_raw_indices(ratio) + topk = min(indexer.index_topk, width) + columns = torch.arange(width, device=lens.device) + for rows, plan in metadata.row_chunks(): + logits = deep_gemm_fp4_paged_mqa_logits( + (q_fp4[rows], q_sf[rows]), + k_cache, + weights[rows], + lens[rows], + page_table[rows], + plan, + width, + ) + lens_c = lens[rows].unsqueeze(-1) + # Columns past a row's length hold garbage. + s = logits.masked_fill(columns[None, :] >= lens_c, -torch.inf) + idx = s.topk(topk, dim=-1, sorted=False).indices.sort(dim=-1).values + reach = idx < lens_c + slots = page_table[rows].gather(-1, idx // page_size) * page_size + ( + idx % page_size + ) + page_indices[rows, :topk] = torch.where(reach, slots, -1).to(torch.int32) + if raw_indices is not None: + raw_indices[rows, :topk] = torch.where(reach, idx, -1).to(torch.int32) + + def _low_ratio_index_topk_decode(self, layer, x, q_lora, pos, req=None) -> None: + from sglang.srt.model_executor.runner_utils.capture_mode import ( + skip_low_ratio_indexer, + ) + + if skip_low_ratio_indexer(layer.compress_ratio): + # The compressor still writes index K for later, longer contexts. + return + + from sglang.kernels.ops.attention.dsv4.fp4_indexer import ( + quantize_fp4_indexer_tensor, + ) + + pool = self.token_to_kv_pool + core = self.forward_metadata.core_metadata + ratio = layer.compress_ratio + indexer = layer.indexer + metadata = ( + self.forward_metadata.c1_indexer_metadata + if ratio == 1 + else self.forward_metadata.c2_indexer_metadata + ) + assert metadata is not None, f"no decode indexer metadata for {ratio = }" + + # fp4 query as (payload, scale), kernel layout [bs, 1, n_heads, dim]. The + # kernel sums head scores locally, so the indexer heads must be replicated. + assert indexer.n_local_heads == indexer.n_heads + if ( + x.is_cuda + and torch.version.cuda is not None + and x.dtype == torch.bfloat16 + and indexer.index_head_dim == 128 + ): + from sglang.kernels.ops.attention.dsv4.fp4_indexer_rope import ( + index_q_rope_pack_weights, + ) + + q, _ = indexer.wq_b(q_lora) + q = q.view(q.shape[0], indexer.n_local_heads, indexer.index_head_dim) + # The fused pack also computes head_weights(x).float(), same rounding. + q_fp4, q_sf, weights = index_q_rope_pack_weights( + q, + torch.view_as_real(layer.freqs_cis).flatten(-2), + pos, + indexer.head_weights_raw(x), # [bs, n_local] bf16, n32k5120 + indexer.head_weight_scale, + ) + else: + q = indexer.queries(q_lora, layer.freqs_cis[pos]) + q_fp4, q_sf = quantize_fp4_indexer_tensor(q.flatten(0, 1), rne=True) + weights = indexer.head_weights(x).float() # [bs, n_local] + bs = q.shape[0] + q_fp4 = q_fp4.view(bs, 1, indexer.n_local_heads, 64) + q_sf = q_sf.view(bs, 1, indexer.n_local_heads) + + k_cache = pool.get_index_k_with_scale_buffer(layer.layer_id) + assert k_cache.dim() == 2 + # Index pool page (64 slots); metadata.page_table is expanded to match. + page_size = metadata.compressed_page_size + k_cache = k_cache.view( + k_cache.shape[0], page_size, 1, 68 + ) # fp4: 64 payload + 4 scale + + page_indices = core.sparse_page_indices(ratio) + raw_indices = core.sparse_raw_indices(ratio) + inputs = IndexerInputs( + q_fp4, + q_sf, + k_cache, + weights, + metadata, + request_ids=req, # one per query row; verify rows of a request share one + ) + candidate_layer = not _every_request_fits() + # use special selection for candidate layers + if indexer.uses_candidates and candidate_layer: + return self.candidate_indexer.select_decode( + self.forward_metadata.candidate_metadata, + inputs, + page_indices, + raw_indices, + ) + if indexer.is_candidate_source and candidate_layer: + self.forward_metadata.candidate_metadata = ( + self.candidate_indexer.publish_decode(inputs, page_indices, raw_indices) + ) + return + logits = deep_gemm_fp4_paged_mqa_logits( + (q_fp4, q_sf), + k_cache, + weights, + metadata.compressed_seq_lens, + metadata.page_table, + metadata.deep_gemm_metadata, + metadata.max_compressed_seq_len, + ) + # TODO(dark): add bf16 topk + topk_transform_paged_from_metadata(logits, metadata, page_indices, raw_indices) + + # TODO(candidate): Hopper decode still publishes / consumes masks inline (torch + # top-k); move into the candidate indexer with the prefill paths. + def _low_ratio_index_topk_sm90_decode(self, layer, x, q_lora, req, pos) -> None: + pool = self.token_to_kv_pool + core = self.forward_metadata.core_metadata + ratio = layer.compress_ratio + indexer = layer.indexer + page_indices = core.sparse_page_indices(ratio) + raw_indices = core.sparse_raw_indices(ratio) + page_indices.fill_(-1) + if raw_indices is not None: + raw_indices.fill_(-1) + bs = req.shape[0] + assert pos.shape[0] == bs, ( + f"decode expects one token per request, {pos.shape=} {bs=}" + ) + if bs == 0: + return + lens = (pos + 1) // ratio + metadata = ( + self.forward_metadata.c1_indexer_metadata + if ratio == 1 + else self.forward_metadata.c2_indexer_metadata + ) + assert metadata is not None + # V4 reserves the replay bound in metadata; visibility stays on device. + # A capture-time length read would both synchronize and truncate replay. + lmax = min(metadata.max_compressed_seq_len, self.req_to_token.shape[1] // ratio) + if lmax == 0: + return + q = indexer.queries(q_lora, layer.freqs_cis[pos]) + weights = indexer.head_weights(x) + j = torch.arange(lmax, device=pos.device) + valid = j[None, :] < lens[:, None] + slots = ( + self.req_to_token[req[:, None], (j * ratio)[None, :]].to(torch.int64) + // ratio + ) + slots = slots.masked_fill(~valid, 0) + table = pool.get_index_k_with_scale_buffer(layer.layer_id) + s = fp4_index_logits_decode( + q, weights, slots, lens, table, table.shape[1] // 68 + ) + if indexer.is_candidate_source: + mask = select_candidate_blocks( + s, + lens[:, None], + topk_blocks=indexer.candidate_topk_blocks, + block_size=indexer.candidate_block_size, + ) + self.forward_metadata.candidate_metadata = CandidateMasks(mask=mask) + elif indexer.uses_candidates: + # Published this step by the candidate-source layer's decode pass above. + consume = published_masks(self.forward_metadata.candidate_metadata).mask + assert torch.is_tensor(consume) and consume.shape[0] == bs, ( + "candidate mask missing for decode" + ) + s = s.masked_fill(~consume[:, :lmax], -torch.inf) + k = min(indexer.index_topk, lmax) + idx = s.topk(k, dim=-1, sorted=False).indices + if indexer.uses_candidates and not indexer.is_candidate_source: + idx = mask_topk_scores(s, idx) + idx = idx.masked_fill(idx < 0, lmax) + idx = idx.sort(dim=-1).values + reach = idx < lens[:, None] + page_indices[:bs, :k] = torch.where( + reach, slots.gather(1, idx.clamp_max(lmax - 1)), -1 + ).to(torch.int32) + if raw_indices is not None: + raw_indices[:bs, :k] = torch.where(reach, idx, -1).to(torch.int32) + + # TODO(candidate): torch prefill still publishes / consumes masks inline; same + # move as above. + def _low_ratio_index_topk_torch(self, layer, x, q_lora, req, pos) -> None: + pool = self.token_to_kv_pool + core = self.forward_metadata.core_metadata + ratio = layer.compress_ratio + indexer = layer.indexer + # Attention scans sparse_topk_lengths slots and skips -1 entries. + page_indices = core.sparse_page_indices(ratio) + raw_indices = core.sparse_raw_indices(ratio) + page_indices.fill_(-1) + if raw_indices is not None: + raw_indices.fill_(-1) + q = indexer.queries(q_lora, layer.freqs_cis[pos]) + weights = indexer.head_weights(x) + # A compressed position is visible once the query has passed its last token. + compress_lens = (pos + 1) // ratio + topk = indexer.index_topk + publish = [] if indexer.is_candidate_source else None + consume = ( + published_masks(self.forward_metadata.candidate_metadata).request_masks + if indexer.uses_candidates + else None + ) + for b, r in enumerate(torch.unique_consecutive(req).tolist()): + tok = (req == r).nonzero().squeeze(1) + lens = compress_lens[tok] + lc = int(lens.max().item()) + if lc == 0: + # Consumers address masks by request position, including empty requests. + if publish is not None: + publish.append( + torch.zeros(0, 0, dtype=torch.bool, device=pos.device) + ) + continue + j = torch.arange(lc, device=pos.device) + slots_j = self.req_to_token[r, j * ratio].to(torch.int64) // ratio + # Dequantize only this request's visible K rows; the table is pool-sized. + index_k = pool.get_low_ratio_index_k_dequant(layer.layer_id, slots_j) + k = min(topk, lc) + # Every step below is per query row; chunk rows so the [rows, heads, lc] + # bf16 scores stay under the budget (16 GiB at once for a 16k-token prompt). + rows_per_chunk = max( + 1, + _TORCH_INDEXER_SCORE_BUDGET_BYTES // (q.shape[1] * lc * 2), + ) + masks = [] if publish is not None else None + for start in range(0, tok.numel(), rows_per_chunk): + rows = slice(start, start + rows_per_chunk) + tok_c, lens_c = tok[rows], lens[rows] + s = indexer.scores(q[tok_c], index_k, weights[tok_c]) + s = s.masked_fill(j[None, :] >= lens_c[:, None], -torch.inf) + if masks is not None: + masks.append( + select_candidate_blocks( + s, + lens_c[:, None], + topk_blocks=indexer.candidate_topk_blocks, + block_size=indexer.candidate_block_size, + ) + ) + elif consume is not None: + s = s.masked_fill(~consume[b][rows], -torch.inf) + idx = s.topk(k, dim=-1, sorted=False).indices + if consume is not None and masks is None: + idx = mask_topk_scores(s, idx) + idx = idx.masked_fill(idx < 0, lc) + idx = idx.sort(dim=-1).values + reach = idx < lens_c[:, None] + page_indices[tok_c, :k] = torch.where( + reach, slots_j[idx.clamp_max(lc - 1)], -1 + ).to(torch.int32) + if raw_indices is not None: + raw_indices[tok_c, :k] = torch.where(reach, idx, -1).to(torch.int32) + if masks is not None: + publish.append(torch.cat(masks) if len(masks) > 1 else masks[0]) + if publish is not None: + self.forward_metadata.candidate_metadata = CandidateMasks( + request_masks=publish + ) + def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: - # Idle metadata may be stale; zero-padded locations target the dummy slot. + """Idle always re-translates at store time: its metadata may be stale, and + translating the zero-padded out_cache_loc writes to the dummy slot.""" + metadata = self.forward_metadata + if self.token_to_kv_pool.request_window is not None: + layout = metadata.core_attn_metadata.request_window_layout + self.token_to_kv_pool.request_window.activate(layout) + return layout.write_loc + if isinstance(metadata, DSV4Metadata) and metadata.late_layer_tail is not None: + # The tail's q rows are a subset of the extend, so the full + # out_cache_loc below would be the wrong length; the tail owns its own. + return metadata.late_layer_tail.swa_out_cache_loc out_cache_loc = forward_batch.out_cache_loc core = getattr(self.forward_metadata, "core_attn_metadata", None) cached = core.swa_out_cache_loc if core is not None else None @@ -1859,14 +3701,25 @@ class DeepseekV4AttnBackend( cache_k=swa_k, ) - def forward( + def forward(self, q, k, v, layer, forward_batch, *args, **kwargs): + result = self._forward_attention(q, k, v, layer, forward_batch, *args, **kwargs) + window = self.token_to_kv_pool.request_window + if ( + window is not None + and not self.is_dspark_draft + and not forward_batch.forward_mode.is_idle() + ): + window.commit(self.token_to_kv_pool._swa_local_layer_id(layer.layer_id)) + return result + + def _forward_attention( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer: RadixAttention, forward_batch: ForwardBatch, - compress_ratio: Literal[0, 4, 128], + compress_ratio: Literal[0, 1, 2, 4, 128], save_kv_cache: bool = True, attn_sink: Optional[torch.Tensor] = None, **_, @@ -1898,7 +3751,8 @@ class DeepseekV4AttnBackend( swa_page_size = token_to_kv_pool.swa_page_size assert swa_k_cache.ndim == 2 - # The kernel detects each cache's format from the last dim of this view. + # The kernel detects each cache's format from the last dim of this + # view: 584 (V4), 528 (V4.1 fp8) or 288 (V4.1 fp4, extra cache only). k_cache_total_dim = token_to_kv_pool.get_swa_key_bytes_per_token() swa_k_cache = swa_k_cache[:, : swa_page_size * k_cache_total_dim].view( swa_k_cache.shape[0], swa_page_size, 1, k_cache_total_dim @@ -1965,10 +3819,13 @@ class DeepseekV4AttnBackend( f"{extra_indices.shape=}'s last dimension is not aligned to 64" ) - # sparse_prefill_fwd does not support SM120. + # sparse_prefill_fwd does not support SM120. The tail stays dense: its + # window floor lives in swa_page_indices, which the chunk cache ignores. if ( forward_batch.forward_mode.is_extend_without_speculative() and not get_platform().is_sm120 + and self.forward_metadata.late_layer_tail is None + and token_to_kv_pool.request_window is None and ( q.shape[0] > _LARGE_INDEXER_QUERY_THRESHOLD or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get() @@ -1994,6 +3851,38 @@ class DeepseekV4AttnBackend( attn_sink=attn_sink, ) + if ( + self.is_dsv41 + and get_platform().is_sm100 + and can_use_swapab_attention( + q, + swa_k_cache, + extra_k_cache, + layer.tp_q_head_num, + self.head_dim_v, + self.softmax_scale, + ) + and ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_draft_extend_v2() + ) + ): + from sglang.kernels.ops.attention.dsv4.decode_attention_sm100 import ( + swapab_attention, + ) + + return swapab_attention( + q[..., :16, :], + swa_k_cache, + swa_page_indices, + swa_topk_lengths, + attn_sink, + extra_k_cache, + extra_indices, + extra_topk_lengths, + ) + if get_platform().is_sm120: from sglang.kernels.ops.attention.flash_mla_sm120 import ( SM120_DECODE_MAX_TOKENS, @@ -2029,6 +3918,15 @@ class DeepseekV4AttnBackend( else: from sgl_kernel.flash_mla import flash_mla_with_kvcache + if self.is_dsv41: + _maybe_precompute_flashmla_sched_meta( + flashmla_metadata, + q=q, + indices=swa_page_indices, + topk_length=swa_topk_lengths, + extra_indices=extra_indices, + extra_topk_length=extra_topk_lengths, + ) o = flash_mla_with_kvcache( q=q, k_cache=swa_k_cache, @@ -2055,7 +3953,7 @@ class DeepseekV4AttnBackend( self, q: torch.Tensor, layer_id: int, - compress_ratio: Literal[0, 4, 128], + compress_ratio: Literal[0, 1, 2, 4, 128], forward_batch: ForwardBatch, token_to_kv_pool: DeepSeekV4TokenToKVPool, core_attn_metadata: DSV4AttnMetadata, @@ -2192,7 +4090,7 @@ class DeepseekV4AttnBackend( self, q: torch.Tensor, layer_id: int, - compress_ratio: Literal[0, 4, 128], + compress_ratio: Literal[0, 1, 2, 4, 128], forward_batch: ForwardBatch, token_to_kv_pool: DeepSeekV4TokenToKVPool, core_attn_metadata: DSV4AttnMetadata, @@ -2387,10 +4285,33 @@ class DeepseekV4AttnBackend( is_prefill: bool = False, dspark_block_size: Optional[int] = None, num_tokens: Optional[int] = None, + swa_replay_start: Optional[torch.Tensor] = None, + num_groups: Optional[int] = None, + dspark_swa_buffers: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> DSV4AttnMetadata: assert self.swa_page_size == SWA_WINDOW - prep = BuildPageTablePositions.execute( + small_metadata = ( + not is_prefill + and seq_lens_casual.is_cuda + and 0 < seq_lens_casual.numel() <= 8 + and out_loc.numel() == seq_lens_casual.numel() + and self.low_ratios == (1, 2) + and set(self.present_ratios) == {1, 2} + and get_parallel().attn_cp_size == 1 + and self.token_to_kv_pool.request_window is None + ) + build_pages = BuildPageTablePositions.execute + if small_metadata: + from sglang.kernels.ops.attention.dsv4.metadata_kernel import ( + build_low_ratio_metadata, + ) + from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import ( + build_page_table_positions_small, + ) + + build_pages = build_page_table_positions_small + prep = build_pages( req_to_token=req_to_token, req_pool_indices_repeated=req_pool_indices_repeated, seq_lens_casual=seq_lens_casual, @@ -2401,7 +4322,30 @@ class DeepseekV4AttnBackend( seq_lens_casual = prep.seq_lens_casual raw_positions = prep.positions_casual - if dspark_block_size is not None: + request_layout = None + if self.token_to_kv_pool.request_window is not None: + from sglang.srt.mem_cache.dsv41_request_window import window_layout + + if self.encoder_replay: + starts = torch.ones_like(raw_positions, dtype=torch.bool) + starts[1:] = ( + req_pool_indices_repeated[1:] != req_pool_indices_repeated[:-1] + ) + offset = torch.arange( + raw_positions.numel(), device=raw_positions.device + ) + group_first = torch.cummax(torch.where(starts, offset, 0), dim=0).values + swa_replay_start = raw_positions - (offset - group_first) + request_layout = window_layout( + req_pool_indices_repeated, + raw_positions, + capacity=self.token_to_kv_pool.request_window.capacity, + floor=swa_replay_start, + num_groups=num_groups, + ) + swa_page_indices = _pad_last_dim(request_layout.indices) + swa_topk_lengths = request_layout.lengths + elif dspark_block_size is not None: assert ( self.is_dspark_draft and dspark_block_size == self.speculative_num_draft_tokens - 1 @@ -2411,12 +4355,18 @@ class DeepseekV4AttnBackend( f"and is only valid on the DSpark draft backend " f"(is_dspark_draft={self.is_dspark_draft})." ) - swa_page_indices, swa_topk_lengths = self.get_dspark_swa_page_indices( - seq_lens_casual=seq_lens_casual, - req_pool_indices_repeated=req_pool_indices_repeated, - out_loc=out_loc, - block_size=dspark_block_size, + assert swa_replay_start is None, ( + "swa_replay_start is not wired for the DSpark draft window" ) + if dspark_swa_buffers is None: + swa_page_indices, swa_topk_lengths = self.get_dspark_swa_page_indices( + seq_lens_casual=seq_lens_casual, + req_pool_indices_repeated=req_pool_indices_repeated, + out_loc=out_loc, + block_size=dspark_block_size, + ) + else: + swa_page_indices, swa_topk_lengths = dspark_swa_buffers else: swa_page_indices = BuildCausalSwaPageIndices.execute( req_to_token=self.req_to_token, @@ -2425,8 +4375,15 @@ class DeepseekV4AttnBackend( seq_lens_casual=seq_lens_casual, swa_window=SWA_WINDOW, page_index_aligned_size=PAGE_INDEX_ALIGNED_SIZE, + swa_replay_start=swa_replay_start, ) swa_topk_lengths = prep.swa_topk_lengths + if swa_replay_start is not None: + # Slots below the floor are -1; the valid count shrinks to match. + floored = raw_positions - swa_replay_start.to(raw_positions.dtype) + 1 + swa_topk_lengths = torch.minimum( + swa_topk_lengths, floored.clamp_min(0).to(swa_topk_lengths.dtype) + ) page_table = prep.page_table @@ -2441,11 +4398,23 @@ class DeepseekV4AttnBackend( swa_topk_lengths=swa_topk_lengths, index_topk=self.index_topk, present_ratios=self.present_ratios, + low_ratios=self.low_ratios, + request_window_layout=request_layout, + swa_out_cache_loc=( + request_layout.write_loc if request_layout is not None else None + ), ) if need_compress: - core_attn_metadata.init_compression_metadata(num_tokens) - core_attn_metadata.init_flashmla_related(is_prefill=is_prefill) + low_ratio_buffers = ( + build_low_ratio_metadata(seq_lens_casual, out_loc, self.index_topk) + if small_metadata + else None + ) + core_attn_metadata.init_compression_metadata(num_tokens, low_ratio_buffers) + core_attn_metadata.init_flashmla_related( + is_prefill=is_prefill, low_ratio_buffers=low_ratio_buffers + ) if self.trtllm_attn: core_attn_metadata.init_trtllm_sparse_buffers() else: diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index f369cc50f..34c2d6cc7 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -11,6 +11,7 @@ from sglang.kernels.ops.attention.dsv4 import ( compress_forward, compress_norm_rope_store, ) +from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout from sglang.srt.environ import envs if TYPE_CHECKING: @@ -158,6 +159,7 @@ class CompressorBackendMixin: bf16_store: bool = False, kv_scale_cache: Optional[torch.Tensor] = None, rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + kv_layout: KVLayout = KVLayout.V4, fp8_2buff: bool = False, kv_cache_rope: Optional[torch.Tensor] = None, ) -> None: @@ -215,6 +217,7 @@ class CompressorBackendMixin: bf16_store=bf16_store, kvcache_scale=kv_scale_cache, rope_cache=rope_cache, + layout=kv_layout, # Derived once per forward by the backend; every C4 layer writes the # same rows to the same slots. fp4_k_write_metadata=( @@ -268,6 +271,7 @@ class CompressorBackendMixin: ) use_hip_fp4 = _is_hip and use_fp4_indexer bf16_store = False + kv_layout = KVLayout.V4 kv_scale_cache = None fp8_2buff = False kv_cache_rope = None @@ -295,6 +299,8 @@ class CompressorBackendMixin: assert compress_kv_pool is not None kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id) page_size = token_to_kv_pool.get_extra_key_page_size(layer_id) + # The pool's page format (V4, or the V4.1 fp8 / fp4 layouts). + kv_layout = token_to_kv_pool.get_extra_key_layout(layer_id) if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"): out_loc = compress_kv_pool._translate_loc_to_hisparse_device(out_loc) self._forward_compress_all_in_one( @@ -316,6 +322,7 @@ class CompressorBackendMixin: rope_cache=( (compressor.fp4_cos, compressor.fp4_sin) if use_hip_fp4 else None ), + kv_layout=kv_layout, fp8_2buff=fp8_2buff, kv_cache_rope=( None if kv_cache_rope is None else kv_cache_rope.view(dtype=torch.uint8) diff --git a/python/sglang/srt/layers/attention/dsv4/dsv41_sparse.py b/python/sglang/srt/layers/attention/dsv4/dsv41_sparse.py new file mode 100644 index 000000000..797294c4d --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/dsv41_sparse.py @@ -0,0 +1,266 @@ +"""DeepSeek V4.1 ratio-1/2 compressors and indexers. + +Only kv_source layers own compressed latents; later layers of the same ratio +share that storage. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +from torch import nn + +from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32 +from sglang.kernels.ops.attention.dsv4.torch_quant import ( + fake_quant_compressed_kv, + fake_quant_fp4, +) +from sglang.kernels.ops.layernorm.rmsnorm_fp32 import rmsnorm_fp32 +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.utils import add_prefix + + +def _rope_fq4(x, freqs, rope_dim, *, compressed_kv=False): + if x.is_cuda and torch.version.cuda is not None and x.dtype == torch.bfloat16: + from sglang.kernels.ops.attention.dsv4.fp4_rope_fake_quant import ( + rope_tail_fake_quant_fp4, + ) + + return rope_tail_fake_quant_fp4(x, freqs, rope_dim, compressed_kv=compressed_kv) + quant = fake_quant_compressed_kv if compressed_kv else fake_quant_fp4 + return quant(rope_tail(x, freqs, rope_dim)) + + +class RMSNorm(nn.Module): + """fp32 statistics and fp32 weight multiply, cast back at the very end.""" + + def __init__(self, dim: int, eps: float): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if ( + x.is_cuda + and torch.version.cuda is not None + and x.dtype in (torch.bfloat16, torch.float32) + and self.weight.dtype in (torch.bfloat16, torch.float32) + and x.shape[-1] in (128, 512) + and x.is_contiguous() + and self.weight.is_contiguous() + ): + return rmsnorm_fp32(x, self.weight, self.eps) + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +def token_req_indices(forward_batch, *, num_tokens=None) -> torch.Tensor: + req = forward_batch.req_pool_indices.to(torch.int64) + if forward_batch.forward_mode.is_decode(): + return req + if forward_batch.forward_mode.is_target_verify(): + return torch.repeat_interleave( + req, int(forward_batch.spec_info.draft_token_num), output_size=num_tokens + ) + assert forward_batch.forward_mode.is_extend(), ( + "the V4.1 torch attention path serves extend, target-verify and decode" + ) + return torch.repeat_interleave( + req, forward_batch.extend_seq_lens.to(torch.int64), output_size=num_tokens + ) + + +def rope_tail( + x: torch.Tensor, freqs: torch.Tensor, rope_dim: int, inverse: bool = False +) -> torch.Tensor: + """Rotate the last rope_dim features of x with complex freqs [T, rope_dim // 2].""" + head, tail = x[..., :-rope_dim], x[..., -rope_dim:] + tc = torch.view_as_complex(tail.float().unflatten(-1, (-1, 2)).contiguous()) + f = freqs.conj() if inverse else freqs + f = f.view(x.shape[0], *([1] * (x.ndim - 2)), rope_dim // 2) + rotated = torch.view_as_real(tc * f).flatten(-2).to(x.dtype) + return torch.cat([head, rotated], dim=-1) + + +def fused_low_ratio_compress_supported() -> bool: + """The fused c1 / c2 / index-K decode kernels pack fp4 with + `cvt.rn.satfinite.e2m1x2`, an sm100+ instruction; the answer also fixes the + ratio-2 weight layout (`wkv_gate`, or `wkv` plus `wgate`).""" + if not torch.cuda.is_available() or torch.version.hip is not None: + return False + return torch.cuda.get_device_capability()[0] >= 10 + + +class DeepseekV41Compressor(nn.Module): + """Pool consecutive tokens into one pre-RoPE KV latent; bf16 weights, fp32 + projection and softmax pooling, rounded back to bf16 in `finish`.""" + + def __init__( + self, + hidden_size: int, + head_dim: int, + compress_ratio: int, + eps: float, + *, + fused_compress: Optional[bool] = None, + ): + super().__init__() + self.compress_ratio = compress_ratio + self.norm = RMSNorm(head_dim, eps) + # The loader concatenates ratio-2 wkv/wgate when it finds wkv_gate.weight; + # ratio 1 must retain wkv.weight because there is no gate half to load. + self.use_fused_compress = ( + fused_low_ratio_compress_supported() + if fused_compress is None + else bool(fused_compress) + ) + self.use_fused_gate = compress_ratio > 1 and self.use_fused_compress + if self.use_fused_gate: + self.wkv_gate = nn.Linear( + hidden_size, 2 * head_dim, bias=False, dtype=torch.bfloat16 + ) + else: + self.wkv = nn.Linear( + hidden_size, head_dim, bias=False, dtype=torch.bfloat16 + ) + if compress_ratio > 1: + self.wgate = nn.Linear( + hidden_size, head_dim, bias=False, dtype=torch.bfloat16 + ) + + def project(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if self.compress_ratio == 1: + return self.wkv(x), None + if self.use_fused_gate: + fused = self.project_fused(x) + head_dim = fused.shape[-1] // 2 + return fused[..., :head_dim], fused[..., head_dim:] + # Two GEMMs, not one fused [2D, K] projection: c2_decode_pool reads kv and + # score as contiguous [n, D] rows; column slices of a fused output are not. + kv = linear_bf16_fp32(x, self.wkv.weight) + score = linear_bf16_fp32(x, self.wgate.weight) + return kv, score + + def project_fused(self, x: torch.Tensor) -> torch.Tensor: + """`[n, 2D]` fp32, `| kv | score |`.""" + return linear_bf16_fp32(x, self.wkv_gate.weight) + + def finish(self, kv: torch.Tensor) -> torch.Tensor: + return self.norm(kv.to(torch.bfloat16)) + + @staticmethod + def pool_pairs(kv2: torch.Tensor, score2: torch.Tensor) -> torch.Tensor: + """kv2, score2 [n, 2, D] fp32 -> [n, D]""" + return (kv2 * score2.softmax(dim=1)).sum(dim=1) + + +def _small_weights_proj_max_m(n_heads: int, hidden_size: int) -> int: + # -1 means the device or checkpoint shape requires the linear fallback. + if not torch.cuda.is_available() or torch.version.hip is not None: + return -1 + from sglang.kernels.ops.gemm.small_gemm_bf16 import MAX_M, can_use_n32k5120_gemm + + return MAX_M if can_use_n32k5120_gemm(n_heads, hidden_size, 1) else -1 + + +class DeepseekV41Indexer(nn.Module): + """Scores compressed positions with a small fp4 side attention; only a + kv_source layer owns index keys. Projections are replicated across TP: every + rank scores with all heads, so the top-k needs no cross-rank reduction.""" + + def __init__( + self, + config, + layer_id: int, + head_dim: int, + quant_config: Optional[QuantizationConfig], + prefix: str, + ): + super().__init__() + self.n_heads = config.index_n_heads + self.n_local_heads = self.n_heads + self.index_head_dim = config.index_head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.index_topk = config.index_topk + self.owns_k = layer_id in config.kv_source_layer_ids + self.is_candidate_source = layer_id == config.candidate_source_layer_id + self.uses_candidates = 0 <= config.candidate_source_layer_id < layer_id + self.candidate_topk_blocks = config.candidate_topk_blocks + self.candidate_block_size = config.candidate_block_size + self.softmax_scale = self.index_head_dim**-0.5 + self.wq_b = ReplicatedLinear( + config.q_lora_rank, + self.n_heads * self.index_head_dim, + bias=False, + quant_config=quant_config, + params_dtype=torch.bfloat16, + prefix=add_prefix("wq_b", prefix), + ) + self.weights_proj = ReplicatedLinear( + config.hidden_size, + self.n_heads, + bias=False, + params_dtype=torch.bfloat16, + quant_config=None, + prefix=add_prefix("weights_proj", prefix), + ) + # The decode GEMM matches tiny_gemm's reduction order, not cuBLAS's. + self.weights_proj_small_max_m = _small_weights_proj_max_m( + self.n_heads, config.hidden_size + ) + if self.owns_k: + self.wk = nn.Linear( + head_dim, self.index_head_dim, bias=False, dtype=torch.bfloat16 + ) + self.k_norm = RMSNorm(self.index_head_dim, config.rms_norm_eps) + + def forward_wk(self, latent: torch.Tensor) -> torch.Tensor: + from sglang.kernels.ops.gemm.small_gemm_bf16 import ( + can_use_n128k512_gemm, + n128k512_gemm_bf16, + ) + + # The JIT kernel raises rather than falling back on an unsupported shape. + if can_use_n128k512_gemm( + self.index_head_dim, latent.shape[-1], latent.shape[0] + ): + return n128k512_gemm_bf16(latent, self.wk.weight) + return self.wk(latent) + + def index_keys(self, latent: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Pre-RoPE latents [n, D] -> fp4-rounded index keys [n, index_head_dim].""" + k = self.k_norm(self.forward_wk(latent)) + return _rope_fq4(k, freqs, self.rope_head_dim) + + def queries(self, q_lora: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + q, _ = self.wq_b(q_lora) + q = q.view(q.shape[0], self.n_local_heads, self.index_head_dim) + return _rope_fq4(q, freqs, self.rope_head_dim) + + def head_weights_raw(self, x: torch.Tensor) -> torch.Tensor: + """`weights_proj(x)` before the scale, [tokens, n_heads] bf16.""" + if 0 < x.shape[0] <= self.weights_proj_small_max_m and x.is_cuda: + from sglang.kernels.ops.gemm.small_gemm_bf16 import n32k5120_gemm_bf16 + + return n32k5120_gemm_bf16(x, self.weights_proj.weight) + w, _ = self.weights_proj(x) + return w + + @property + def head_weight_scale(self) -> float: + return self.softmax_scale * self.n_heads**-0.5 + + def head_weights(self, x: torch.Tensor) -> torch.Tensor: + return self.head_weights_raw(x) * self.head_weight_scale + + def scores( + self, q: torch.Tensor, k: torch.Tensor, weights: torch.Tensor + ) -> torch.Tensor: + """q [t, H, d], k [n, d], weights [t, H] -> [t, n], summed over all heads.""" + s = torch.einsum("bhd,nd->bhn", q, k) + s = (s.relu() * weights.unsqueeze(-1)).sum(dim=1) + return s.float() diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 4c4741bde..0080ad024 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -116,6 +116,11 @@ class PagedIndexerMetadata: use_topk_v2: bool force_deep_gemm_metadata: bool = False use_prefill_cuda_graph: bool = False + # Indexer source compression ratio: 4 for c4, 1 or 2 for the dsv41 sources. + compress_ratio: int = 4 + # Rows per logits chunk for the prefill CUDA graph low-ratio indexer; 0 plans + # all rows at once. + row_chunk: int = 0 deep_gemm_metadata: Any = field(init=False, repr=False) topk_metadata: torch.Tensor = field(init=False, repr=False) nonpaged_plan: Optional[NonPagedIndexerPlan] = field( @@ -144,7 +149,18 @@ class PagedIndexerMetadata: compressed_seq_lens = self.compressed_seq_lens.to(torch.int32) if compressed_seq_lens.dim() == 1: compressed_seq_lens = compressed_seq_lens.unsqueeze(-1) - if _IS_SM120 and compressed_seq_lens.shape[0] > _SM120_INDEXER_M_CHUNK: + if self.row_chunk > 0: + self.deep_gemm_metadata = torch.stack( + [ + get_paged_mqa_logits_metadata( + compressed_seq_lens[_s : _s + self.row_chunk], + self.compressed_page_size, + deep_gemm.get_num_sms(), + ) + for _s in range(0, compressed_seq_lens.shape[0], self.row_chunk) + ] + ) + elif _IS_SM120 and compressed_seq_lens.shape[0] > _SM120_INDEXER_M_CHUNK: # Chunk metadata is shared by all indexer layers in this forward. self.deep_gemm_metadata = [ get_paged_mqa_logits_metadata( @@ -173,6 +189,9 @@ class PagedIndexerMetadata: self.topk_metadata = torch.empty((0,)) assert self.page_size == 256, "the system hardcodes page_size=256" + assert self.page_size % self.compress_ratio == 0, ( + f"compress_ratio {self.compress_ratio} must divide page_size {self.page_size}" + ) @property def max_seq_len(self) -> int: @@ -182,6 +201,17 @@ class PagedIndexerMetadata: def max_compressed_seq_len(self) -> int: return self.page_table.shape[1] * self.compressed_page_size + def row_chunks(self): + num_rows = self.compressed_seq_lens.shape[0] + if self.row_chunk <= 0: + return [(slice(0, num_rows), self.deep_gemm_metadata)] + return [ + (slice(start, min(start + self.row_chunk, num_rows)), plan) + for start, plan in zip( + range(0, num_rows, self.row_chunk), self.deep_gemm_metadata + ) + ] + def copy_(self, other: PagedIndexerMetadata): if is_hip(): copy_fields = ["page_table", "compressed_seq_lens"] @@ -196,6 +226,8 @@ class PagedIndexerMetadata: check_eq_fields=[ "page_size", "compressed_page_size", + "compress_ratio", + "row_chunk", "force_deep_gemm_metadata", "use_prefill_cuda_graph", "use_topk_v2", diff --git a/python/sglang/srt/layers/attention/graph_variants.py b/python/sglang/srt/layers/attention/graph_variants.py index 14ba77c7b..6c34620d2 100644 --- a/python/sglang/srt/layers/attention/graph_variants.py +++ b/python/sglang/srt/layers/attention/graph_variants.py @@ -57,3 +57,90 @@ def create_attention_graph_variants(hf_config) -> Optional[AttentionGraphVariant ) return DsaGraphVariants(index_topk) return None + + +DSV41_CANDIDATE_FILTERED = "candidate_filtered" + + +@dataclass(frozen=True) +class Dsv41CandidateGraphVariants: + """Candidate-indexer graphs keyed by the batch's longest request; a variant + below its limit skips low-ratio scoring or candidate filtering.""" + + # (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: + 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 + return DSV41_CANDIDATE_FILTERED + + +def create_dsv41_candidate_graph_variants( + model_runner, capture_forward_mode, captured_req_width: int = 0 +) -> Optional[Dsv41CandidateGraphVariants]: + 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 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 + and getattr(text_config, "model_type", None) == "deepseek_v41" + and getattr(text_config, "candidate_source_layer_id", -1) >= 0 + ): + return None + span = text_config.candidate_topk_blocks * text_config.candidate_block_size + if span <= 0: + return None + ratios = set(text_config.compress_ratios) & {1, 2} + topk = text_config.index_topk + variants = [] + # Verify needs per-query causal top-k, so it always keeps candidate filtering. + 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)) + variants.append(("candidate_unfiltered", span)) + graph_limits = [] + for variant, limit in variants: + graph_limits.append((variant, min(limit, span))) + if limit >= span: + break + logger.info( + "Candidate indexer graph limits: %s; use full filtering above %s.", + graph_limits, + span, + ) + 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/cp/interleave.py b/python/sglang/srt/layers/cp/interleave.py index 53da7f6c1..b5542de61 100644 --- a/python/sglang/srt/layers/cp/interleave.py +++ b/python/sglang/srt/layers/cp/interleave.py @@ -55,6 +55,9 @@ class InterleaveContextParallelMetadata(BaseContextParallelMetadata): per_rank_actual_token: Optional[List[int]] = None max_rank_len: Optional[List[int]] = None per_rank_logical_token: Optional[List[int]] = None + # Tail row -> packed all-gather slot; local tail metadata rows include padding. + gather_index: Optional[torch.Tensor] = None + local_index: Optional[torch.Tensor] = None class InterleaveCPStrategy(ContextParallelStrategy): @@ -213,6 +216,9 @@ class InterleaveCPStrategy(ContextParallelStrategy): gathered = x.new_empty((self.cp_size * physical_rank_len, *x.shape[1:])) attn_cp_all_gather_into_tensor(gathered, padded_x.contiguous()) + if metadata.gather_index is not None: + return gathered.index_select(0, metadata.gather_index) + # Equal per-rank lengths: one interleave copy restores the original # token order; cheaper than the index_select fallback below. actual = metadata.per_rank_actual_token @@ -290,3 +296,15 @@ class InterleaveCPStrategy(ContextParallelStrategy): k_nope = full_latent[..., :kv_lora_rank].unsqueeze(1) k_rope = full_latent[..., kv_lora_rank:].unsqueeze(1) return k_nope, k_rope + + +def interleave_rows_per_request( + extend_lens: List[int], cp_rank: int, cp_size: int +) -> List[int]: + """Rows of each request a CP rank holds: global token index congruent to cp_rank.""" + counts, start = [], 0 + for n in extend_lens: + end = start + n + counts.append((end - 1 - cp_rank) // cp_size - (start - 1 - cp_rank) // cp_size) + start = end + return counts diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 48cf1bc0c..9cc7b795e 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -939,6 +939,15 @@ def can_use_flashinfer_allreduce( # Dynamo, so statically-off configs must short-circuit before reaching them # (same ordering rule as apply_flashinfer_allreduce_fusion). token_num, hidden_dim = input_.shape + + # MNNVL hard-fails instead of falling back when the width is not float4-aligned + # (FlashInfer csrc/trtllm_mnnvl_allreduce.cu). + if ( + workspace_manager.backend == "mnnvl" + and hidden_dim % (16 // input_.element_size()) != 0 + ): + return False + if torch.compiler.is_compiling(): # Don't call into the flashinfer workspace object while tracing. The # workspace was allocated for (max_token_num, hidden_dim, dtype) and diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index b488bfa9d..9ebd61ba9 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -209,6 +209,9 @@ class LogitsProcessorOutput: # The last hidden layers hidden_states: Optional[torch.Tensor] = None + # Original flattened token indices when only a subset of hidden rows is captured. + hidden_states_token_indices: Optional[torch.Tensor] = None + ## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler # he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature. next_token_logprobs: Optional[torch.Tensor] = None diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index 5e37f4aa4..bd7c2a82d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -1505,7 +1505,7 @@ class FusedMoE(torch.nn.Module): self, hidden_states: torch.Tensor, topk_output: TopKOutput, - pre_quant_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + pre_quant_input: Optional[Tuple] = None, ): if self._use_ascend_fuseep: from sglang.srt.hardware_backend.npu.moe.fuseep import forward_fuseep @@ -1546,7 +1546,7 @@ class FusedMoE(torch.nn.Module): self, hidden_states: torch.Tensor, topk_output: TopKOutput, - pre_quant_input: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + pre_quant_input: Optional[Tuple] = None, ): origin_hidden_states_dim = hidden_states.shape[-1] assert self.quant_method is not None @@ -1555,21 +1555,9 @@ class FusedMoE(torch.nn.Module): dwdp_mgr = get_global_dwdp_manager() dwdp_mgr.wait_prefetch(self.layer_id) - dispatch_output = self.dispatcher.dispatch( - hidden_states=hidden_states, topk_output=topk_output + dispatch_output = self._dispatch_with_pre_quant( + hidden_states, topk_output, pre_quant_input ) - if ( - pre_quant_input is not None - and dispatch_output.format.is_standard() - and dispatch_output.hidden_states_scale is None - ): - # SGLANG_OPT_MOE_QUANT_ONCE: the standard dispatch was a pure - # passthrough, so the caller's pre-quantized (q, scale) pair still - # matches dispatch_output.hidden_states; attach it for the triton - # fused runner to skip its own activation quant. - dispatch_output = dispatch_output._replace( - hidden_states_pre_quant=pre_quant_input - ) combine_input = self.run_moe_core( dispatch_output=dispatch_output, @@ -1593,16 +1581,40 @@ class FusedMoE(torch.nn.Module): return final_hidden_states + def _dispatch_with_pre_quant( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + pre_quant_input: Optional[Tuple], + ) -> DispatchOutput: + dispatch_output = self.dispatcher.dispatch( + hidden_states=hidden_states, topk_output=topk_output + ) + if ( + pre_quant_input is not None + and dispatch_output.format.is_standard() + and dispatch_output.hidden_states_scale is None + ): + # Dropping an Mxfp8RoutedInputPreQuant here would leave its side + # stream unjoined under CUDA-graph capture. + dispatch_output = dispatch_output._replace( + hidden_states_pre_quant=pre_quant_input + ) + return dispatch_output + def forward_deferred_finalize( - self, hidden_states: torch.Tensor, topk_output: TopKOutput + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + pre_quant_input: Optional[Tuple] = None, ): assert self.quant_method is not None from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( flashinfer_trtllm_deferred_finalize_context, ) - dispatch_output = self.dispatcher.dispatch( - hidden_states=hidden_states, topk_output=topk_output + dispatch_output = self._dispatch_with_pre_quant( + hidden_states, topk_output, pre_quant_input ) with flashinfer_trtllm_deferred_finalize_context(): diff --git a/python/sglang/srt/layers/moe/mhc_post_fusion.py b/python/sglang/srt/layers/moe/mhc_post_fusion.py new file mode 100644 index 000000000..64b90d32e --- /dev/null +++ b/python/sglang/srt/layers/moe/mhc_post_fusion.py @@ -0,0 +1,48 @@ +"""Scoped handoff of a decoder's HC post operands to deferred MoE finalize.""" + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Callable, Optional + +import torch + + +@dataclass +class MhcPostFusion: + residual: torch.Tensor + post: Optional[torch.Tensor] + comb: Optional[torch.Tensor] + stats_stream: Optional[torch.cuda.Stream] + output: Optional[torch.Tensor] = None + pre: Optional[torch.Tensor] = None + norm_weight: Optional[torch.Tensor] = None + norm_eps: float = 0.0 + normalized: Optional[torch.Tensor] = None + quantized: Optional[tuple[torch.Tensor, torch.Tensor]] = None + record_stats: Optional[ + Callable[[], tuple[torch.Tensor, torch.Tensor, torch.Tensor]] + ] = None + + def materialize_stats(self): + # Record after the main parent; graph replay must keep the join on the + # main stream. + if self.record_stats is not None: + self.pre, self.post, self.comb = self.record_stats() + self.record_stats = None + + +_current: ContextVar[Optional[MhcPostFusion]] = ContextVar("moe_mhc_post", default=None) + + +def current_mhc_post_fusion(): + return _current.get() + + +@contextmanager +def use_mhc_post_fusion(state): + token = _current.set(state) + try: + yield + finally: + _current.reset(token) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 4d307df24..cc25486b9 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -75,6 +75,10 @@ def flashinfer_trtllm_deferred_finalize_context( _deferred_finalize_enabled.reset(token) +def is_deferred_finalize_enabled() -> bool: + return _deferred_finalize_enabled.get() + + def finalize_flashinfer_trtllm_deferred_output( deferred_output: FlashInferTrtllmDeferredFinalizeOutput, shared_output: torch.Tensor, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/standard.py b/python/sglang/srt/layers/moe/token_dispatcher/standard.py index 89a4bd0a1..b712c8128 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/standard.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/standard.py @@ -68,11 +68,10 @@ class StandardDispatchOutput(NamedTuple): hidden_states: torch.Tensor hidden_states_scale: Optional[torch.Tensor] topk_output: TopKOutput - # SGLANG_OPT_MOE_QUANT_ONCE: optional pre-quantized (q, scale) pair for - # ``hidden_states`` (per-token-group-128 fp8, q rows possibly padded to a - # multiple of 4). Consumed by the standard->triton fused runner so it can - # skip its own activation quant; ``hidden_states`` itself stays bf16. - hidden_states_pre_quant: Optional[Tuple[torch.Tensor, torch.Tensor]] = None + # Pre-quantized activation for ``hidden_states``, which itself stays bf16: + # either a (q, scale) pair (per-token-group-128 fp8, q rows padded to a + # multiple of 4) or an ``Mxfp8RoutedInputPreQuant``. + hidden_states_pre_quant: Optional[Tuple] = None @property def format(self) -> DispatchOutputFormat: diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index f90273088..a734ca560 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -26,6 +26,7 @@ from typing import ( Protocol, Tuple, TypeGuard, + Union, runtime_checkable, ) @@ -235,6 +236,10 @@ class TopKConfig: fused_shared_experts_scaling_factor: Optional[float] = None output_format: Optional[TopKOutputFormat] = None scoring_func: str = "softmax" + # sqrtsoftplus through log1p with NaNs ranked first (DeepSeek-V4.1 routing). + sqrtsoftplus_log1p: bool = False + # Let the fused router also emit FlashInfer routed-MoE packed ids. + fused_gate_packed_ids: bool = False # Draft-side MoE blocks set this False so they never write the target's # process-global routed-experts capture buffer. allow_routed_experts_capture: bool = True @@ -271,15 +276,10 @@ class TopKConfig: class TopKOutputChecker: @staticmethod - def format_is_standard(topk_output: TopKOutput) -> TypeGuard[StandardTopKOutput]: - # ===== TO BE REFACTORED ==== - # The experimental fused topk+pack carrier only exists under the master switch. - if _SGLANG_EXPERIMENTAL_LORA_OPTI: - return isinstance( - topk_output, (StandardTopKOutput, StandardTopKOutputPacked) - ) - # ===== END TO BE REFACTORED ==== - return isinstance(topk_output, StandardTopKOutput) + def format_is_standard( + topk_output: TopKOutput, + ) -> TypeGuard[Union[StandardTopKOutput, StandardTopKOutputPacked]]: + return isinstance(topk_output, (StandardTopKOutput, StandardTopKOutputPacked)) @staticmethod def format_is_triton_kernels( @@ -325,11 +325,8 @@ class StandardTopKOutput(NamedTuple): return TopKOutputFormat.STANDARD -# ===== TO BE REFACTORED ==== -# Experimental fused topk+pack (SGLANG_OPT_LORA_FUSED_TOPK_PACK) carrier: the FlashInfer -# routed-MoE packed topk produced fused in the gating kernel. Kept a SEPARATE type rather -# than a 4th StandardTopKOutput field so the OSS `a, b, _ = topk_output` 3-tuple unpack -# stays valid; only the gated experimental MoE dispatch reads .packed_topk_ids (getattr). +# Standard top-k output plus the FlashInfer routed-MoE packed ids that +# ``moe_fused_gate`` writes; a separate type keeps the 3-tuple unpack valid. class StandardTopKOutputPacked(NamedTuple): topk_weights: torch.Tensor topk_ids: torch.Tensor @@ -341,9 +338,6 @@ class StandardTopKOutputPacked(NamedTuple): return TopKOutputFormat.STANDARD -# ===== END TO BE REFACTORED ==== - - class TritonKernelTopKOutput(NamedTuple): """Triton kernel top-k output format.""" @@ -546,6 +540,8 @@ class TopK(BaseFusedOp): fused_shared_experts_scaling_factor: Optional[float] = None, is_fp4_experts: bool = False, allow_routed_experts_capture: bool = True, + sqrtsoftplus_log1p: bool = False, + fused_gate_packed_ids: bool = False, ): # NOTE: scoring_func is not used for now, but we keep it for future use # see https://github.com/sgl-project/sglang/pull/4505 for more details @@ -584,6 +580,8 @@ class TopK(BaseFusedOp): fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor, output_format=output_format, scoring_func=scoring_func, + sqrtsoftplus_log1p=sqrtsoftplus_log1p, + fused_gate_packed_ids=fused_gate_packed_ids, allow_routed_experts_capture=allow_routed_experts_capture, ) @@ -1387,10 +1385,13 @@ def biased_topk_jit_kernel_impl( num_token_non_padded: Optional[torch.Tensor] = None, expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, apply_routed_scaling_factor_on_output: Optional[bool] = False, + packed_out: Optional[torch.Tensor] = None, + sqrtsoftplus_log1p: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" if _use_aiter and scoring_func == "sqrtsoftplus" and num_fused_shared_experts == 0: + assert packed_out is None, "aiter topk_gating cannot emit packed ids" from aiter import topk_gating num_tokens = gating_output.shape[0] @@ -1429,6 +1430,14 @@ def biased_topk_jit_kernel_impl( renormalize=renormalize, routed_scaling_factor=routed_scaling_factor, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + num_token_non_padded=( + num_token_non_padded + if _fused_gate_masks_padded_rows(scoring_func) + else None + ), + # Optional FlashInfer routed-MoE packed ids, written in the same launch. + packed_out=packed_out, + sqrtsoftplus_log1p=sqrtsoftplus_log1p, ) topk_weights, topk_ids = ( topk_weights.to(torch.float32), @@ -1586,6 +1595,31 @@ def _eplb_remap_enabled() -> bool: ) +def _fused_gate_masks_padded_rows(scoring_func: str) -> bool: + # Sigmoid is excluded: a padding count bypasses moe_fused_gate's radix fast + # path, and HIP fills padded ids with 0, not -1. + return _is_cuda and not _use_aiter and scoring_func == "sqrtsoftplus" + + +def _fused_gate_emits_packed_ids( + scoring_func: str, + num_fused_shared_experts: int, + expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo], + routing_overridden: bool, + enabled: bool, +) -> bool: + # The pack is taken from the router's final values, so every condition past + # the caller's opt-in rules out a later rewrite of ids or weights. + return ( + enabled + and _fused_gate_masks_padded_rows(scoring_func) + and get_moe_runner_backend().is_flashinfer_mxfp4() + and expert_location_dispatch_info is None + and num_fused_shared_experts == 0 + and not routing_overridden + ) + + def _mask_topk_ids_padded_region( topk_ids: torch.Tensor, num_token_non_padded: Optional[torch.Tensor] = None, @@ -2162,6 +2196,7 @@ def _post_process_topk_ids( layer_id: int, num_token_non_padded: Optional[torch.Tensor] = None, expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None, + padded_rows_masked: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: num_fused_shared_experts = topk_config.num_fused_shared_experts use_per_rank_shared_slots = has_per_rank_fused_shared_slots( @@ -2204,8 +2239,16 @@ def _post_process_topk_ids( # ExpertDistributionRecorder tracks only EPLB physical routed experts. recorder_topk_ids = routed_cols else: + # A remap table indexed by -1 aliases its last entry, so only the + # identity-remap branch may drop the padded-row mask. topk_ids = _biased_grouped_topk_postprocess( - topk_ids, expert_location_dispatch_info, num_token_non_padded + topk_ids, + expert_location_dispatch_info, + ( + None + if padded_rows_masked and expert_location_dispatch_info is None + else num_token_non_padded + ), ) elif _is_hip: # On AMD HIP the aiter MoE kernels do not handle topk_ids=-1 safely @@ -2374,8 +2417,19 @@ def select_experts( scoring_func = topk_config.scoring_func - # Set by the fused-gating+pack branch below; None everywhere else. + # Set by the fused-gating+pack branches below; None everywhere else. packed_topk = None + # True when the router itself masked rows >= num_token_non_padded. + padded_rows_masked = False + + simulate_uniform_experts = envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() + simulate_round_robin_experts = envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() + if simulate_uniform_experts and simulate_round_robin_experts: + raise ValueError( + "SGLANG_SIMULATE_UNIFORM_EXPERTS and " + "SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS are mutually exclusive" + ) + routing_overridden = simulate_uniform_experts or simulate_round_robin_experts ( router_logits, @@ -2481,6 +2535,22 @@ def select_experts( scoring_func == "sqrtsoftplus" or scoring_func == "sigmoid" ): _biased_topk = biased_topk_xpu if _is_xpu else biased_topk_jit_kernel_impl + _packed_kwargs = {} + if _fused_gate_emits_packed_ids( + scoring_func, + num_fused_shared_experts, + expert_location_dispatch_info, + routing_overridden, + topk_config.fused_gate_packed_ids, + ): + packed_topk = torch.empty( + (hidden_states.shape[0], top_k), + dtype=torch.int32, + device=hidden_states.device, + ) + _packed_kwargs = dict(packed_out=packed_topk) + if topk_config.sqrtsoftplus_log1p: + _packed_kwargs["sqrtsoftplus_log1p"] = True topk_weights, topk_ids = _biased_topk( hidden_states=hidden_states, gating_output=router_logits, @@ -2493,7 +2563,9 @@ def select_experts( num_token_non_padded=num_token_non_padded, expert_location_dispatch_info=expert_location_dispatch_info, apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + **_packed_kwargs, ) + padded_rows_masked = _fused_gate_masks_padded_rows(scoring_func) elif ( get_moe_runner_backend().is_flashinfer_trtllm_routed() and scoring_func == "softmax" @@ -2522,8 +2594,7 @@ def select_experts( and correction_bias is None and expert_location_dispatch_info is None and num_fused_shared_experts == 0 - and not envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() - and not envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() + and not routing_overridden ): num_experts = router_logits.shape[-1] if num_experts & (num_experts - 1) == 0 and num_experts <= 512: @@ -2569,15 +2640,7 @@ def select_experts( renormalize=renormalize, ) - simulate_uniform_experts = envs.SGLANG_SIMULATE_UNIFORM_EXPERTS.get() - simulate_round_robin_experts = envs.SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS.get() - if simulate_uniform_experts and simulate_round_robin_experts: - raise ValueError( - "SGLANG_SIMULATE_UNIFORM_EXPERTS and " - "SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS are mutually exclusive" - ) - - if simulate_uniform_experts or simulate_round_robin_experts: + if routing_overridden: # Benchmark-only: override gating with a balanced expert assignment (so # dummy/random benchmark tokens don't skew MoE load) via a single fused # Triton kernel — one launch instead of the ~5-7 small elementwise ops it @@ -2600,6 +2663,8 @@ def select_experts( token_shard_rank=token_shard_rank, num_token_shards=num_token_shards, ) + # The override rewrote every row, including the router-masked ones. + padded_rows_masked = False topk_ids, topk_weights, recorder_topk_ids = _post_process_topk_ids( topk_ids=topk_ids, @@ -2609,18 +2674,17 @@ def select_experts( num_token_non_padded=num_token_non_padded, layer_id=layer_id, expert_location_dispatch_info=expert_location_dispatch_info, + padded_rows_masked=padded_rows_masked, ) get_global_expert_distribution_recorder().on_select_experts( topk_ids=recorder_topk_ids ) - # ===== TO BE REFACTORED ==== if packed_topk is not None: return StandardTopKOutputPacked( topk_weights, topk_ids, router_logits, packed_topk ) - # ===== END TO BE REFACTORED ==== return StandardTopKOutput(topk_weights, topk_ids, router_logits) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 053638c66..9edc1d962 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -516,6 +516,8 @@ class Fp8LinearMethod(LinearMethodBase): 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. + 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() @@ -1154,6 +1156,19 @@ class Fp8LinearMethod(LinearMethodBase): if mxfp8_view: backend = self.mxfp8_dense_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(): @@ -1172,7 +1187,7 @@ class Fp8LinearMethod(LinearMethodBase): bias=bias, **extra_kwargs, ) - return self.w8a8_mxfp8_linear( + out = self.w8a8_mxfp8_linear( input=x, weight=layer.weight, weight_scale=weight_scale, @@ -1180,6 +1195,7 @@ class Fp8LinearMethod(LinearMethodBase): bias=bias, **extra_kwargs, ) + return out if self.block_quant: if use_intel_amx_backend(layer): diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 976b609af..0810193cb 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -746,9 +746,13 @@ def can_serve_block_fp8_as_mxfp8( def dispatch_block_fp8_mxfp8_linear(backend: Mxfp8DenseGemmBackend) -> Callable: """The MXFP8 linear for a block-fp8 weight served as MXFP8.""" if backend.is_flashinfer_cutlass(): - return partial(flashinfer_mxfp8_blockscaled_linear, backend="cutlass") + return partial( + flashinfer_mxfp8_blockscaled_linear, backend="cutlass", pin_tactic=True + ) if backend.is_flashinfer_cutedsl(): - return partial(flashinfer_mxfp8_blockscaled_linear, backend="cute-dsl") + return partial( + flashinfer_mxfp8_blockscaled_linear, backend="cute-dsl", pin_tactic=True + ) return _unsupported_mxfp8_linear @@ -1473,9 +1477,14 @@ def flashinfer_mxfp8_blockscaled_linear( bias: Optional[torch.Tensor] = None, output_dtype: Optional[torch.dtype] = None, backend: str = "cutlass", + pin_tactic: bool = False, ) -> torch.Tensor: """MXFP8 dense linear via FlashInfer mm_mxfp8. `weight_scale` must be the layout - the backend expects, prepared at load time.""" + the backend expects, prepared at load time. + + pin_tactic skips autotuning: tactics tuned per M bucket change the fp32 + reduction order, breaking row-wise batch invariance. + """ input_2d = input.view(-1, input.shape[-1]) output_shape = [*input.shape[:-1], weight.shape[0]] @@ -1507,15 +1516,29 @@ def flashinfer_mxfp8_blockscaled_linear( else: weight_scale_t = weight_scale.t() if weight_scale.ndim == 2 else weight_scale - output = flashinfer_mm_mxfp8( - q_input, - weight.t(), - x_scale_u8, - weight_scale_t, - out_dtype=output_dtype, - use_8x4_sf_layout=False, - backend=backend, - ) + if pin_tactic: + from flashinfer.autotuner import autotune + + with autotune(False, skip_ops={"mxfp8_gemm"}): + output = flashinfer_mm_mxfp8( + q_input, + weight.t(), + x_scale_u8, + weight_scale_t, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend=backend, + ) + else: + output = flashinfer_mm_mxfp8( + q_input, + weight.t(), + x_scale_u8, + weight_scale_t, + out_dtype=output_dtype, + use_8x4_sf_layout=False, + backend=backend, + ) if bias is not None: output += bias diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index 53d088690..1b948e017 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple import torch from torch.nn import Module @@ -97,6 +97,21 @@ def _pad_intermediate_size(layer: Module) -> None: ) +def routed_hidden_size(layer: Module) -> int: + """Hidden size the routed GEMM1 expects (uint8 weights hold two fp4/row).""" + w13 = layer.w13_weight + return w13.shape[2] * 2 if w13.dtype == torch.uint8 else w13.shape[2] + + +class Mxfp8RoutedInputPreQuant(NamedTuple): + """MXFP8 linear-layout quant of the routed MoE input. ``ready`` is recorded on + the producing stream; the consumer must wait on it before the routed MoE op.""" + + x_q: torch.Tensor + x_sf: torch.Tensor + ready: Optional[torch.cuda.Event] + + class Mxfp4FlashinferTrtllmMoEMethod: fuse_routed_scaling_factor_in_topk = True @@ -311,6 +326,24 @@ class Mxfp4FlashinferTrtllmMoEMethod: persistent=False, ) + def quantize_routed_input( + self, hidden_states: torch.Tensor, hidden_size: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + """MXFP8 quant of the routed input, with the scale in the linear + [tokens, hidden // 32] layout the routed MoE op requires.""" + from sglang.srt.layers.quantization.fp8_utils import flashinfer_mxfp8_quantize + + x_quant, x_scale = flashinfer_mxfp8_quantize( + hidden_states, + False, + alignment=hidden_size, + backend=_MXFP8_QUANTIZE_BACKEND, + ) + x_scale = x_scale.view(torch.float8_e4m3fn).reshape( + *hidden_states.shape[:-1], -1 + ) + return x_quant, x_scale + def apply( self, layer: Module, @@ -321,6 +354,7 @@ class Mxfp4FlashinferTrtllmMoEMethod: hidden_states = dispatch_output.hidden_states topk_output = dispatch_output.topk_output + pre_quant = getattr(dispatch_output, "hidden_states_pre_quant", None) w13 = layer.w13_weight w2 = layer.w2_weight @@ -328,7 +362,7 @@ class Mxfp4FlashinferTrtllmMoEMethod: w2_scale = layer.w2_weight_scale_inv intermediate_size = w2.shape[2] * 2 if w2.dtype == torch.uint8 else w2.shape[2] - hidden_size = w13.shape[2] * 2 if w13.dtype == torch.uint8 else w13.shape[2] + hidden_size = routed_hidden_size(layer) num_local_experts = layer.num_local_experts if w13_scale.dim() == 2: @@ -336,17 +370,18 @@ class Mxfp4FlashinferTrtllmMoEMethod: if w2_scale.dim() == 2: w2_scale = w2_scale.reshape(num_local_experts, hidden_size, -1) - if TopKOutputChecker.format_is_standard(topk_output): - topk_ids = topk_output.topk_ids - topk_weights = topk_output.topk_weights - elif TopKOutputChecker.format_is_bypassed(topk_output): + if TopKOutputChecker.format_is_bypassed(topk_output): raise NotImplementedError( "the old code in this branch is WRONG. e.g. it does not consider HashTopK, and may miss args" ) - else: + if not TopKOutputChecker.format_is_standard(topk_output): raise ValueError(f"Unsupported topk output format: {topk_output.format}") + topk_ids = topk_output.topk_ids + topk_weights = topk_output.topk_weights + precision = self.flashinfer_mxfp4_moe_precision + input_ready: Optional[torch.cuda.Event] = None if precision == "bf16": assert hidden_states.dtype == torch.bfloat16 x_quant = hidden_states @@ -360,40 +395,48 @@ class Mxfp4FlashinferTrtllmMoEMethod: value=0.0, ) elif precision == "default": - from sglang.srt.layers.quantization.fp8_utils import ( - flashinfer_mxfp8_quantize, - ) - - x_quant, x_scale = flashinfer_mxfp8_quantize( - hidden_states, - False, - alignment=hidden_size, - backend=_MXFP8_QUANTIZE_BACKEND, - ) - x_scale = x_scale.view(torch.float8_e4m3fn).reshape( - *hidden_states.shape[:-1], -1 - ) + if isinstance(pre_quant, Mxfp8RoutedInputPreQuant): + assert pre_quant.x_q.shape[0] == hidden_states.shape[0] + x_quant, x_scale, input_ready = pre_quant + else: + x_quant, x_scale = self.quantize_routed_input( + hidden_states, hidden_size + ) else: raise NotImplementedError(f"Unsupported mxfp4 moe precision: {precision}") from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + _make_deferred_finalize_output, + is_deferred_finalize_enabled, trtllm_moe_enable_pdl, ) - with use_symmetric_memory( - get_tp_group(), disabled=not is_allocation_symmetric() - ): - num_tokens = x_quant.shape[0] - out_hidden_size = ( - x_quant.shape[-1] * 2 - if x_quant.dtype == torch.uint8 - else x_quant.shape[-1] - ) - symm_output = torch.empty( - num_tokens, out_hidden_size, dtype=torch.bfloat16, device=x_quant.device - ) + num_tokens = x_quant.shape[0] + # Deferred finalize returns the permuted GEMM2 output plus the routing + # triple instead of the finalized [T, hidden] tensor. + defer_finalize = is_deferred_finalize_enabled() + symm_output = None + if not defer_finalize: + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + out_hidden_size = ( + x_quant.shape[-1] * 2 + if x_quant.dtype == torch.uint8 + else x_quant.shape[-1] + ) + symm_output = torch.empty( + num_tokens, + out_hidden_size, + dtype=torch.bfloat16, + device=x_quant.device, + ) - output = trtllm_fp4_block_scale_routed_moe( + if input_ready is not None: + # The op launches the routing kernel, so the join must precede it. + torch.cuda.current_stream().wait_event(input_ready) + + result = trtllm_fp4_block_scale_routed_moe( topk_ids=(topk_ids, topk_weights), routing_bias=None, hidden_states=x_quant, @@ -419,11 +462,15 @@ class Mxfp4FlashinferTrtllmMoEMethod: local_num_experts=num_local_experts, routed_scaling_factor=1.0, routing_method_type=int(RoutingMethodType.TopK), - do_finalize=True, - tune_max_num_tokens=next_power_of_2(x_quant.shape[0]), + do_finalize=not defer_finalize, + tune_max_num_tokens=next_power_of_2(num_tokens), output=symm_output, enable_pdl=trtllm_moe_enable_pdl(num_tokens), - )[0] + ) + if defer_finalize: + output = _make_deferred_finalize_output(result, top_k=topk_ids.shape[1]) + else: + output = result[0] return StandardCombineInput(hidden_states=output) @@ -467,3 +514,58 @@ def maybe_fuse_routed_scale_and_shared_add( if shared is not None: routed += shared return routed + + +# Fused finalize + shared add + TP all-reduce +_fused_finalize_all_reduce_world_size: Optional[int] = None +_fused_finalize_all_reduce_probed = False + + +def _fused_finalize_all_reduce_comm_world_size() -> Optional[int]: + global _fused_finalize_all_reduce_world_size, _fused_finalize_all_reduce_probed + if not _fused_finalize_all_reduce_probed: + _fused_finalize_all_reduce_probed = True + from sglang.kernels.ops.communication import all_reduce_fusion + from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import ( + CustomAllReduceV2, + ) + + ca_comm = get_tp_group().ca_comm + if isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled: + all_reduce_fusion.register_comm(ca_comm.obj) + _fused_finalize_all_reduce_world_size = ca_comm.world_size + else: + log_info_on_rank0( + logger, + "Fused MoE finalize: TP group has no " + "CustomAllReduceV2 push plane; keeping the unfused finalize path", + ) + return _fused_finalize_all_reduce_world_size + + +def should_use_fuse_finalize_all_reduce( + experts, num_tokens: int, hidden_dim: int +) -> bool: + """Capability only; the batch-size policy cap lives at the call site. The + kernel never rescales, so the expert weights must carry the routed scaling.""" + if not isinstance(experts.quant_method, Mxfp4FlashinferTrtllmMoEMethod): + return False + if experts.quant_method.flashinfer_mxfp4_moe_precision != "default": + return False + if not experts.should_fuse_routed_scaling_factor_in_topk: + return False + if num_tokens <= 0: + return False + from sglang.kernels.ops.communication import all_reduce_fusion + + if not all_reduce_fusion.valid_cluster_sizes(hidden_dim): + return False + tp_group = get_tp_group() + if _fused_finalize_all_reduce_comm_world_size() != tp_group.world_size: + return False + # one push phase counter per row (the plane has num_sm of them) + if num_tokens > tp_group.ca_comm.config.num_push_blocks: + return False + return all_reduce_fusion.fits_push_slot( + tp_group.ca_comm.max_push_size, num_tokens, hidden_dim + ) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f154d741a..51167bc60 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2386,6 +2386,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # DeepSeek-V4.1 engram, extend batches only: [bs, n - 1] int32 predecessors # of each request's first extend token (NgramEmbeddingManager). engram_history: Optional[torch.Tensor] = None + encoder_swa_reset: Optional[List[bool]] = None req_pool_indices: torch.Tensor = None # shape: [b], int64 seq_lens: torch.Tensor = None # shape: [b], int64 @@ -2723,6 +2724,26 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.seq_lens_cpu = seq_lens_cpu self.extend_num_tokens = extend_num_tokens + if get_exec().features.enable_encoder_swa_bounded_replay: + for req in reqs: + if ( + req.multimodal_inputs is not None + or req.input_embeds is not None + or req.positional_embed_overrides is not None + ): + raise ValueError( + "encoder SWA replay currently supports token-only text requests" + ) + if req.return_logprob and req.logprob_start_len not in ( + -1, + len(req.origin_input_ids), + ): + raise ValueError( + "encoder SWA replay cannot return cached prompt logprobs" + ) + self.encoder_swa_reset = [ + r.kv.req_pool_idx is None or r.is_retracted for r in reqs + ] # Allocate memory out_cache_loc, req_pool_indices_tensor, req_pool_indices_cpu = alloc_for_extend( self diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index f4082c592..f83b18178 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -7,6 +7,7 @@ from sglang.srt.environ import envs from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor from sglang.srt.runtime_context import ( get_disagg, + get_exec, get_schedule, ) from sglang.srt.utils import get_bool_env_var, is_gfx95_supported, is_hip @@ -667,6 +668,7 @@ class PrefillAdder: self.log_host_hit_tokens = 0 self.log_storage_hit_tokens = 0 self.log_input_tokens = 0 + self.log_replay_tokens = 0 self.reprocessed_log_input_tokens = 0 if running_batch is not None: @@ -897,6 +899,14 @@ class PrefillAdder: self.reprocessed_log_input_tokens += raw_extend_input_len def _account_prefill_cache_admission(self, req: Req, prefix_len: int) -> None: + if get_exec().features.enable_encoder_swa_bounded_replay and ( + req.kv.req_pool_idx is None or req.is_retracted + ): + replay_tokens = min(prefix_len, 128) + self.log_replay_tokens += replay_tokens + self.rem_input_tokens -= replay_tokens + if self.rem_chunk_tokens is not None: + self.rem_chunk_tokens -= replay_tokens if req.retracted_stain: # Retraction attribution is intentionally omitted for now; discard # its lifecycle state so a later abort cannot report it as a drop. diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 10e9426c2..317e4e9aa 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -103,6 +103,7 @@ class PrefillStats: log_host_hit_tokens: int = 0 log_storage_hit_tokens: int = 0 num_pending_tokens: int = 0 + log_replay_tokens: int = 0 @classmethod def from_adder( @@ -114,6 +115,7 @@ class PrefillStats: ): return cls( log_input_tokens=adder.log_input_tokens, + log_replay_tokens=adder.log_replay_tokens, log_hit_tokens=adder.log_hit_tokens, reprocessed_log_input_tokens=adder.reprocessed_log_input_tokens, reprocessed_log_hit_tokens=adder.reprocessed_log_hit_tokens, @@ -660,7 +662,10 @@ class SchedulerMetricsReporter: gap_latency = now - self.last_prefill_stats_tic self.last_prefill_stats_tic = now self.last_input_throughput = ( - prefill_stats.log_input_tokens / gap_latency if gap_latency > 0 else 0.0 + (prefill_stats.log_input_tokens + prefill_stats.log_replay_tokens) + / gap_latency + if gap_latency > 0 + else 0.0 ) pool_stats = self.scheduler.pool_stats_observer.get_pool_stats() @@ -685,6 +690,8 @@ class SchedulerMetricsReporter: f"#pending-token: {prefill_stats.num_pending_tokens}, " ) + if prefill_stats.log_replay_tokens: + msg += f"#replay-token: {prefill_stats.log_replay_tokens}, " if self.scheduler.disaggregation_mode == DisaggregationMode.PREFILL: msg += f"#bootstrap-req: {len(self.scheduler.disagg_prefill_bootstrap_queue.queue)}, " msg += ( @@ -728,7 +735,9 @@ class SchedulerMetricsReporter: value=can_run_cuda_graph ) self.metrics_collector.increment_realtime_tokens( - prefill_compute_tokens=prefill_stats.log_input_tokens, + prefill_compute_tokens=( + prefill_stats.log_input_tokens + prefill_stats.log_replay_tokens + ), prefill_cache_tokens=prefill_stats.log_hit_tokens, dp_cooperation_info=dp_cooperation_info, ) diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 24da7ed13..45b70e172 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1257,6 +1257,28 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ) -> None: """Validates that the input token count and the requested token count doesn't exceed the model's context length.""" # FIXME: unify the length validation logic with the one in the scheduler. + if get_exec().features.enable_encoder_swa_bounded_replay: + if any( + value is not None + for value in ( + obj.image_data, + obj.video_data, + obj.audio_data, + obj.input_embeds, + obj.positional_embed_overrides, + ) + ): + raise ValueError( + "encoder SWA replay currently supports token-only text requests" + ) + if ( + isinstance(obj, GenerateReqInput) + and obj.return_logprob + and obj.logprob_start_len not in (None, -1, len(input_ids)) + ): + raise ValueError( + "encoder SWA replay cannot return cached prompt logprobs" + ) _max_req_len = self.context_len input_token_num = len(input_ids) if input_ids is not None else 0 input_token_num += self.num_reserved_tokens diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index a9067ac1f..cad2ec201 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -642,6 +642,14 @@ class TpModelWorker(BaseTpWorker): # update the consumer index of hicache to the running batch self.set_hicache_consumer(batch.hicache_consumer_index) + if get_exec().features.enable_encoder_swa_bounded_replay: + from sglang.srt.model_executor.encoder_swa_replay import ( + run_encoder_swa_replay, + ) + + # Replay reads restored main/indexer KV before the normal extend. + run_encoder_swa_replay(self, batch) + forward_batch = ForwardBatch.init_new( batch, self.model_runner, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 27713ddc5..22dd0729a 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -194,6 +194,19 @@ def _evict_until_allocatable( return +def dsv41_dspark_needs_rebootstrap( + token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator, +) -> bool: + """V4.1's request-scoped pair ring and draft KV cannot use CPU tensor backup.""" + if str(get_spec().speculative_algorithm).upper() != "DSPARK": + return False + + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + pool = token_to_kv_pool_allocator.get_kvcache() + return isinstance(pool, DeepSeekV4TokenToKVPool) and 2 in pool.compression_ratios + + def retraction_backup( req: Req, tree_cache: BasePrefixCache, @@ -203,6 +216,11 @@ def retraction_backup( ) -> bool: """Returns False when the host pool cannot hold the backup; the caller aborts the request since its KV cannot be preserved.""" + if dsv41_dspark_needs_rebootstrap(token_to_kv_pool_allocator): + # Drain the in-flight verify before its slots can receive recomputed KV. + device = token_to_kv_pool_allocator.get_kvcache().device + torch.get_device_module(device).synchronize(device) + return True if backend == "cpu_tensor": req.offload_kv_cache(req_to_token_pool, token_to_kv_pool_allocator) return True diff --git a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py index dc4375b05..ce27f6fd9 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_compress_state.py @@ -183,13 +183,8 @@ class CompressStatePool: dtype=dtype, device=device, enable_memory_saver=enable_memory_saver ) if not online: - if _is_hip and ratio == 128: - # Request-scoped C128 state is addressed by req_pool_idx (or a - # per-request ring). The pool is allocated with torch.empty(), - # so a cold server can otherwise read uninitialized partial - # states before a request slot has been written for the first - # time. Initialize all C128 rows to the empty-state sentinel; - # C4 keeps the historical last-row sentinel behavior. + if ratio == 2 or (_is_hip and ratio == 128): + # Request-scoped rings reset all rows; C4 only its -1 sentinel row. self.kv_score_buffer.clear() else: self.kv_score_buffer[-1].clear() @@ -197,6 +192,11 @@ class CompressStatePool: def transfer_indices(self, req_pool_idx: int, seq_len: int) -> np.ndarray: """PD transfer indices of this pool's state for one request.""" assert self.request_scoped, "page-scoped state travels with the SWA pages" + if self.ratio == 2: + # Only an odd prefix leaves a pending half-pair for decode to read. + if seq_len % 2 == 0: + return np.empty((0,), dtype=np.int32) + return np.array([int(req_pool_idx)], dtype=np.int32) return request_scoped_state_transfer_indices( req_pool_idx, seq_len, @@ -262,15 +262,16 @@ class CompressStatePool: ) -> torch.Tensor: swa_pages = swa_loc // self.swa_page_size state_loc = swa_pages * self.ring_size + (swa_loc % self.ring_size) - state_loc = torch.where(swa_loc < 0, -1, state_loc) - return state_loc + # Not where(cond, -1, x): its scalar overload may stage a host tensor, + # which a CUDA graph capture cannot run. + return state_loc.masked_fill_(swa_loc < 0, -1) def translate_from_req_position_to_state_loc( self, req_pool_indices: torch.Tensor, positions: torch.Tensor ) -> torch.Tensor: state_loc = req_pool_indices * self.ring_size + positions % self.ring_size - state_loc = torch.where(positions < 0, -1, state_loc) - return state_loc + # A negative position means "no slot"; it lands on the empty row -1. + return state_loc.masked_fill_(positions < 0, -1) def get_state_by_state_loc(self, state_loc: torch.Tensor) -> KVAndScore: return self.kv_score_buffer[state_loc] diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index c9412c50b..388bd1578 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from contextlib import nullcontext -from typing import List, NamedTuple, Optional, Sequence, Tuple +from typing import List, Literal, NamedTuple, Optional, Sequence, Tuple, Union import torch @@ -16,7 +16,10 @@ from sglang.kernels.ops.attention.dsv4 import ( index_buf_accessor as dsv4_index_buf_accessor, ) from sglang.kernels.ops.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack -from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout +from sglang.kernels.ops.attention.dsv4.kv_layout import ( + KVLayout, + is_valid_kv_layout_pair, +) from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import layout from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs @@ -41,11 +44,17 @@ def get_dsv4_indexer_bytes_per_token(index_head_dim: int, use_fp4_indexer: bool) def get_compress_state_ring_size( - compress_ratio: int, is_speculative: bool = False + compress_ratio: int, is_speculative: bool = False, num_draft_tokens: int = 0 ) -> int: - assert compress_ratio in [4, 128], f"Unsupported {compress_ratio = }" - # Online C128 stores one (max, sum, kv) state per index; - # speculative decoding requires the experimental online C128 MTP path. + assert compress_ratio in [2, 4, 128], f"Unsupported {compress_ratio = }" + if compress_ratio == 2: + # Two positions are one pair, addressed by position % ring_size; a + # speculative ring must be wider than the draft window: pow2 >= 2 + drafts. + if not is_speculative: + return 2 + return 1 << (num_draft_tokens + 1).bit_length() + # Online c128 keeps one (max, sum, kv) state per index instead of a 128-slot + # ring of raw tokens, so ring_size collapses to 1. if compress_ratio == 128 and ONLINE_C128: if is_speculative and not envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get(): raise AssertionError("online c128 does not support MTP") @@ -57,8 +66,8 @@ def get_compress_state_ring_size( def get_compress_state_write_pad(compress_ratio: int, ring_size: int) -> int: - # Draft-token capacity must match mtp_pad in c_plan.cuh; - # a non-speculative ring has no write padding. + """Largest draft-token count this ring can serve; mirrors `mtp_pad` in + `c_plan.cuh`, where the bound is derived.""" window_size = compress_ratio * (2 if compress_ratio == 4 else 1) return ring_size - window_size + 2 if ring_size > window_size else 0 @@ -69,6 +78,69 @@ def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int: return sliding_window + spec_extra +def resolve_compressed_kv_layout( + kv_layout: KVLayout, compress_ratio: int, option: Optional[str] = None +) -> KVLayout: + """Layout of one compress ratio's cache next to a ``kv_layout`` main cache. + The ratio-1/2 latents are already e2m1 with per-16 e4m3 scales, so ``V41_FP4`` + is lossless for them; ratios 4 / 128 are not fp4-rounded and stay fp8.""" + if option is not None: + option = option.lower() + assert option in ( + "auto", + "fp8", + "fp4", + ), f"unknown compressed KV layout {option!r}" + if option == "auto": + option = None + if kv_layout is KVLayout.V4: + assert option in (None, "fp8"), "the V4 main cache only pairs with V4 caches" + return KVLayout.V4 + assert kv_layout is KVLayout.V41, f"{kv_layout} is not a main-cache layout" + if option == "fp8": + return KVLayout.V41 + if option == "fp4": + return KVLayout.V41_FP4 + return KVLayout.V41_FP4 if compress_ratio in (1, 2) else KVLayout.V41 + + +def flashmla_supports_v41_kv_layouts() -> bool: + """Whether the installed FlashMLA decode kernel reads the V41 / V41_FP4 + formats; its docstring lists the bytes-per-token it detects.""" + try: + from sgl_kernel.flash_mla import flash_mla_with_kvcache + except Exception: + return False + return "528" in (flash_mla_with_kvcache.__doc__ or "") + + +def select_dsv4_kv_layout() -> Tuple[KVLayout, Optional[str]]: + """The (main-cache layout, compressed-cache option) for a new DeepSeek-V4 + family pool; the V4.1 layouts exist only in SM100 / SM103 FlashMLA.""" + mode = envs.SGLANG_DSV4_KV_LAYOUT.get().lower() + option = envs.SGLANG_DSV4_COMPRESSED_KV_LAYOUT.get().lower() + if mode == "v4": + return KVLayout.V4, None if option == "auto" else option + assert mode in ("v41", "auto"), f"unknown SGLANG_DSV4_KV_LAYOUT={mode!r}" + is_sm100 = ( + torch.cuda.is_available() + and torch.version.cuda is not None + and torch.cuda.get_device_capability()[0] == 10 + ) + supported = flashmla_supports_v41_kv_layouts() + if mode == "auto": + if is_sm100 and supported: + return KVLayout.V41, option + return KVLayout.V4, None + assert is_sm100, "the V4.1 KV cache layouts need an SM100 / SM103 GPU" + if not supported: + logger.warning( + "SGLANG_DSV4_KV_LAYOUT=v41 but the installed FlashMLA does not advertise " + "the V4.1 KV cache formats; the attention kernel will reject the cache." + ) + return KVLayout.V41, option + + class DeepSeekV4SingleKVPool(KVCache): # Paged FlashMLA main-KV format of this pool's rows. kv_layout: KVLayout = KVLayout.V4 @@ -85,6 +157,7 @@ class DeepSeekV4SingleKVPool(KVCache): enable_memory_saver: bool, start_layer: Optional[int] = None, end_layer: Optional[int] = None, + kv_layout: Union[str, KVLayout] = KVLayout.V4, ): super().__init__( size, @@ -99,8 +172,11 @@ class DeepSeekV4SingleKVPool(KVCache): self.qk_nope_head_dim = qk_nope_head_dim self.qk_rope_head_dim = qk_rope_head_dim + # Paged FlashMLA layout of this pool's pages; see KVLayout. + self.kv_layout = KVLayout.parse(kv_layout) self.scale_pad = 1 - self.quantize_block_size = 64 + self.quantize_block_size = self.kv_layout.tile_size + # V4 keeps its 64 RoPE dims in bf16; the V4.1 layouts quantize them too. self.rope_storage_dtype = torch.bfloat16 self.k_with_scale_buffer_dtype = torch.int8 self._create_buffers() @@ -120,6 +196,9 @@ class DeepSeekV4SingleKVPool(KVCache): ] def get_bytes_per_token(self) -> int: + if self.kv_layout is not KVLayout.V4: + assert self.qk_nope_head_dim + self.qk_rope_head_dim == 512 + return self.kv_layout.bytes_per_token dim_per_token = ( self.qk_nope_head_dim + self.qk_rope_head_dim * self.rope_storage_dtype.itemsize @@ -131,13 +210,17 @@ class DeepSeekV4SingleKVPool(KVCache): def create_buffer(self, *, num_pages: int): bytes_per_token = self.get_bytes_per_token() self.kv_cache_total_dim = bytes_per_token - bytes_per_page_non_padded = self.page_size * bytes_per_token - self.bytes_per_page_padded = ceil_div(bytes_per_page_non_padded, 576) * 576 + self.bytes_per_page_padded = self.kv_layout.page_bytes(self.page_size) - assert bytes_per_token == 448 + 64 * 2 + 8, ( - "DSV4 KV layout: qk_nope_head_dim FP8 (448) + qk_rope_head_dim BF16 " - "(64*2) + nope FP8 scales + scale_pad = 584 bytes/token" - ) + if self.kv_layout is KVLayout.V4: + assert bytes_per_token == 448 + 64 * 2 + 8, ( + "DSV4 KV layout: qk_nope_head_dim FP8 (448) + qk_rope_head_dim BF16 " + "(64*2) + nope FP8 scales + scale_pad = 584 bytes/token" + ) + assert ( + self.bytes_per_page_padded + == ceil_div(self.page_size * bytes_per_token, 576) * 576 + ) assert self.store_dtype == torch.uint8 return torch.zeros( @@ -153,6 +236,10 @@ class DeepSeekV4SingleKVPool(KVCache): loc: torch.Tensor, cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack, ): + assert self.kv_layout is KVLayout.V4, ( + "the (fp8 nope, bf16 rope, 7 scales) pack is the V4 layout; " + f"a {self.kv_layout.value} pool is written through set_key_buffer_fused" + ) dsv4_index_buf_accessor.SetKAndS.execute( pool=self, buf=self.kv_buffer[layer_id], @@ -165,13 +252,19 @@ class DeepSeekV4SingleKVPool(KVCache): layer_id: int, loc: torch.Tensor, cache_k: torch.Tensor, + freqs_cis: Optional[torch.Tensor] = None, ) -> None: + """Quantize ``cache_k`` ``[n, 512]`` bf16 into this pool's layout at ``loc``. + ``freqs_cis`` (V4.1 only) rotates the RoPE tail in-kernel, so the input is + the un-rotated latent and the fp4 / fp8 rounding happens once.""" return fused_store_cache( input=cache_k, cache=self.kv_buffer[layer_id], indices=loc, page_size=self.page_size, type="flashmla", + layout=self.kv_layout, + freqs_cis=freqs_cis, ) def get_key_buffer(self, layer_id: int): @@ -235,12 +328,14 @@ class DeepSeekV4UniformFP8KVPool(DeepSeekV4SingleKVPool): layer_id: int, loc: torch.Tensor, cache_k: torch.Tensor, + freqs_cis: Optional[torch.Tensor] = None, ) -> None: """Store normed/roped rows as e4m3 with the backend's fixed unit scale. uint8 views work around index_put not supporting FP8 dtypes. """ + assert freqs_cis is None, "the uniform-FP8 pool takes finished (rotated) rows" assert cache_k.dim() == 2 and cache_k.shape[1] == self.kv_cache_total_dim self.kv_buffer[layer_id].view(torch.uint8).view(-1, self.kv_cache_total_dim)[ loc.long() @@ -260,6 +355,7 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool): enable_memory_saver: bool, start_layer: int | None = None, end_layer: int | None = None, + kv_layout: Union[str, KVLayout] = KVLayout.V4, ): super().__init__( size, @@ -272,6 +368,11 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool): enable_memory_saver, start_layer, end_layer, + kv_layout=kv_layout, + ) + # The HiSparse transfer kernels hardcode the V4 token layout. + assert self.kv_layout is KVLayout.V4, ( + f"HiSparse C4 pools support the V4 layout only, got {self.kv_layout}" ) self.data_ptrs = torch.tensor( @@ -318,9 +419,10 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool): layer_id: int, loc: torch.Tensor, cache_k: torch.Tensor, + freqs_cis: Optional[torch.Tensor] = None, ) -> None: loc = self.translate_loc_to_hisparse_device(loc) - return super().set_key_buffer_fused(layer_id, loc, cache_k) + return super().set_key_buffer_fused(layer_id, loc, cache_k, freqs_cis) def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None): raise NotImplementedError("HiSparseC4DevicePool does not support get_cpu_copy") @@ -331,6 +433,21 @@ class HiSparseC4DevicePool(DeepSeekV4SingleKVPool): raise NotImplementedError("HiSparseC4DevicePool does not support load_cpu_copy") +# Low-ratio indexer-K pool page, in compressed slots: the DeepGEMM indexer reads +# K in blocks of at most 128 and sglang's JIT metadata builder asserts 64. +def dsv41_index_page_size() -> int: + from sglang.srt.layers.deep_gemm_wrapper.configurer import ( + DEEPGEMM_PAGED_SPARSE_MQA_LOGITS, + ) + + if DEEPGEMM_PAGED_SPARSE_MQA_LOGITS: + return 128 + return 64 + + +DSV41_INDEX_PAGE_SIZE = dsv41_index_page_size() + + class DeepSeekV4IndexerPool(KVCache): quant_block_size = 128 index_k_with_scale_buffer_dtype = torch.uint8 @@ -346,6 +463,7 @@ class DeepSeekV4IndexerPool(KVCache): enable_memory_saver: bool, start_layer: Optional[int] = None, end_layer: Optional[int] = None, + use_fp4_indexer: Optional[bool] = None, ): super().__init__( size, @@ -358,8 +476,12 @@ class DeepSeekV4IndexerPool(KVCache): end_layer, ) self.index_head_dim = index_head_dim - self.use_fp4_indexer = get_exec().kernel.enable_deepseek_v4_fp4_indexer + if use_fp4_indexer is None: + use_fp4_indexer = get_exec().kernel.enable_deepseek_v4_fp4_indexer + self.use_fp4_indexer = use_fp4_indexer self.uses_aiter_fp4_layout = _is_hip and self.use_fp4_indexer + # Low-ratio pools round to nearest even; c4 keeps threshold rounding. + self.index_k_rne = False self._create_buffer() @@ -500,8 +622,54 @@ class DeepSeekV4IndexerPool(KVCache): cache=self.index_k_with_scale_buffer[layer_id - self.start_layer], loc=loc, page_size=self.page_size, + rne=self.index_k_rne, ) + def get_index_k_fp4( + self, layer_id: int, slots: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Packed fp4 rows at `slots`: (payload int8 [n, 64], scales int32 [n]), + from the page layout [page_size * 64 payload | page_size * 4 scale].""" + assert self.use_fp4_indexer, "packed readback only applies to the fp4 layout" + buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + slots = slots.to(torch.int64) + p = self.page_size + page, off = (slots // p).unsqueeze(-1), slots % p + payload_cols = (off * 64).unsqueeze(-1) + torch.arange(64, device=buf.device) + scale_cols = (p * 64 + off * 4).unsqueeze(-1) + torch.arange( + 4, device=buf.device + ) + payload = buf[page, payload_cols].view(torch.int8) # [n, 64] + scales = buf[page, scale_cols].contiguous().view(torch.int32).squeeze(-1) + return payload, scales + + def get_index_k_dequant( + self, layer_id: int, slots: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """Dequantized bf16 [n, index_head_dim] index K; `slots` None reads the pool.""" + from sglang.srt.layers.quantization.fp8 import DSV4_DEQUANT_FP4_TABLE + + assert self.use_fp4_indexer, "dequant readback only applies to the fp4 layout" + buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] + if slots is None: + slots = torch.arange(self.size, device=buf.device) + slots = slots.to(torch.int64) + # Page layout: see get_index_k_fp4. + p = self.page_size + page, off = (slots // p).unsqueeze(-1), slots % p + payload_cols = (off * 64).unsqueeze(-1) + torch.arange(64, device=buf.device) + scale_cols = (p * 64 + off * 4).unsqueeze(-1) + torch.arange( + 4, device=buf.device + ) + u = buf[page, payload_cols].view(torch.uint8) # [n, 64] + codes = torch.stack([u & 0x0F, (u >> 4) & 0x0F], dim=-1) # [n, 64, 2] + vals = DSV4_DEQUANT_FP4_TABLE.to(buf.device)[codes.long()].flatten( + 1 + ) # [n, 128] + exps = buf[page, scale_cols].to(torch.int32) & 0xFF # [n, 4] + scales = torch.exp2(exps.float() - 127).repeat_interleave(32, dim=-1) + return (vals * scales).to(torch.bfloat16) + class _CompressedPoolConfig(NamedTuple): kv_size: int @@ -511,7 +679,9 @@ class _CompressedPoolConfig(NamedTuple): class DeepSeekV4LayerItem(NamedTuple): - compress_ratio: int + compress_ratio: Literal[0, 1, 2, 4, 128] + # Layer index inside compress_kv_pool. Ratios 1/2 share a pool layer across the + # kv_source layer that writes it and the layers that read it. compress_layer_id: int compress_kv_pool: Optional[DeepSeekV4SingleKVPool] = None @@ -683,6 +853,11 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_hisparse: bool = False, online_mtp_max_draft_tokens: int = 0, num_req_slots: Optional[int] = None, + kv_source_layers: Sequence[int] = (), + full_size: Optional[int] = None, + is_draft_worker: bool = False, + kv_layout: Union[str, KVLayout] = KVLayout.V4, + compressed_kv_layout: Optional[str] = None, ): super().__init__( swa_size, @@ -694,6 +869,14 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): start_layer, end_layer, ) + # Layout of the SWA (main) cache; compressed caches follow + # resolve_compressed_kv_layout, so valid (main, extra) pairs form only here. + self.kv_layout = KVLayout.parse(kv_layout) + assert self.kv_layout in ( + KVLayout.V4, + KVLayout.V41, + ), f"{self.kv_layout} is only valid for a compressed (extra) cache" + self.compressed_kv_layout_option = compressed_kv_layout c4_logical_size = c128_size * 32 logger.info( @@ -723,6 +906,11 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.uniform_fp8 = ( not self._unified_kv ) and get_exec().kernel.dsv4_attn_backend == "trtllm" + if self.uniform_fp8: + assert self.kv_layout is KVLayout.V4, ( + "--dsv4-attn-backend trtllm keeps its own uniform 512-byte pages; " + f"it cannot be combined with SGLANG_DSV4_KV_LAYOUT={self.kv_layout.value}" + ) c4_ring_size = self.get_ring_size(4) if self._unified_kv: # Unified C4 state is request-addressed: one ring per req slot, @@ -738,18 +926,26 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): c128_state_pool_size = max( c128_state_pool_size, self.num_req_slots * c128_ring_size ) + # Only the ratios the model has anywhere get a pool config: the backend, PD + # state transfer and HiCache read the registries as "the ratios this model + # has", and a PP stage missing one keeps its empty pool so the PD wire aligns. + model_ratios = set(compression_ratios) self.compressed_pool_configs = { - 4: _CompressedPoolConfig( - kv_size=c4_size, - state_size=c4_state_pool_size, - state_dtype=c4_state_dtype, - indexer_size=c4_logical_size, - ), - 128: _CompressedPoolConfig( - kv_size=c128_size, - state_size=c128_state_pool_size, - state_dtype=c128_state_dtype, - ), + ratio: config + for ratio, config in { + 4: _CompressedPoolConfig( + kv_size=c4_size, + state_size=c4_state_pool_size, + state_dtype=c4_state_dtype, + indexer_size=c4_logical_size, + ), + 128: _CompressedPoolConfig( + kv_size=c128_size, + state_size=c128_state_pool_size, + state_dtype=c128_state_dtype, + ), + }.items() + if ratio in model_ratios } self.compression_ratios = compression_ratios self.online_mtp_max_draft_tokens = online_mtp_max_draft_tokens @@ -786,7 +982,51 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): stage_layer_num = len(stage_ratios) kv_pool_cls: type = DeepSeekV4SingleKVPool - if self._unified_kv: + self.request_window = None + encoder_replay = get_exec().features.enable_encoder_swa_bounded_replay + # DSpark's draft shares the target's full-to-SWA mapping, so the target + # keeps its paged SWA allocator even under encoder replay. + self.needs_paged_swa_allocator = ( + not encoder_replay + or is_draft_worker + or get_spec().speculative_algorithm is not None + ) + if encoder_replay and not is_draft_worker: + from sglang.srt.mem_cache.dsv41_request_window import RequestWindow + + def make_window_pool(size, layers): + return self._make_kv_pool( + size=size, + page_size=swa_page_size, + dtype=dtype, + layer_num=layers, + device=device, + enable_memory_saver=enable_memory_saver, + global_page_size=swa_page_size, + kv_layout=self.kv_layout, + ) + + self.swa_kv_pool = None + self.unified_kv_pool = None + from sglang.srt.runtime_context import get_schedule + + chunk = get_schedule().chunked_prefill_size or 0 + self.request_window = RequestWindow( + make_window_pool, + num_slots=self.num_req_slots, + layers=stage_layer_num, + page_size=swa_page_size, + capacity=self.sliding_window + (online_mtp_max_draft_tokens or 0), + workspace_rows=(self.num_req_slots + 1) * self.sliding_window + + max( + chunk, + (self.num_req_slots + 1) * (1 + (online_mtp_max_draft_tokens or 0)), + ), + ) + elif self._unified_kv: + assert self.kv_layout is KVLayout.V4, ( + "unified_kv keeps bf16 rows, not a paged FlashMLA layout" + ) self.swa_kv_pool = None swa_ring_size = get_swa_ring_size( self.sliding_window, get_spec().speculative_algorithm is not None @@ -825,9 +1065,19 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): device=device, enable_memory_saver=enable_memory_saver, global_page_size=swa_page_size, + kv_layout=self.kv_layout, cls=kv_pool_cls, ) + logger.info( + "DSV4 SWA storage: worker=%s, storage=%s, paged_allocator=%s", + "draft" if is_draft_worker else "target", + "request_window" if self.request_window is not None else "paged", + self.needs_paged_swa_allocator, + ) + self.full_size = full_size + self.kv_source_layers = list(kv_source_layers) + self.sources_by_ratio = self._collect_sources_by_ratio() self._init_compressed_pools( stage_ratios=stage_ratios, page_size=page_size, @@ -863,8 +1113,12 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.full_to_swa_index_mapping = full_to_swa_index_mapping def get_ring_size(self, compress_ratio: int) -> int: - is_speculative = get_spec().speculative_algorithm is not None - return get_compress_state_ring_size(compress_ratio, is_speculative) + spec = get_spec() + return get_compress_state_ring_size( + compress_ratio, + spec.speculative_algorithm is not None, + spec.speculative_num_draft_tokens or 0, + ) def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): assert self.full_to_swa_index_mapping is not None @@ -912,14 +1166,31 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): data_ptrs.append(buf.data_ptr() + swa_pages * row_bytes) data_lens.append(compress_rows * row_bytes) item_lens.append(rows_per_page * row_bytes) - else: + elif kv_pool is not None: for buf in kv_pool.kv_buffer: append_page_buffer(buf) indexer_pool = self.index_pools.get(ratio) - if indexer_pool is not None: - for buf in indexer_pool.contiguous_page_row_buffers(): - append_page_buffer(buf) + if indexer_pool is None: + continue + # The transfer addresses every buffer by FULL page id; ratio-1/2 index + # pools page at DSV41_INDEX_PAGE_SIZE, so one item is the run of index + # pages holding a FULL page's page_size // ratio slots. + index_pages_per_full_page = 1 + if ratio in (1, 2): + slots_per_full_page = self.page_size // ratio + assert slots_per_full_page % indexer_pool.page_size == 0, ( + f"ratio-{ratio} index pages of {indexer_pool.page_size} slots do not " + f"tile a FULL page of {slots_per_full_page} slots" + ) + index_pages_per_full_page = ( + slots_per_full_page // indexer_pool.page_size + ) + for buf in indexer_pool.contiguous_page_row_buffers(): + assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D" + data_ptrs.append(buf.data_ptr()) + data_lens.append(buf.nbytes) + item_lens.append(buf[0].nbytes * index_pages_per_full_page) return data_ptrs, data_lens, item_lens @@ -1011,7 +1282,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): data_lens: List[int] = [] item_lens: List[int] = [] - if not self._unified_kv: + if self.swa_kv_pool is not None: for buf in self.swa_kv_pool.kv_buffer: assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D" data_ptrs.append(buf.data_ptr()) @@ -1023,6 +1294,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.indexer_compress_state_pools, ]: for pool in pools: + # Request-scoped state ships as C128_STATE, not with the SWA ring. if pool is None or pool.request_scoped: continue t = pool.kv_score_buffer.kv_score @@ -1036,6 +1308,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): def get_request_state_buf_infos( self, ) -> Tuple[List[int], List[int], List[int]]: + """Request-scoped state: the c128 raw-token ring (or its single online row) + and the ratio-2 pending-pair ring. One item is one c128 page / pair ring.""" data_ptrs: List[int] = [] data_lens: List[int] = [] item_lens: List[int] = [] @@ -1046,7 +1320,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): assert t.ndim == 2, f"expected 2D buffer, got {t.ndim}D" data_ptrs.append(t.data_ptr()) data_lens.append(t.nbytes) - item_lens.append(t[0].nbytes if ONLINE_C128 else t[0].nbytes * 128) + if pool.ratio == 2: + item_lens.append(t[0].nbytes * pool.ring_size) + else: + item_lens.append(t[0].nbytes if ONLINE_C128 else t[0].nbytes * 128) return data_ptrs, data_lens, item_lens def _init_compressed_pools( @@ -1060,12 +1337,25 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_hisparse: bool, kv_pool_cls: type, ) -> None: + """One KV pool (plus packed indexer-K pool) per compress ratio in this stage: + slot = full-pool loc // ratio, page = page_size // ratio, so pages line up.""" configs = self.compressed_pool_configs layer_counts = {ratio: stage_ratios.count(ratio) for ratio in configs} # Keep empty pools and allocation order for PP stages without a given ratio. self.kv_pools: dict[int, Optional[DeepSeekV4SingleKVPool]] = { ratio: None for ratio in configs } + # The PD wire order stays C4, C128, then the ratio-1/2 kv_source layers. + low_ratio_sources = { + ratio: sources + for ratio, sources in getattr(self, "sources_by_ratio", {}).items() + if ratio in (1, 2) + } + if low_ratio_sources: + assert self.full_size is not None, ( + "low compress ratios need the full pool size" + ) + assert not self._unified_kv, "unified_kv has no low compress ratio layout" if not self._unified_kv: for ratio, config in configs.items(): @@ -1084,6 +1374,19 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_memory_saver=enable_memory_saver, global_page_size=page_size, cls=pool_cls, + kv_layout=self.compressed_kv_layout(ratio), + ) + for ratio, sources in low_ratio_sources.items(): + self.kv_pools[ratio] = self._make_kv_pool( + size=self.full_size // ratio, + page_size=page_size // ratio, + dtype=dtype, + layer_num=len(sources), + device=device, + enable_memory_saver=enable_memory_saver, + global_page_size=page_size, + cls=kv_pool_cls, + kv_layout=self.compressed_kv_layout(ratio), ) self.index_pools: dict[int, DeepSeekV4IndexerPool] = { @@ -1099,11 +1402,24 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): for ratio, config in configs.items() if config.indexer_size is not None } + for ratio, sources in low_ratio_sources.items(): + # Reserved FULL page 0 pushes real slots past full_size, and one index + # padding page is too small to cover that gap. + self.index_pools[ratio] = self._make_indexer_pool( + (self.full_size + page_size) // ratio, + DSV41_INDEX_PAGE_SIZE, + dtype, + self.indexer_head_dim, + len(sources), + device, + enable_memory_saver, + force_fp4=True, + ) - # HiCache and hardware backends still access the per-ratio attributes. - self.c4_kv_pool = self.kv_pools[4] - self.c128_kv_pool = self.kv_pools[128] - self.c4_indexer_kv_pool = self.index_pools[4] + # HiCache and hardware backends still read these per-ratio attributes. + self.c4_kv_pool = self.kv_pools.get(4) + self.c128_kv_pool = self.kv_pools.get(128) + self.c4_indexer_kv_pool = self.index_pools.get(4) def _make_kv_pool( self, @@ -1116,6 +1432,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_memory_saver: bool, global_page_size: int, cls: type = DeepSeekV4SingleKVPool, + kv_layout: KVLayout = KVLayout.V4, ) -> DeepSeekV4SingleKVPool: """Build a full / SWA / c4 / c128 single-KV pool. ``global_page_size`` is the model-wide page_size (== ``page_size`` for the SWA pool, larger @@ -1132,8 +1449,17 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): layer_num, device, enable_memory_saver, + kv_layout=kv_layout, ) + def compressed_kv_layout(self, compress_ratio: int) -> KVLayout: + """See :func:`resolve_compressed_kv_layout`.""" + layout = resolve_compressed_kv_layout( + self.kv_layout, compress_ratio, self.compressed_kv_layout_option + ) + assert is_valid_kv_layout_pair(self.kv_layout, layout) + return layout + def _make_indexer_pool( self, size: int, @@ -1143,10 +1469,25 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): layer_num: int, device: str, enable_memory_saver: bool, + force_fp4: bool = False, ) -> DeepSeekV4IndexerPool: """Build the c4 lightning-indexer K pool (packed CUDA layout). Overridden by :class:`DSV4NPUTokenToKVPool` to swap in the - dedicated-buffer NPU variant (int8 K + fp16 scale).""" + dedicated-buffer NPU variant. ``force_fp4`` forces the fp4 low-ratio layout.""" + if force_fp4: + pool = DeepSeekV4IndexerPool( + size, + page_size, + dtype, + index_head_dim, + layer_num, + device, + enable_memory_saver, + use_fp4_indexer=True, + ) + # The dsv41 low-ratio indexer rounds to nearest even (reference rounding). + pool.index_k_rne = True + return pool return DeepSeekV4IndexerPool( size, page_size, @@ -1172,13 +1513,30 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_memory_saver=enable_memory_saver, ratio=ratio, online=(ratio == 128 and ONLINE_C128), - request_scoped=ratio == 128, + request_scoped=ratio in (2, 128), swa_page_size=self.swa_page_size, online_mtp_max_draft_tokens=( self.online_mtp_max_draft_tokens if ratio == 128 else 0 ), ) + def _make_pair_state_pool(self, enable_memory_saver: bool) -> CompressStatePool: + """Ratio-2 pending-pair state: one position ring per request slot, holding + the fp32 (kv, score) of an even token until its odd partner arrives.""" + ring_size = self.get_ring_size(2) + return CompressStatePool( + size=self.num_req_slots * ring_size, + ring_size=ring_size, + overlap=False, + head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim, + dtype=torch.float32, + device=self.device, + enable_memory_saver=enable_memory_saver, + ratio=2, + request_scoped=True, + online=False, + ) + def _init_paged_compress_states(self, enable_memory_saver: bool): total_L = len(self.compression_ratios) self.compress_state_pools: List[Optional[CompressStatePool]] = [None] * total_L @@ -1188,7 +1546,15 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): for idx in range(self._stage_start, self._stage_end): ratio = self.compression_ratios[idx] - if ratio == 0: + if ratio in (0, 1): + continue + + if ratio == 2: + # Only a kv_source layer compresses; later ratio-2 layers read it. + if idx in self.sources_by_ratio.get(2, []): + self.compress_state_pools[idx] = self._make_pair_state_pool( + enable_memory_saver + ) continue self.compress_state_pools[idx] = self._make_compress_state_pool( @@ -1204,6 +1570,36 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_memory_saver=enable_memory_saver, ) + def _collect_sources_by_ratio(self) -> dict[int, List[int]]: + """Layers owning compressed storage: all of ratios 4/128, kv_sources of 1/2.""" + stage = range(self._stage_start, self._stage_end) + for idx in stage: + ratio = self.compression_ratios[idx] + if ratio not in (0, 1, 2, 4, 128): + raise ValueError(f"Unsupported compression ratio: {ratio}") + + sources_by_ratio: dict[int, List[int]] = {} + for ratio in (4, 128, 1, 2): + if ratio in (1, 2): + layers = [ + l + for l in self.kv_source_layers + if l in stage and self.compression_ratios[l] == ratio + ] + else: + layers = [l for l in stage if self.compression_ratios[l] == ratio] + if layers: + sources_by_ratio[ratio] = layers + return sources_by_ratio + + def source_layer_of(self, layer_id: int) -> int: + """The layer owning this layer's compressed storage: itself for ratios 4/128, + the nearest preceding kv_source layer for ratios 1/2.""" + ratio = self.compression_ratios[layer_id] + sources = [l for l in self.sources_by_ratio[ratio] if l <= layer_id] + assert sources, f"layer {layer_id} (ratio {ratio}) has no kv_source layer" + return max(sources) + def _init_compressed_layer_mapping(self): layer_counts = {0: 0, **{ratio: 0 for ratio in self.kv_pools}} total_L = len(self.compression_ratios) @@ -1213,12 +1609,17 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): ratio = self.compression_ratios[idx] if ratio not in layer_counts: raise ValueError(f"Unsupported compression ratio: {ratio}") + if ratio in (1, 2): + sources = self.sources_by_ratio[ratio] + compress_layer_id = sources.index(self.source_layer_of(idx)) + else: + compress_layer_id = layer_counts[ratio] + layer_counts[ratio] += 1 self.layer_mapping[idx] = DeepSeekV4LayerItem( compress_ratio=ratio, - compress_layer_id=layer_counts[ratio], + compress_layer_id=compress_layer_id, compress_kv_pool=self.kv_pools.get(ratio), ) - layer_counts[ratio] += 1 def wait_layer_transfer(self, layer_id: int) -> None: if self.layer_transfer_counter is not None: @@ -1228,7 +1629,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.wait_layer_transfer(layer_id) compress_state_pool = self.compress_state_pools[layer_id] assert compress_state_pool is not None, ( - "Only c4/c128 layers have attention states." + "Only c4/c128 layers and ratio-2 kv_source layers have attention states." ) return compress_state_pool @@ -1292,23 +1693,20 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): return pools[0].transfer_indices(req_pool_idx, seq_len) def clear_request_scoped_state(self, req_pool_idx: int) -> None: - """Reset request-scoped state for one req slot.""" + """Reset one req slot's C128 ring and ratio-2 pending-pair state.""" for pool in self.compress_state_pools: if pool is None or not pool.request_scoped: continue - state = pool.kv_score_buffer.kv_score - if ONLINE_C128: - row = state[req_pool_idx] + if pool.ratio == 128 and ONLINE_C128: + row = pool.kv_score_buffer.kv_score[req_pool_idx] head_dim = row.shape[-1] // 3 row[:head_dim].fill_(float("-inf")) row[head_dim:].zero_() - else: - start = req_pool_idx * pool.ring_size - rows = state[start : start + pool.ring_size] - half = rows.shape[-1] // 2 - rows[:, :half].zero_() - rows[:, half:].fill_(float("-inf")) + continue + + start = req_pool_idx * pool.ring_size + pool.kv_score_buffer[start : start + pool.ring_size].clear() def clear_unaccepted_c128_draft_states( self, @@ -1349,8 +1747,40 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): return layer_id - self._stage_start def get_swa_raw_buffer(self, layer_id: int) -> torch.Tensor: + if self.request_window is not None: + return self.request_window.buffer(self._swa_local_layer_id(layer_id)) return self.swa_kv_pool.kv_buffer[self._swa_local_layer_id(layer_id)] + def get_swa_key_buffer(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) + if self.request_window is not None: + return self.get_swa_raw_buffer(layer_id).view( + self.request_window.state.dtype + ) + return self.swa_kv_pool.get_key_buffer(self._swa_local_layer_id(layer_id)) + + def set_swa_key_buffer( + self, + layer_id: int, + loc: torch.Tensor, + cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack, + ) -> None: + assert self.kv_layout is KVLayout.V4, ( + "the (fp8 nope, bf16 rope, 7 scales) pack is the V4 layout; " + f"a {self.kv_layout.value} pool is written through the fused setters" + ) + if self.request_window is not None: + dsv4_index_buf_accessor.SetKAndS.execute( + pool=self.request_window.state, + buf=self.get_swa_raw_buffer(layer_id), + loc=loc, + nope_fp8_rope_bf16_pack=cache_nope_fp8_rope_bf16_pack, + ) + else: + self.swa_kv_pool.set_key_buffer( + self._swa_local_layer_id(layer_id), loc, cache_nope_fp8_rope_bf16_pack + ) + def get_extra_key_page_size(self, layer_id: int) -> int: _, _, compress_kv_pool = self.layer_mapping[layer_id] assert compress_kv_pool is not None @@ -1369,12 +1799,16 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): return compress_kv_pool.kv_cache_total_dim def get_swa_key_layout(self) -> KVLayout: - return self.swa_kv_pool.kv_layout + # swa_kv_pool is None under the request window and unified_kv. + return self.kv_layout def get_swa_key_bytes_per_token(self) -> int: """Last dim of the ``(pages, page_size, 1, bytes)`` view the attention kernel detects the SWA cache's format from.""" - return self.swa_kv_pool.kv_cache_total_dim + if self.uniform_fp8: + # The trtllm uniform-FP8 pool has no paged FlashMLA layout: 512 B/token. + return self.swa_kv_pool.kv_cache_total_dim + return self.kv_layout.bytes_per_token def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor | None: self.wait_layer_transfer(layer_id) @@ -1401,6 +1835,25 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): ) return pool + def get_low_ratio_index_k_dequant( + self, layer_id: int, slots: Optional[torch.Tensor] = None + ) -> torch.Tensor: + """Index-K rows at `slots` from the layer's latent source.""" + compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] + return self._indexer_pool(compress_ratio).get_index_k_dequant( + compress_layer_id, slots + ) + + def get_low_ratio_index_k_fp4( + self, layer_id: int, slots: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Packed fp4 index-K rows at `slots`: (payload int8 [n, 64], ue8m0 scales + packed int32 [n]), the input layout of quantize_fp4_indexer_tensor.""" + compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] + return self._indexer_pool(compress_ratio).get_index_k_fp4( + compress_layer_id, slots + ) + def get_index_k_page_size(self, compress_ratio: int = 4) -> int: return self._indexer_pool(compress_ratio).page_size @@ -1473,12 +1926,14 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): swa_loc: torch.Tensor, cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack, ) -> None: - self.swa_kv_pool.set_key_buffer( - self._swa_local_layer_id(layer_id), swa_loc, cache_nope_fp8_rope_bf16_pack - ) + self.set_swa_key_buffer(layer_id, swa_loc, cache_nope_fp8_rope_bf16_pack) def get_swa_key_buffer_radix(self, layer_id: int) -> torch.Tensor: self.wait_layer_transfer(layer_id) + if self.request_window is not None: + return self.get_swa_raw_buffer(layer_id).view( + self.request_window.state.dtype + ) return self.swa_kv_pool.get_key_buffer(self._swa_local_layer_id(layer_id)) def set_swa_key_buffer_radix_fused( @@ -1487,8 +1942,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): swa_loc: torch.Tensor, cache_k: torch.Tensor, ) -> None: - return self.swa_kv_pool.set_key_buffer_fused( - self._swa_local_layer_id(layer_id), swa_loc, cache_k + return fused_store_cache( + input=cache_k, + cache=self.get_swa_raw_buffer(layer_id), + indices=swa_loc, + page_size=self.swa_page_size, + type="flashmla", + layout=self.kv_layout, ) def set_swa_key_buffer_radix_fused_norm_rope( @@ -1528,8 +1988,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): freqs_cis=freqs_cis, positions=positions, out_loc=swa_loc, - kvcache=self.swa_kv_pool.kv_buffer[self._swa_local_layer_id(layer_id)], - page_size=self.swa_kv_pool.page_size, + kvcache=self.get_swa_raw_buffer(layer_id), + page_size=self.swa_page_size, + layout=self.kv_layout, ) def set_unified_key_buffer_radix_fused_norm_rope( @@ -1565,10 +2026,20 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): layer_id: int, loc: torch.Tensor, cache_k: torch.Tensor, + freqs_cis: Optional[torch.Tensor] = None, ) -> None: + """Write ``cache_k`` ``[n, 512]`` bf16 into the layer's compressed cache. + For an fp4 (``V41_FP4``) cache pass the *un-quantized* latent, plus + ``freqs_cis`` if it is not rotated yet: the kernel rounds to e2m1 once.""" _, compress_layer_id, compress_kv_pool = self.layer_mapping[layer_id] assert compress_kv_pool is not None - return compress_kv_pool.set_key_buffer_fused(compress_layer_id, loc, cache_k) + if freqs_cis is not None: + assert compress_kv_pool.kv_layout is KVLayout.V41_FP4, ( + "in-kernel RoPE is for the fp4 cache; fp8 caches take the finished value" + ) + return compress_kv_pool.set_key_buffer_fused( + compress_layer_id, loc, cache_k, freqs_cis + ) def set_index_k_fused( self, diff --git a/python/sglang/srt/mem_cache/dsv41_request_window.py b/python/sglang/srt/mem_cache/dsv41_request_window.py new file mode 100644 index 000000000..f336d6ae7 --- /dev/null +++ b/python/sglang/srt/mem_cache/dsv41_request_window.py @@ -0,0 +1,269 @@ +from typing import Optional + +import msgspec +import torch + +from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout +from sglang.srt.model_executor.runner_utils.capture_mode import get_is_capture_mode + + +class WindowLayout(msgspec.Struct, frozen=True): + req: torch.Tensor + pos: torch.Tensor + write_loc: torch.Tensor + indices: torch.Tensor + lengths: torch.Tensor + history_req: torch.Tensor + history_pos: torch.Tensor + history_loc: torch.Tensor + history_valid: torch.Tensor + commit_mask: torch.Tensor + size: int + + def copy_(self, other: "WindowLayout") -> None: + # Captured copy kernels read these tensors by address, so a graph replay + # must refresh their contents in place, not rebind the object. + assert self.size == other.size, (self.size, other.size) + self.req.copy_(other.req) + self.pos.copy_(other.pos) + self.write_loc.copy_(other.write_loc) + self.indices.copy_(other.indices) + self.lengths.copy_(other.lengths) + self.history_req.copy_(other.history_req) + self.history_pos.copy_(other.history_pos) + self.history_loc.copy_(other.history_loc) + self.history_valid.copy_(other.history_valid) + self.commit_mask.copy_(other.commit_mask) + + +def _first_row_offsets( + req: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + n = req.numel() + offset = torch.arange(n, device=req.device) + starts = torch.ones(n, dtype=torch.bool, device=req.device) + starts[1:] = req[1:] != req[:-1] + group = starts.cumsum(0) - 1 + group_first = torch.cummax(torch.where(starts, offset, 0), dim=0).values + ends = torch.ones(n, dtype=torch.bool, device=req.device) + ends[:-1] = starts[1:] + group_last = torch.cummin( + torch.where(ends, offset, n - 1).flip(0), dim=0 + ).values.flip(0) + return group, group_first, group_last + + +def window_layout( + req, + pos, + *, + window: int = 128, + capacity: int = 256, + floor: Optional[torch.Tensor] = None, + num_groups: Optional[int] = None, +): + n = pos.numel() + if n == 0: + raise ValueError("request-window layout needs at least one query") + req = req.to(torch.int64) + pos = pos.to(torch.int64) + device = pos.device + groups = n if num_groups is None else int(num_groups) + offset = torch.arange(n, device=device) + group, group_first, group_last = _first_row_offsets(req) + first_pos = pos - (offset - group_first) + history_rows = groups * window + + write_loc = (history_rows + offset).to(torch.int32) + lookback = torch.arange(window, device=device) + seen_pos = pos[:, None] - lookback + old = seen_pos < first_pos[:, None] + old_loc = group[:, None] * window + (seen_pos - (first_pos[:, None] - window)) + new_loc = history_rows + group_first[:, None] + seen_pos - first_pos[:, None] + valid = seen_pos >= 0 + if floor is not None: + floor = floor.to(torch.int64) + valid &= seen_pos >= floor[:, None] + indices = torch.where(valid, torch.where(old, old_loc, new_loc), -1).to(torch.int32) + lengths = valid.sum(-1).to(torch.int32) + + g_req = torch.zeros(groups, dtype=torch.int64, device=device).scatter_( + 0, group, req + ) + g_first = torch.zeros(groups, dtype=torch.int64, device=device).scatter_( + 0, group, first_pos + ) + g_live = torch.zeros(groups, dtype=torch.bool, device=device).scatter_( + 0, group, torch.ones_like(group, dtype=torch.bool) + ) + history_pos = (g_first[:, None] - window + lookback[None, :]).flatten() + history_valid = (history_pos >= 0) & g_live.repeat_interleave(window) + if floor is not None: + g_floor = torch.zeros(groups, dtype=torch.int64, device=device).scatter_( + 0, group, floor + ) + history_valid &= history_pos >= g_floor.repeat_interleave(window) + history_req = g_req.repeat_interleave(window) + history_loc = torch.arange(history_rows, device=device) + + commit_mask = (group_last - offset) < capacity + return WindowLayout( + req, + pos, + write_loc, + indices, + lengths, + history_req, + history_pos, + history_loc, + history_valid, + commit_mask, + history_rows + n, + ) + + +def copy_packed_tokens(src, dst, src_loc, dst_loc, *, page_size, layout=KVLayout.V4): + """Move tokens between paged buffers of ``layout``: a data row and a scale row.""" + if not src_loc.numel(): + return + src_loc, dst_loc = src_loc.long(), dst_loc.long() + for width, base in ( + (layout.data_bytes, 0), + (layout.scale_bytes, page_size * layout.data_bytes), + ): + cols = torch.arange(width, device=src.device) + values = src[ + src_loc[:, None] // page_size, + base + (src_loc[:, None] % page_size) * width + cols, + ] + dst[ + dst_loc[:, None] // page_size, + base + (dst_loc[:, None] % page_size) * width + cols, + ] = values + + +def _capturing() -> bool: + return torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() + + +class RequestWindow: + def __init__( + self, + pool_factory, + *, + num_slots, + layers, + page_size, + capacity, + workspace_rows: Optional[int] = None, + ): + self.capacity = ((capacity + page_size - 1) // page_size) * page_size + self.page_size = page_size + self.pool_factory = pool_factory + self.num_slots = num_slots + self.rows = num_slots * self.capacity + + self.state = pool_factory(self.rows + page_size, layers) + self.zero_row = self.rows + self.sink_row = self.rows + 1 + self.tags = torch.full( + (layers, self.rows + page_size), + -1, + dtype=torch.int64, + device=self.state.kv_buffer[0].device, + ) + self.workspace = None + if workspace_rows: + self._ensure_workspace(workspace_rows) + self.layout = None + self.prepared = None + + def _ensure_workspace(self, rows: int) -> None: + if self.workspace is not None and self.workspace.size >= rows: + return + assert not _capturing(), "request-window workspace must be sized before capture" + size = ((rows + self.page_size - 1) // self.page_size) * self.page_size + self.workspace = self.pool_factory(size, 1) + + def reset(self, slots): + loc = slots.to(torch.int64)[:, None] * self.capacity + torch.arange( + self.capacity, device=slots.device + ) + self.tags[:, loc.flatten()] = -1 + self.prepared = None + + def activate(self, layout): + if self.layout is layout: + return + self.layout = layout + self.prepared = None + if self.workspace is None: + self._ensure_workspace(layout.size) + elif self.workspace.size < layout.size: + # Captured graphs hold the workspace address; growing it strands them. + raise RuntimeError( + f"request-window workspace too small: {self.workspace.size} rows " + f"for a layout of {layout.size}" + ) + + def initialize_dummy_history(self): + layout = self.layout + self.tags.fill_(-1) + loc = layout.history_req * self.capacity + layout.history_pos % self.capacity + for buf in self.state.kv_buffer: + buf.zero_() + self.tags[:, loc] = layout.history_pos + self.prepared = None + + def _history_src(self, layout): + return torch.where( + layout.history_valid, + layout.history_req * self.capacity + layout.history_pos % self.capacity, + self.zero_row, + ) + + def buffer(self, layer): + # The runner's capture scope includes eager warmups before CUDA capture + # starts, so the phase is part of the key: leaving the scope revalidates. + in_capture = get_is_capture_mode() or _capturing() + prepared_key = (layer, in_capture) + if self.prepared != prepared_key: + layout = self.layout + if layout is None: + raise RuntimeError("request-window metadata was not activated") + src = self._history_src(layout) + if not in_capture: + valid = layout.history_valid + if not torch.equal( + self.tags[layer, src][valid], layout.history_pos[valid] + ): + raise RuntimeError( + "SWA history is missing: replay or window ownership is invalid" + ) + copy_packed_tokens( + self.state.kv_buffer[layer], + self.workspace.kv_buffer[0], + src, + layout.history_loc, + page_size=self.page_size, + layout=self.state.kv_layout, + ) + self.prepared = prepared_key + return self.workspace.kv_buffer[0] + + def commit(self, layer): + layout = self.layout + dst = torch.where( + layout.commit_mask, + layout.req * self.capacity + layout.pos % self.capacity, + self.sink_row, + ) + copy_packed_tokens( + self.buffer(layer), + self.state.kv_buffer[layer], + layout.write_loc, + dst, + page_size=self.page_size, + layout=self.state.kv_layout, + ) + self.tags[layer, dst] = layout.pos diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 2b41258fd..be69c4803 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -66,6 +66,12 @@ class PoolName(str, Enum): INDEXER = "indexer" # TODO(hzh0425): Current DeepSeek V4 pool naming is verbose; will be normalized to # 'COMPRESSED_KV / COMPRESSED_INDEXER / COMPRESSED_STATE' in the next PR. + DEEPSEEK_V4_C1 = "deepseek_v4_c1" + DEEPSEEK_V4_C1_INDEXER = "deepseek_v4_c1_indexer" + DEEPSEEK_V4_C1_INDEXER_SCALE = "deepseek_v4_c1_indexer_scale" + DEEPSEEK_V4_C2 = "deepseek_v4_c2" + DEEPSEEK_V4_C2_INDEXER = "deepseek_v4_c2_indexer" + DEEPSEEK_V4_C2_INDEXER_SCALE = "deepseek_v4_c2_indexer_scale" DEEPSEEK_V4_C4 = "deepseek_v4_c4" DEEPSEEK_V4_C4_INDEXER = "deepseek_v4_c4_indexer" # FP4 indexer splits the indexer cache into separate payload/scale buffers, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index e7705d9ea..78acf93e5 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -113,7 +113,7 @@ def _resolve_deepseek_v4_layer_mappings( ) -> _DeepSeekV4LayerMappings: transfer_layer_num = kvcache.end_layer - kvcache.start_layer full = {layer: layer for layer in range(transfer_layer_num)} - swa = {} if getattr(kvcache, "_unified_kv", False) else full.copy() + swa = full.copy() if kvcache.swa_kv_pool is not None else {} c4, c128, c4_state_global_layers = {}, {}, [] for local_layer, item in enumerate( @@ -483,6 +483,13 @@ def _dsv4_compressed_region_buffers(kvcache: Any, ratio: int) -> tuple[list, int return pool.kv_buffer, pool.bytes_per_page_padded +def _dsv4_page_aligned_only(pool: Any) -> bool: + """Whether a pool may only move whole pages: the token-granular copy + (``transfer_cache_dsv4_mla``) hardcodes the V4 data/scale row split.""" + layout = getattr(pool, "kv_layout", None) + return layout is not None and layout.value != "v4" + + @dataclass(frozen=True) class _IndexerRegion: """One page-contiguous indexer buffer group to mirror on the host.""" @@ -578,6 +585,85 @@ def _dsv4_indexer_regions(kvcache: Any, page_size: int) -> list[_IndexerRegion]: ] +def _dsv4_low_ratio_entries( + kvcache: Any, page_size: int, num_host_pages: int, transfer_layer_num: int +): + """Mirror each shared source once, in FULL-page units. Prefixes end on an even + page boundary, so ratio-2's request-scoped ring is rebuilt, not cached.""" + import torch + + entries = [] + for ratio, names in ( + ( + 1, + ( + PoolName.DEEPSEEK_V4_C1, + PoolName.DEEPSEEK_V4_C1_INDEXER, + PoolName.DEEPSEEK_V4_C1_INDEXER_SCALE, + ), + ), + ( + 2, + ( + PoolName.DEEPSEEK_V4_C2, + PoolName.DEEPSEEK_V4_C2_INDEXER, + PoolName.DEEPSEEK_V4_C2_INDEXER_SCALE, + ), + ), + ): + sources = getattr(kvcache, "sources_by_ratio", {}).get(ratio, []) + if not sources: + continue + kv_pool = kvcache.kv_pools[ratio] + index_pool = kvcache.index_pools[ratio] + assert page_size % ratio == 0 + slots_per_page = page_size // ratio + assert slots_per_page % index_pool.page_size == 0 + index_pages_per_full_page = slots_per_page // index_pool.page_size + layer_mapping = { + source - kvcache.start_layer: index for index, source in enumerate(sources) + } + regions = [(names[0], kv_pool, kv_pool.kv_buffer)] + if index_pool.index_k_with_scale_buffer is not None: + index_regions = [(names[1], index_pool.index_k_with_scale_buffer)] + else: + index_regions = [ + (names[1], index_pool.index_k_payload_buffer), + (names[2], index_pool.index_k_scale_buffer), + ] + for name, buffers in index_regions: + # Drop only the padding rows past the FULL page address space. + rows = [] + for buffer in buffers: + full_pages = buffer.shape[0] // index_pages_per_full_page + rows.append( + buffer[: full_pages * index_pages_per_full_page] + .view(torch.uint8) + .reshape(full_pages, -1) + ) + regions.append((name, index_pool, rows)) + for name, device_pool, buffers in regions: + entries.append( + build_pool_entry( + name=name, + host_pool=DeepSeekV4PagedHostPool( + pool_name=str(name), + device_buffers=buffers, + item_bytes=buffers[0].shape[1] * buffers[0].element_size(), + num_host_pages=num_host_pages, + slot_page_size=page_size, + layout=get_memory().hicache_mem_layout, + allocator_type=_get_allocator_type(), + page_aligned_only=True, + ), + device_pool=device_pool, + layer_mapping=layer_mapping, + transfer_layer_num=transfer_layer_num, + ) + ) + return entries + + def _dsv4_rope_sibling( kvcache: Any, ratio: int ) -> Optional[tuple[PoolName, list, int]]: @@ -651,10 +737,10 @@ def build_deepseek_v4_hicache_stack( full_layer_mapping = layer_mappings.full is_unified_kv = getattr(kvcache, "_unified_kv", False) + has_paged_swa = not is_unified_kv and kvcache.swa_kv_pool is not None mtp_swa_device_buffers = [] - if is_unified_kv: - # unified_kv keeps the SWA ring inside the unified pool and never offloads it, - # so there is no separate SWA host pool to map. + if not has_paged_swa: + # Unified KV and encoder replay rebuild SWA state; keep it out of host cache. swa_layer_mapping = {} else: if len(kvcache.swa_kv_pool.kv_buffer) != transfer_layer_num: @@ -707,7 +793,7 @@ def build_deepseek_v4_hicache_stack( ), ] - if not is_unified_kv: + if has_paged_swa: swa_host_pool = DeepSeekV4PagedHostPool( pool_name=str(PoolName.SWA), device_buffers=[ @@ -719,6 +805,7 @@ def build_deepseek_v4_hicache_stack( slot_page_size=kvcache.swa_page_size, layout=get_memory().hicache_mem_layout, allocator_type=_get_allocator_type(), + page_aligned_only=_dsv4_page_aligned_only(kvcache.swa_kv_pool), ) swa_attn_allocator = params.token_to_kv_pool_allocator.swa_attn_allocator entries.append( @@ -749,7 +836,8 @@ def build_deepseek_v4_hicache_stack( slot_page_size=page_size, layout=get_memory().hicache_mem_layout, allocator_type=_get_allocator_type(), - page_aligned_only=is_unified_kv, + page_aligned_only=is_unified_kv + or _dsv4_page_aligned_only(kvcache.c4_kv_pool), ) entries.append( build_pool_entry( @@ -855,7 +943,8 @@ def build_deepseek_v4_hicache_stack( slot_page_size=c128_slot_page_size, layout=get_memory().hicache_mem_layout, allocator_type=_get_allocator_type(), - page_aligned_only=is_unified_kv, + page_aligned_only=is_unified_kv + or _dsv4_page_aligned_only(kvcache.c128_kv_pool), ) # C128 state pool is intentionally not registered with hicache. # page_size=256 % 128 == 0, so state pool is not consumed on load. @@ -897,6 +986,10 @@ def build_deepseek_v4_hicache_stack( if c128_rope_entry is not None: entries.append(c128_rope_entry) + entries.extend( + _dsv4_low_ratio_entries(kvcache, page_size, num_host_pages, transfer_layer_num) + ) + host_pool_group = HostPoolGroup(entries) cache_controller = HybridCacheController( params.token_to_kv_pool_allocator, @@ -1459,10 +1552,12 @@ class _DeepSeekV4Strategy(StackStrategy): def matches(self, kvcache, components): from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool - return isinstance(kvcache, DeepSeekV4TokenToKVPool) and components in ( + if not isinstance(kvcache, DeepSeekV4TokenToKVPool): + return False + return components in ( {ComponentType.FULL, ComponentType.SWA}, {ComponentType.FULL, ComponentType.SWA, ComponentType.C128}, - ) + ) or (components == {ComponentType.FULL} and kvcache.swa_kv_pool is None) def build_direct_linker_pool_group(self, *, kvcache, params, page_size): from sglang.srt.mem_cache.hybrid_cache.linker_pool_assembler import ( @@ -1513,6 +1608,12 @@ class _DeepSeekV4Strategy(StackStrategy): # The *_ROPE entries only resolve under unified fp8 kv; entry_map filters # them out everywhere else. _sidecar_srcs = [ + (PoolName.DEEPSEEK_V4_C1, PoolName.KV), + (PoolName.DEEPSEEK_V4_C1_INDEXER, PoolName.KV), + (PoolName.DEEPSEEK_V4_C1_INDEXER_SCALE, PoolName.KV), + (PoolName.DEEPSEEK_V4_C2, PoolName.KV), + (PoolName.DEEPSEEK_V4_C2_INDEXER, PoolName.KV), + (PoolName.DEEPSEEK_V4_C2_INDEXER_SCALE, PoolName.KV), (PoolName.DEEPSEEK_V4_C4, PoolName.KV), (PoolName.DEEPSEEK_V4_C4_ROPE, PoolName.KV), (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index d6e448158..d13240aba 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -39,6 +39,7 @@ from sglang.srt.environ import envs from sglang.srt.hardware_backend.mlx.runtime import use_mlx from sglang.srt.managers.mm_schedule import init_mm_embedding_cache from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool @@ -233,7 +234,11 @@ def build_kv_cache( ) # Hybrid memory pool - is_hybrid_swa = tp_worker.is_hybrid_swa + token_to_kv_pool = tp_worker.model_runner.token_to_kv_pool + is_hybrid_swa = tp_worker.is_hybrid_swa and ( + not isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) + or token_to_kv_pool.needs_paged_swa_allocator + ) is_hybrid_ssm = uses_ssm_state(tp_worker.model_runner.model_config) is_dsa = is_deepseek_dsa(model_config.hf_config) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index e73db8fcf..f01c1c04e 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -56,7 +56,10 @@ from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, ) -from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + DeepSeekV4TokenToKVPool, + select_dsv4_kv_layout, +) from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool from sglang.srt.mem_cache.memory_pool import ( DSATokenToKVPool, @@ -1237,6 +1240,7 @@ class KVCacheConfigurator: if is_dsv4_model: token_to_kv_pool = self._build_dsv4_kv_pool( max_running_requests=sizes.max_running_requests, + full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, c4_max_total_num_tokens=sizes.c4_max_total_num_tokens, c128_max_total_num_tokens=sizes.c128_max_total_num_tokens, @@ -1340,6 +1344,7 @@ class KVCacheConfigurator: self, *, max_running_requests: int, + full_max_total_num_tokens: int, swa_max_total_num_tokens: Optional[int], c4_max_total_num_tokens: int, c128_max_total_num_tokens: int, @@ -1361,8 +1366,10 @@ class KVCacheConfigurator: compression_ratios = [ COMPRESS_RATIO_NEXTN_LAYER ] * self.layer_info.num_effective_layers + kv_source_layers = [] else: compression_ratios = self.model_config.compress_ratios + kv_source_layers = list(self.model_config.hf_config.kv_source_layer_ids) # NPU keeps its PA_ND KV-pool subclass, while Compressor state sizing # follows the same fixed ring ownership as GPU. Do not replace the @@ -1374,8 +1381,13 @@ class KVCacheConfigurator: ) pool_cls = DSV4NPUTokenToKVPool + kv_layout_kwargs = {} else: pool_cls = DeepSeekV4TokenToKVPool + kv_layout, compressed_kv_layout = select_dsv4_kv_layout() + kv_layout_kwargs = dict( + kv_layout=kv_layout, compressed_kv_layout=compressed_kv_layout + ) token_to_kv_pool = pool_cls( max_num_reqs=max_running_requests, @@ -1404,6 +1416,10 @@ class KVCacheConfigurator: end_layer=self.layer_info.end_layer, enable_hisparse=get_memory().enable_hisparse, online_mtp_max_draft_tokens=(max_speculative_num_draft_tokens() or 0), + kv_source_layers=kv_source_layers, + full_size=full_max_total_num_tokens, + **({"is_draft_worker": self.is_draft_worker} if not _is_npu else {}), + **kv_layout_kwargs, ) if not self.is_draft_worker and token_to_kv_pool._unified_kv: # The draft pool has no C4 layers and shares this req pool, so only @@ -2071,7 +2087,19 @@ class KVCacheConfigurator: need_sort=need_sort, ) else: - if self.is_hybrid_swa and sizes.full_max_total_num_tokens == 0: + if ( + isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) + and not token_to_kv_pool.needs_paged_swa_allocator + ): + token_to_kv_pool_allocator = PagedTokenToKVPoolAllocator( + sizes.full_max_total_num_tokens, + page_size=get_schedule().page_size, + dtype=self.kv_cache_dtype, + device=self.device, + kvcache=token_to_kv_pool, + need_sort=need_sort, + ) + elif self.is_hybrid_swa and sizes.full_max_total_num_tokens == 0: token_to_kv_pool_allocator = PureSWATokenToKVPoolAllocator( sizes.swa_max_total_num_tokens, page_size=get_schedule().page_size, @@ -2141,7 +2169,10 @@ class KVCacheConfigurator: else: assert self.is_draft_worker - if self.is_hybrid_swa: + if self.is_hybrid_swa and ( + not isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) + or token_to_kv_pool.needs_paged_swa_allocator + ): if isinstance( token_to_kv_pool_allocator, DeepSeekV4HiSparseTokenToKVPoolAllocator, diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 8d89774a8..5fbb2e589 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -71,6 +71,7 @@ from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.runtime_context import get_parallel @@ -162,6 +163,10 @@ class KVIndexTranslator: self._swa_write_loc_from_full = ( token_to_kv_pool.translate_loc_from_full_to_swa if isinstance(token_to_kv_pool, BaseSWAKVPool) + and ( + not isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool) + or token_to_kv_pool.request_window is None + ) else None ) diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index e49c88e9a..a9d96cf78 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -814,6 +814,12 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): elif pool_name in ( PoolName.INDEXER, PoolName.DRAFT_INDEXER, + PoolName.DEEPSEEK_V4_C1, + PoolName.DEEPSEEK_V4_C1_INDEXER, + PoolName.DEEPSEEK_V4_C1_INDEXER_SCALE, + PoolName.DEEPSEEK_V4_C2, + PoolName.DEEPSEEK_V4_C2_INDEXER, + PoolName.DEEPSEEK_V4_C2_INDEXER_SCALE, PoolName.DEEPSEEK_V4_C4, PoolName.DEEPSEEK_V4_C4_ROPE, PoolName.DEEPSEEK_V4_C4_INDEXER, diff --git a/python/sglang/srt/model_executor/cuda_graph_config.py b/python/sglang/srt/model_executor/cuda_graph_config.py index a61505bae..760659ede 100644 --- a/python/sglang/srt/model_executor/cuda_graph_config.py +++ b/python/sglang/srt/model_executor/cuda_graph_config.py @@ -84,6 +84,7 @@ ALLOWED_KEYS_PER_PHASE = { "max_context_size", "full_prefill_max_req", "full_prefill_prefix_chunk_tokens", + "max_seq_len", ), } @@ -113,6 +114,9 @@ class PhaseConfig: # chunk variants and chooses the smallest one covering a batch. None uses # the scheduler's aggregate chunked_prefill_size token budget. full_prefill_prefix_chunk_tokens: Optional[int] = None + # Prefill only: a batch whose longest sequence exceeds this replays eagerly, and + # backends that capture context-wide work size it. None defers to token buckets. + max_seq_len: Optional[int] = None def default_prefill_backend() -> str: diff --git a/python/sglang/srt/model_executor/encoder_swa_replay.py b/python/sglang/srt/model_executor/encoder_swa_replay.py new file mode 100644 index 000000000..3f96b7d6e --- /dev/null +++ b/python/sglang/srt/model_executor/encoder_swa_replay.py @@ -0,0 +1,77 @@ +from copy import copy + +import torch + + +def run_encoder_swa_replay(worker, batch): + from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ) + + runner = worker.model_runner + window = runner.token_to_kv_pool.request_window + if window is None or not batch.forward_mode.is_extend_without_speculative(): + return + for i, reset in enumerate(batch.encoder_swa_reset): + if not reset: + continue + slot = batch.req_pool_indices[i : i + 1] + window.reset(slot) + end = batch.prefix_lens[i] + if not end: + continue + if end % 2: + raise ValueError( + "encoder SWA replay requires an even cached-prefix boundary" + ) + start = max(0, end - 128) + req = batch.reqs[i] + replay = copy(batch) + replay.reqs = [req] + replay.input_ids = torch.tensor( + list(req.full_untruncated_fill_ids[start:end]), + dtype=torch.int64, + device=runner.device, + ) + replay.prefill_input_ids_cpu = None + replay.req_pool_indices = slot + replay.req_pool_indices_cpu = batch.req_pool_indices_cpu[i : i + 1] + replay.prefix_lens = [start] + replay.extend_lens = [end - start] + replay.extend_num_tokens = end - start + replay.seq_lens = torch.tensor([end], dtype=torch.int64, device=runner.device) + replay.seq_lens_cpu = torch.tensor([end], dtype=torch.int64) + replay.seq_lens_sum = end + replay.orig_seq_lens = replay.seq_lens + replay.out_cache_loc = runner.req_to_token_pool.req_to_token[ + slot[0], start:end + ].long() + replay.return_logprob = False + replay.top_logprobs_nums = None + replay.token_ids_logprobs = None + replay.extend_logprob_start_lens = [end - start] + replay.extend_input_logprob_token_ids = None + replay.is_prefill_only = True + replay.spec_info = None + replay.sampling_info = None + replay.has_grammar = False + replay.multimodal_inputs = [None] + replay.engram_history = None + hasher = runner.model.model.engram_hasher + if hasher is not None: + n = hasher.max_ngram_size - 1 + ids = list(req.full_untruncated_fill_ids[max(0, start - n) : start]) + replay.engram_history = torch.tensor( + [[0] * (n - len(ids)) + ids], + dtype=torch.int32, + device=runner.device, + ) + fb = ForwardBatch.init_new( + replay, + runner, + capture_hidden_mode=CaptureHiddenMode.NULL, + return_hidden_states_before_norm=False, + ) + fb.encoder_swa_replay = True + runner.forward(fb) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 27322631e..7cc421b14 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -698,6 +698,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For ngram embedding ngram_embedding_info: Optional[NgramEmbeddingInfo] = None + encoder_swa_replay: bool = False # DeepSeek-V4.1 engram, extend only: the n - 1 tokens before each request's # first extend token, oldest first, [bs, n - 1] int32 (see EngramHasher). 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 43d358544..31c7518df 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 @@ -33,7 +33,9 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def _unsupported_derived_weight_cache_error() -> Optional[str]: +def _unsupported_derived_weight_cache_error( + model: Optional[torch.nn.Module] = None, +) -> Optional[str]: """Reject online weight updates that derived-weight caches cannot survive. The HPC-Ops bf16xfp32 GEMM caches the fp32 weight split; in-place loader @@ -41,6 +43,18 @@ 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 model is not None and any( + getattr(module, "_hc_attn_tf32_parts", None) is not None + or getattr(module, "_hc_ffn_tf32_parts", None) is not None + for module in model.modules() + ): + return ( + "Online weight updates are not supported while compensated mHC " + "weight splits are active: captured CUDA graphs retain these derived " + "weights. Restart with SGLANG_OPT_DEEPGEMM_HC_PRENORM=0 to use " + "online weight updates." + ) + from sglang.kernels.ops.attention.dsv4.gemm import hpc_bf16xfp32_gemm_enabled if hpc_bf16xfp32_gemm_enabled(): @@ -148,7 +162,7 @@ class WeightUpdater: ) -> tuple[bool, str]: """Update engine weights in-place from the disk.""" self._assert_weight_cache_inactive("update_weights_from_disk") - error = _unsupported_derived_weight_cache_error() + error = _unsupported_derived_weight_cache_error(self.get_model()) if error is not None: return False, error @@ -238,7 +252,7 @@ class WeightUpdater: shape: the shape of the parameter to be updated. """ self._assert_weight_cache_inactive("update_weights_from_distributed") - error = _unsupported_derived_weight_cache_error() + error = _unsupported_derived_weight_cache_error(self.get_model()) if error is not None: return False, error @@ -322,7 +336,7 @@ class WeightUpdater: named_tensors: List[Tuple[str, Union[torch.Tensor, LocalSerializedTensor]]], load_format: Optional[str] = None, ): - error = _unsupported_derived_weight_cache_error() + error = _unsupported_derived_weight_cache_error(self.get_model()) if error is not None: return False, error @@ -388,7 +402,7 @@ class WeightUpdater: def update_weights_from_ipc(self: WeightUpdater, recv_req): """Update weights from IPC for checkpoint-engine integration.""" self._assert_weight_cache_inactive("update_weights_from_ipc") - error = _unsupported_derived_weight_cache_error() + error = _unsupported_derived_weight_cache_error(self.get_model()) if error is not None: return False, error diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 19c4f3e30..a8f972b45 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -804,6 +804,45 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): return self._solve_pool_sizes(max_total_num_tokens, page_size) +def compute_swa_request_cap(*, page_size: int, window: int, attn_dp_size: int) -> int: + """Worst-case SWA slots the scheduler holds live at max_running_requests.""" + draft_tokens = get_spec().speculative_num_draft_tokens or 1 + eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get()) + + # __________[padding][eviction_interval][window] + # Padding to make sure eviction point is page-aligned. + trailing_tokens = window + eviction_interval * draft_tokens + page_size + if get_spec().speculative_algorithm is None: + decode_alloc = page_size + elif get_schedule().disable_overlap_schedule: + # spec-v1: new_tokens_required_next_decode per request. + decode_alloc = spec_decode_alloc_len_per_request( + page_size=page_size, + speculative_num_steps=get_spec().speculative_num_steps, + speculative_eagle_topk=get_spec().speculative_eagle_topk, + speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens, + ) + else: + # spec-v2: the overlap allocator keeps 2 * alloc_len outstanding + # (eagle_utils.eagle_prepare_for_decode: kv_committed_len + 2 * alloc_len). + decode_alloc = 2 * get_alloc_len_per_decode() + per_request = trailing_tokens + decode_alloc + + num_reqs = get_schedule().max_running_requests // attn_dp_size + if get_disagg().disaggregation_mode == "decode": + return ( + per_request * num_reqs + + (window + page_size) * get_disagg().disaggregation_decode_extra_slots + ) + else: + chunks_in_flight = 1 if get_schedule().disable_overlap_schedule else 2 + return ( + per_request * num_reqs + + chunks_in_flight * get_schedule().chunked_prefill_size + + page_size + ) + + class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): """Hybrid SWA configurator with the SWA pool sized from a fixed token cap. @@ -818,45 +857,11 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): super().__init__(kvc) assert self._full_layers_num > 0 - page_size = kvc.page_size - window = kvc.sliding_window_size - draft_tokens = get_spec().speculative_num_draft_tokens or 1 - eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get()) - - """ - __________[padding][eviction_interval][window] - Padding to make sure eviction point is page-aligned. - """ - trailing_tokens = window + eviction_interval * draft_tokens + page_size - if get_spec().speculative_algorithm is None: - decode_alloc = page_size - elif get_schedule().disable_overlap_schedule: - # spec-v1: new_tokens_required_next_decode per request. - decode_alloc = spec_decode_alloc_len_per_request( - page_size=page_size, - speculative_num_steps=get_spec().speculative_num_steps, - speculative_eagle_topk=get_spec().speculative_eagle_topk, - speculative_num_draft_tokens=get_spec().speculative_num_draft_tokens, - ) - else: - # spec-v2: the overlap allocator keeps 2 * alloc_len outstanding - # (eagle_utils.eagle_prepare_for_decode: kv_committed_len + 2 * alloc_len). - decode_alloc = 2 * get_alloc_len_per_decode() - per_request = trailing_tokens + decode_alloc - - num_reqs = get_schedule().max_running_requests // kvc.ps.attn_dp_size - if get_disagg().disaggregation_mode == "decode": - self._swa_cap = ( - per_request * num_reqs - + (window + page_size) * get_disagg().disaggregation_decode_extra_slots - ) - else: - chunks_in_flight = 1 if get_schedule().disable_overlap_schedule else 2 - self._swa_cap = ( - per_request * num_reqs - + chunks_in_flight * get_schedule().chunked_prefill_size - + page_size - ) + self._swa_cap = compute_swa_request_cap( + page_size=kvc.page_size, + window=kvc.sliding_window_size, + attn_dp_size=kvc.ps.attn_dp_size, + ) @staticmethod def is_applicable(kvc: KVCacheConfigurator) -> bool: @@ -915,6 +920,18 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): ) +# Used when --swa-full-tokens-ratio is at its default and cap mode is unusable. +DSV4_DEFAULT_SWA_FULL_TOKENS_RATIO = 0.1 + + +def _operator_swa_full_tokens_ratio() -> Optional[float]: + """The operator's --swa-full-tokens-ratio, or None when it was not given.""" + schedule = get_schedule() + if not schedule._swa_full_tokens_ratio_explicitly_set: + return None + return schedule.swa_full_tokens_ratio + + @dataclass class _DSV4PoolSizes: full_max_total_num_tokens: int @@ -940,6 +957,28 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): self.qk_nope_head_dim = cfg.qk_nope_head_dim self.qk_rope_head_dim = cfg.qk_rope_head_dim self.indexer_head_dim = cfg.index_head_dim + self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, + is_unified_kv_triton, + ) + from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + dsv4_unified_row_bytes, + ) + + # Resolve the unified-kv gate before any sizing so the two cannot drift. + self._unified = is_unified_kv_triton() + self._unified_fp8 = is_unified_kv_fp8() + # Row width across both unified pools: 1024 B bf16, 640 B fp8. + self._unified_row_bytes = dsv4_unified_row_bytes( + self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8 + ) + if self._unified: + # Unified_kv stores the whole latent: one bf16 row, or fp8 nope + bf16 rope. + self.kv_bytes = self._unified_row_bytes + else: + # One FlashMLA-layout latent slot, in bytes. + self.kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 # HIP takes the FP4-accurate byte count here. The NVIDIA FP4 path # keeps the FP8 estimate. self.indexer_bytes_per_token = get_dsv4_indexer_bytes_per_token( @@ -958,10 +997,17 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): f"local={len(self.compression_ratios)}/{len(cfg.compress_ratios)}" ) self.swa_page_size = cfg.window_size + self.operator_swa_ratio = _operator_swa_full_tokens_ratio() + self.swa_ratio = ( + self.operator_swa_ratio + if self.operator_swa_ratio is not None + else DSV4_DEFAULT_SWA_FULL_TOKENS_RATIO + ) self.sliding_window_size = kvc.sliding_window_size - self.swa_ratio = get_schedule().swa_full_tokens_ratio + self.page_size = kvc.page_size self.is_speculative = get_spec().speculative_algorithm is not None self.online_c128_mtp_max_draft_tokens = max_speculative_num_draft_tokens() or 0 + self.attn_dp_size = kvc.ps.attn_dp_size self.requested_max_running_requests_per_worker = ( get_schedule().max_running_requests // kvc.ps.attn_dp_size if get_schedule().max_running_requests is not None @@ -987,31 +1033,21 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): self.num_layers_total = len(self.compression_ratios) self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4) self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128) - - # Unified-KV uses a different physical layout than the non-unified V4 path: - # * one row carries the full latent -- 1024 B bf16, or 640 B under - # SGLANG_DSV4_UNIFIED_KV_FP8 (512 B fp8 nope + 128 B bf16 rope) -- not - # that path's 584-byte fp8(nope) + bf16(rope) + scales cell. - # * SWA is a fixed per-request ring (num_req_slots * ring_size), - # independent of full_token, so it is a fixed *bias* rather than a - # per-token term. Gate on the same switch the pool itself uses so the - # sizing and the allocation never drift apart. - from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( - is_unified_kv_fp8, - is_unified_kv_triton, + # The low-ratio indexer pools are built with force_fp4=True + # (deepseek_v4_memory_pool), so they are fp4 whatever dtype c4 uses. + low_ratio_index_bytes = get_dsv4_indexer_bytes_per_token( + self.indexer_head_dim, use_fp4_indexer=True + ) + self.low_ratio_bytes_per_full_token = sum( + (self.kv_bytes + low_ratio_index_bytes) / cfg.compress_ratios[l] + for l in cfg.hf_config.kv_source_layer_ids + if kvc.layer_info.start_layer <= l < kvc.layer_info.end_layer + and cfg.compress_ratios[l] in (1, 2) ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( dsv4_unified_row_bytes, ) - self._unified = is_unified_kv_triton() - self._unified_fp8 = is_unified_kv_fp8() - self.attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - # Row width across both pools: 1024 B bf16, 640 B fp8. Read from the pool - # module so sizing can't drift from the allocation. - self._unified_row_bytes = dsv4_unified_row_bytes( - self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8 - ) # swa_page_size is the model's sliding window (cfg.window_size). self._swa_ring_size = get_swa_ring_size(self.swa_page_size, self.is_speculative) self._spec_infl = 1.0 @@ -1044,8 +1080,47 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): max_speculative_num_draft_tokens() or 0 ) + from sglang.srt.runtime_context import get_exec + + self.encoder_replay = get_exec().features.enable_encoder_swa_bounded_replay + self.paged_draft_layers = 0 + if self.encoder_replay and kvc.spec_algorithm.is_dspark(): + self.paged_draft_layers = int( + kvc.spec_aux_config.dflash_draft_num_layers or 0 + ) + assert self.paged_draft_layers > 0, "DSpark draft layer count is required" + self.request_window_bytes = 0 + if self.encoder_replay: + slots = self.requested_max_running_requests_per_worker + 1 + capacity = ceil_align( + self.sliding_window_size + self.online_c128_mtp_max_draft_tokens, + self.page_size, + ) + layers = self.num_layers_total + scratch = ( + max( + get_schedule().chunked_prefill_size, + slots * max(128, self.online_c128_mtp_max_draft_tokens), + ) + + slots * 128 + + self.page_size + ) + self.request_window_bytes = ( + (slots * capacity + self.page_size) * layers * (self.kv_bytes + 16) + + 4 * scratch * (self.kv_bytes + 16) + + slots * 3 * 16 * self.attn_head_dim * 8 + ) + if not self.paged_draft_layers: + self.swa_ratio = 0 + self.swa_prefix_tails = self._resolve_swa_prefix_tails() + self.swa_cap_tokens = ( + 0 + if self.encoder_replay and not self.paged_draft_layers + else self._resolve_swa_cap_tokens() + ) + self.bytes_per_swa_token = self._get_bytes_per_swa_token() self.bytes_per_full_token = self._get_bytes_per_full_token() - if self.is_speculative: + if self.is_speculative and not self.encoder_replay: # Reserve memory for the speculative draft worker by inflating # per-token bytes by (target+draft)/target. Equivalent to dflash's # scale_kv_cell_size_per_token_for_dflash but applied to @@ -1054,6 +1129,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): target_layers = self.num_layers_total self._spec_infl = (target_layers + draft_layers) / target_layers self.bytes_per_full_token *= self._spec_infl + self.bytes_per_swa_token *= self._spec_infl # Online c128 keeps a single in-progress (max, sum, kv) state per index # and assumes a strict forward-only schedule. Speculative decode (MTP) @@ -1105,72 +1181,133 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): f"get_compress_state_ring_size()." ) - def _get_bytes_per_full_token(self) -> float: - if self._unified: - # Unified_kv stores the whole latent: one bf16 pool, or an fp8 nope - # pool plus a bf16 rope pool. kv_bytes also prices the compressed - # c4/c128 rows below, which live in the same pool(s). - kv_bytes = self._unified_row_bytes - else: - kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 + def _resolve_swa_prefix_tails(self) -> int: + """Cached prefix tails cap mode keeps addressable: a prefix is reusable only + while its last sliding_window tokens still hold SWA slots.""" + prefix_tails = get_schedule().swa_prefix_tails + if prefix_tails is not None: + return prefix_tails + if get_memory().disable_radix_cache: + # Nothing is kept for reuse, so the request cap alone bounds the pool. + return 0 + max_running_requests = self.requested_max_running_requests_per_worker + return 4 * max_running_requests if max_running_requests is not None else 0 - attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim - c4_state_dtype_size, c128_state_dtype_size = ( - _get_dsv4_compress_state_dtype_sizes() + def _resolve_swa_cap_tokens(self) -> Optional[int]: + """SWA slots to reserve in cap mode, None to keep ratio sizing. Cap mode + budgets from the request cap plus radix headroom, not full_tokens.""" + if self.operator_swa_ratio is not None: + return None + if self._unified: + # Ring mode: SWA is a fixed per-request ring, with no paged pool to size. + return None + max_running_requests = self.requested_max_running_requests_per_worker + if max_running_requests is None or self.sliding_window_size is None: + return None + chunked_prefill_size = get_schedule().chunked_prefill_size + if self.disaggregation_mode != "decode" and ( + chunked_prefill_size is None or chunked_prefill_size <= 0 + ): + return None + + cap = compute_swa_request_cap( + page_size=self.page_size, + window=self.sliding_window_size, + attn_dp_size=self.attn_dp_size, ) - c4_state_bytes = 2 * 2 * attn_head_dim * c4_state_dtype_size + headroom = self.swa_prefix_tails * (self.sliding_window_size + self.page_size) + return ceil_align(cap + headroom, self.page_size) + + def _get_bytes_per_swa_token(self) -> float: + """Bytes one SWA slot costs across the stage. c4_state_pool_size = swa_tokens + / swa_page_size * ring, so c4 compress state is priced per SWA slot too.""" + if self.encoder_replay: + # Target SWA lives in the request window; only the draft owns paged SWA + # bytes, and its layers carry no compressed state. + return self.kv_bytes * self.paged_draft_layers + c4_state_dtype_size, _ = _get_dsv4_compress_state_dtype_sizes() + c4_state_bytes = 2 * 2 * self.attn_head_dim * c4_state_dtype_size + c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size + + c4_state_ratio = self.c4_ring_size / self.swa_page_size + return ( + self.kv_bytes * self.num_layers_total + + c4_state_ratio + * (c4_state_bytes + c4_indexer_state_bytes) + * self.num_layers_ca4 + ) + + def _get_bytes_per_full_token(self) -> float: + _, c128_state_dtype_size = _get_dsv4_compress_state_dtype_sizes() # Online c128 stores (max, sum, kv) per slot (3*head_dim) instead of # raw (kv, score) (2*head_dim). Combined with ring_size=1 this still # nets a large reduction (~3/256x) but the per-slot bytes go up. c128_online = envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() c128_state_bytes = ( - (3 if c128_online else 2 * 1) * attn_head_dim * c128_state_dtype_size + (3 if c128_online else 2 * 1) * self.attn_head_dim * c128_state_dtype_size ) - c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size - c4_state_ratio = self.c4_ring_size / self.swa_page_size # C128 state is request-scoped and is finalized after # max_running_requests is known, so it should not scale with # full-token capacity here. c128_state_ratio = 0 + # Cap mode and ring mode both move the SWA pool and the c4 state that + # follows it out of the coefficient and into fixed bytes. + swa_ratio = ( + 0 if self._unified or self.swa_cap_tokens is not None else self.swa_ratio + ) c4_frac = 1 / (4 * self.c4_shrink_factor) return ( - # Ring mode: SWA is a fixed per-request pool (see _fixed_swa_bytes). - ( - 0.0 - if self._unified - else self.swa_ratio * kv_bytes * self.num_layers_total - ) - + c4_frac * kv_bytes * self.num_layers_ca4 - + 1 / 128 * kv_bytes * self.num_layers_ca128 + swa_ratio * self.bytes_per_swa_token + + self.low_ratio_bytes_per_full_token + + c4_frac * self.kv_bytes * self.num_layers_ca4 + + 1 / 128 * self.kv_bytes * self.num_layers_ca128 + 1 / 4 * self.indexer_bytes_per_token * self.num_layers_ca4 - # Ring mode: C4 state is per-request too (see _fixed_c4_state_bytes). - + ( - 0.0 - if self._unified - else self.swa_ratio - * c4_state_ratio - * c4_state_bytes - * self.num_layers_ca4 - ) + c128_state_ratio * c128_state_bytes * self.num_layers_ca128 - + ( - 0.0 - if self._unified - else self.swa_ratio - * c4_state_ratio - * c4_indexer_state_bytes - * self.num_layers_ca4 - ) ) + def _get_swa_fixed_bytes(self) -> float: + """Bias bytes the SWA pool takes in cap mode; 0 when sizing by ratio.""" + paged_bytes = ( + 0 + if self.swa_cap_tokens is None + else self.swa_cap_tokens * self.bytes_per_swa_token + ) + return self.request_window_bytes + paged_bytes + + def _get_swa_tokens(self, full_token: int, page_size: int) -> int: + # swa_cap_tokens was already page-aligned at resolve time. + if self.swa_cap_tokens is None: + return int(full_token * self.swa_ratio) // page_size * page_size + return self.swa_cap_tokens + def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes: full_token = full_token // page_size * page_size - swa_tokens = int(full_token * self.swa_ratio) // page_size * page_size - if not self._unified: - # Ring mode: the paged SWA pool is vestigial, so its floor does not apply. - self.validate_swa_pool_size(swa_tokens, self.sliding_window_size, page_size) + swa_tokens = self._get_swa_tokens(full_token, page_size) + if self.swa_cap_tokens is None: + # Only ratio sizing can under-size a request: cap mode sizes from the + # request floor, and encoder replay deliberately runs swa_tokens == 0. + if not self._unified: + self.validate_swa_pool_size( + swa_tokens, self.sliding_window_size, page_size + ) + source = "explicit" if self.operator_swa_ratio is not None else "default" + mode = ( + "ring (paged swa_tokens vestigial)" + if self._unified + else f"ratio ({source})" + ) + logger.info( + f"DSV4 SWA sizing: mode={mode}, swa_tokens={swa_tokens}, " + f"swa_full_tokens_ratio={self.swa_ratio}" + ) + else: + logger.info( + f"DSV4 SWA sizing: mode=cap, swa_tokens={swa_tokens}, " + f"request_cap+headroom={self.swa_cap_tokens}, " + f"prefix_tails={self.swa_prefix_tails}" + ) return _DSV4PoolSizes( full_max_total_num_tokens=full_token, swa_max_total_num_tokens=swa_tokens, @@ -1195,18 +1332,17 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): return 0 _, c128_state_dtype_size = _get_dsv4_compress_state_dtype_sizes() - attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim num_req_slots = self._get_num_req_slots(max_running_requests) if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(): state_rows = num_req_slots + self.c128_ring_size + 1 state_rows *= 1 + self.online_c128_mtp_max_draft_tokens - state_last_dim = 3 * attn_head_dim + state_last_dim = 3 * self.attn_head_dim else: state_pool_size = num_req_slots * self.c128_ring_size state_rows = state_pool_size + self.c128_ring_size + 1 state_rows = ceil_div(state_rows, 128) * 128 - state_last_dim = 2 * attn_head_dim + state_last_dim = 2 * self.attn_head_dim return ( state_rows * state_last_dim * c128_state_dtype_size * self.num_layers_ca128 @@ -1313,14 +1449,23 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): max_running_requests_per_worker ) - available_bytes_for_tokens = max( - available_bytes - - c128_state_fixed_bytes - - swa_ring_fixed_bytes - - c4_state_fixed_bytes, - 0, + swa_fixed_bytes = self._get_swa_fixed_bytes() + fixed_bytes = ( + c128_state_fixed_bytes + + swa_fixed_bytes + + swa_ring_fixed_bytes + + c4_state_fixed_bytes ) + available_bytes_for_tokens = max(available_bytes - fixed_bytes, 0) full_token = int(available_bytes_for_tokens / self.bytes_per_full_token) + if full_token <= 0 and self.swa_cap_tokens is not None: + raise RuntimeError( + f"The DSV4 SWA pool cap ({self.swa_cap_tokens} tokens, " + f"{swa_fixed_bytes / (1 << 30):.2f} GB) leaves no room for the full " + f"KV pool within the available {available_bytes / (1 << 30):.2f} GB. " + f"Reduce --max-running-requests, lower --swa-prefix-tails " + f"or SGLANG_SWA_EVICTION_INTERVAL, or increase --mem-fraction-static." + ) sizes = self._compute_dsv4_sizes(full_token, page_size) logger.info( @@ -1329,6 +1474,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): f"bytes_per_full_token={self.bytes_per_full_token:.2f}, " f"available_bytes={available_bytes / (1 << 30):.2f} GB, " f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, " + f"swa_fixed={swa_fixed_bytes / (1 << 30):.2f} GB, " f"swa_ring_fixed={swa_ring_fixed_bytes / (1 << 30):.2f} GB, " f"c4_state_fixed={c4_state_fixed_bytes / (1 << 30):.2f} GB, " f"full_token={sizes.full_max_total_num_tokens}" diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index 81951c608..d102d94e2 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -635,7 +635,11 @@ class BaseRunner(ABC): spec_algorithm=mr.spec_algorithm, spec_info=spec_info, capture_hidden_mode=capture_hidden_mode, - num_token_non_padded=buffers.num_token_non_padded, + # Maintained only under expert parallelism; None elsewhere so routing + # does not mask every row against a never-filled zero count. + num_token_non_padded=( + buffers.num_token_non_padded if enable_num_token_non_padded() else None + ), global_forward_mode=capture_forward_mode, lora_ids=lora_ids, ) @@ -649,6 +653,8 @@ class BaseRunner(ABC): forward_batch = mr.prepare_dummy_forward_batch(forward_batch) mr.attn_backend.init_forward_metadata(forward_batch) + if get_exec().features.enable_encoder_swa_bounded_replay: + mr.token_to_kv_pool.request_window.initialize_dummy_history() def run_once(): # Reused dummy batches may carry DP-local lazy caches from a prior 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 0c64e9b42..2bd702619 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 @@ -52,6 +52,7 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.attention.graph_variants import ( AttentionGraphVariants, create_attention_graph_variants, + create_dsv41_candidate_graph_variants, ) from sglang.srt.layers.cp.utils import is_mla_cp_enabled from sglang.srt.layers.dp_attention import ( @@ -300,6 +301,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): 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, self.captured_req_width + ) ) # --- bucket sizes --------------------------------------------- @@ -965,7 +969,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): spec_algorithm=self.model_runner.spec_algorithm, spec_info=spec_info, capture_hidden_mode=self.capture_hidden_mode, - num_token_non_padded=buffers.num_token_non_padded, + # Maintained only under expert parallelism; None elsewhere so routing + # does not mask every row against a never-filled zero count. + num_token_non_padded=( + buffers.num_token_non_padded if enable_num_token_non_padded() else None + ), attn_tp_sequence_sharded=attn_tp_sharded, global_forward_mode=self.capture_forward_mode, lora_ids=lora_ids, diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 94153a6cd..809b4c593 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -342,7 +342,7 @@ def maybe_flashinfer_autotune_speculative_draft( 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 @@ -351,14 +351,27 @@ 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 + # DSpark's dummy forward is TARGET_VERIFY-shaped and misses large prefill GEMMs. + prefill_autotune = getattr(mr.model, "autotune_prefill_kernels", None) + wants_prefill_autotune = getattr(mr.model, "wants_prefill_autotune", None) + if wants_prefill_autotune is not None and not wants_prefill_autotune(): + # Entering the autotune context loads / saves the tactic cache and syncs + # ranks, so a model that has nothing to tune must decline before it. + prefill_autotune = 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/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index f0d3ab989..5ffbf697d 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -84,6 +84,7 @@ from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ForwardMode, + NgramEmbeddingInfo, PPProxyTensors, compute_local_num_token_non_padded, enable_num_token_non_padded, @@ -283,6 +284,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): buffer population, attention metadata init, and output slicing. """ + _backend_can_run_prefill_cuda_graph = None + def __init__(self, model_runner: ModelRunner): if get_schedule().enable_mixed_chunk: backend = get_exec().graph.cuda_graph_config.prefill.backend @@ -291,6 +294,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): f"graph backend; got '{backend}'." ) super().__init__(model_runner) + self._backend_can_run_prefill_cuda_graph = getattr( + model_runner.attn_backend, "can_run_prefill_cuda_graph", None + ) # --- model flags ---------------------------------------------- self.quant_config = getattr(model_runner.model, "quant_config", None) self.is_multimodal = model_runner.model_config.is_multimodal @@ -1332,6 +1338,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): is None ): return False + backend_can_run = self._backend_can_run_prefill_cuda_graph + if backend_can_run is not None and not backend_can_run(forward_batch): + return False # Multi-req replay is supported by body-capture backends via the # layer_model.forward monkey-patch in replay(): the captured graph runs # the transformer stack, then the outer model.forward runs @@ -1497,6 +1506,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): return_pooled_hidden_states=self.capture_return_pooled_hidden_states, max_seq_len_override=self.max_context_size, ) + ngram_manager = self.model_runner.ngram_embedding_manager + if ngram_manager.enabled: + forward_batch.ngram_embedding_info = NgramEmbeddingInfo.create( + ngram_manager.table, + bs, + self.device, + column_starts=0, + req_lens=shape_inputs["extend_seq_lens"], + ) self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens) return forward_batch, self.model_runner.attn_backend @@ -1827,6 +1845,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): ), max_seq_len_override=self.max_context_size, ) + # The n-gram hasher runs outside the graph and reads this at replay. + static_forward_batch.ngram_embedding_info = forward_batch.ngram_embedding_info + static_forward_batch.engram_history = forward_batch.engram_history if self._is_full_backend: forward_batch.next_token_logits_buffer = ( static_forward_batch.next_token_logits_buffer @@ -1901,6 +1922,26 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): return static_forward_batch + def _fill_input_embeds_slot(self, args, layer_kwargs, static_num_tokens: int): + """A text-only batch would otherwise replay the captured input_embeds.""" + ie_idx = self._input_embeds_arg_idx + ie = layer_kwargs.get("input_embeds") + if ie is None and ie_idx is not None and len(args) > ie_idx: + ie = args[ie_idx] + if ie is None: + input_ids = layer_kwargs.get("input_ids") + if input_ids is None and len(args) > 0: + input_ids = args[0] + embed = getattr(self.model_runner.model, "get_input_embeddings", None) + assert input_ids is not None and embed is not None, ( + "prefill CUDA graph replay needs input_embeds for the static " + "slot, and the model exposes no get_input_embeddings()" + ) + ie = embed()(input_ids) + self.buffer_registry.get_slot("input_embeds").slice_for(1, static_num_tokens)[ + : ie.shape[0] + ].copy_(ie) + def _execute_body_capture( self, forward_batch: ForwardBatch, @@ -1913,7 +1954,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # BCG / Full: replay the captured body, run the LM head + # logits_processor eagerly. full_path = self._is_full_backend - ie_idx = self._input_embeds_arg_idx def replay_layer_forward(*args, **layer_kwargs): # The captured body graph reads activations from the static @@ -1925,16 +1965,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # Copy them into the slot before replay so the graph sees the # current request's embeddings (mirrors main's BCG closure). if self.buffer_registry.has_slot("input_embeds"): - ie = layer_kwargs.get("input_embeds") - if ie is None and ie_idx is not None and len(args) > ie_idx: - ie = args[ie_idx] - if ie is None: - # Otherwise the graph replays the previous batch's embeddings. - input_ids = args[0] if args else layer_kwargs["input_ids"] - ie = self.model_runner.model.get_input_embeddings()(input_ids) - self.buffer_registry.get_slot("input_embeds").slice_for( - 1, static_num_tokens - )[: ie.shape[0]].copy_(ie) + self._fill_input_embeds_slot(args, layer_kwargs, static_num_tokens) hs = self.backend.replay(shape_key, static_forward_batch, **kwargs) return _slice_output_rows(hs, raw_num_tokens) if full_path else hs diff --git a/python/sglang/srt/model_executor/runner/shape_key.py b/python/sglang/srt/model_executor/runner/shape_key.py index 7691b7e07..c725e7d51 100644 --- a/python/sglang/srt/model_executor/runner/shape_key.py +++ b/python/sglang/srt/model_executor/runner/shape_key.py @@ -26,5 +26,5 @@ class ShapeKey: stream_idx: Optional[int] = None # LoRA or prefill-prefix variant; None selects the default. variant_label: Optional[str] = None - # Independent attention variant; None selects the default. + # Independent attention variant (DSA dense/sparse, candidate_*); None is default. attention_variant: Optional[str] = None diff --git a/python/sglang/srt/model_executor/runner_utils/capture_mode.py b/python/sglang/srt/model_executor/runner_utils/capture_mode.py index a2195b9d0..38770ea40 100644 --- a/python/sglang/srt/model_executor/runner_utils/capture_mode.py +++ b/python/sglang/srt/model_executor/runner_utils/capture_mode.py @@ -72,6 +72,13 @@ def get_capture_attention_variant() -> Optional[str]: return _capture_attention_variant +def skip_low_ratio_indexer(compress_ratio: int) -> bool: + """Whether the captured candidate variant selects every position for this ratio.""" + return _capture_attention_variant == "candidate_all" or ( + _capture_attention_variant == "candidate_c2_all" and compress_ratio == 2 + ) + + def _set_capture_attention_variant(variant: Optional[str]) -> None: global _capture_attention_variant _capture_attention_variant = variant diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 7c2cd05a0..778fcc16e 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -22,6 +22,7 @@ from __future__ import annotations import logging from contextlib import contextmanager, nullcontext +from functools import cached_property from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import torch @@ -92,6 +93,7 @@ from sglang.srt.layers.moe import ( get_moe_a2a_backend, get_moe_runner_backend, post_experts_all_reduce, + should_skip_post_experts_all_reduce, should_use_flashinfer_cutlass_moe_fp4_allgather, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class @@ -121,7 +123,11 @@ from sglang.srt.layers.quantization.fp8_utils import ( view_aiter_fused_rms_transposed_fp8_scale, ) from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import ( + Mxfp4FlashinferTrtllmMoEMethod, + Mxfp8RoutedInputPreQuant, maybe_fuse_routed_scale_and_shared_add, + routed_hidden_size, + should_use_fuse_finalize_all_reduce, ) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope_wrapper @@ -180,6 +186,7 @@ from sglang.srt.models.deepseek_common.utils import ( quant_blocks_shared_experts_fusion, tiny_router_gemm_max_tokens, ) +from sglang.srt.multimodal.dsv41.vl_routing import vision_topk from sglang.srt.runtime_context import ( attention_backends, get_device, @@ -456,6 +463,7 @@ class MoEGate(nn.Module): prefix: str = "", is_hash_moe: bool = False, is_deepseek_v4: bool = False, + vl_correction_bias: bool = False, ): super().__init__() self.is_deepseek_v4 = is_deepseek_v4 @@ -488,6 +496,12 @@ class MoEGate(nn.Module): self.e_score_correction_bias = nn.Parameter(correction_bias) else: self.e_score_correction_bias = None + self.e_score_correction_bias_vl = None + if vl_correction_bias: + self.e_score_correction_bias_vl = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) if _is_cpu and _is_cpu_amx_available: self.quant_method = PackWeightMethod(weight_names=["weight"]) self.tiny_router_gemm_max_tokens = tiny_router_gemm_max_tokens( @@ -538,6 +552,11 @@ class MoEGate(nn.Module): return logits +# 96 rows of 5120 bf16 fit the 1 MiB CustomAllReduceV2 push slot the whole +# [T, hidden] view is staged through. +_FUSED_FINALIZE_ALL_REDUCE_MAX_TOKENS = 96 + + class DeepseekV2MoE(nn.Module): def __init__( self, @@ -546,8 +565,10 @@ class DeepseekV2MoE(nn.Module): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + routed_quant_stream: Optional[torch.cuda.Stream] = None, is_nextn: bool = False, is_deepseek_v4: bool = False, + vl_correction_bias: bool = False, ): super().__init__() self.tp_size = get_parallel().tp_size @@ -585,7 +606,14 @@ class DeepseekV2MoE(nn.Module): self.config = config self.layer_id = layer_id self.alt_stream = alt_stream + self.routed_quant_stream = routed_quant_stream self.is_nextn = is_nextn + self._fuse_finalize_all_reduce = ( + is_deepseek_v4 + and getattr(config, "hc_pre_from_prev_sublayer", False) + and get_platform().is_blackwell + and self.tp_size == 4 + ) n_hash_layers = getattr(config, "num_hash_layers", 0) self.is_hash = layer_id < n_hash_layers and not (is_deepseek_v4 and is_nextn) @@ -608,6 +636,7 @@ class DeepseekV2MoE(nn.Module): prefix=add_prefix("gate", prefix), is_hash_moe=self.is_hash, is_deepseek_v4=is_deepseek_v4, + vl_correction_bias=vl_correction_bias, ) # scaling factor for fused shared experts on AMD-platform. @@ -681,6 +710,12 @@ class DeepseekV2MoE(nn.Module): topk_kwargs.update( use_grouped_topk=False, scoring_func=config.scoring_func, + sqrtsoftplus_log1p=( + getattr(config, "model_type", None) == "deepseek_v41" + ), + fused_gate_packed_ids=( + getattr(config, "model_type", None) == "deepseek_v41" + ), is_fp4_experts=getattr(quant_config, "is_fp4_experts", False), apply_routed_scaling_factor_on_output=( True @@ -951,6 +986,13 @@ class DeepseekV2MoE(nn.Module): else self._maybe_quant_moe_input_once(hidden_states) ) self.alt_stream.wait_stream(current_stream) + should_quant_routed_input_mxfp8 = ( + not use_flashinfer_trtllm_bypass + and pre_quant_input is None + and self._should_quant_routed_input_mxfp8(hidden_states) + ) + if should_quant_routed_input_mxfp8: + self.routed_quant_stream.wait_stream(current_stream) has_shared_output = ( hidden_states.shape[0] > 0 and self.num_fused_shared_experts == 0 ) @@ -959,6 +1001,7 @@ class DeepseekV2MoE(nn.Module): if get_exec().moe.enable_eplb and not self.is_nextn else None ) + # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, gemm_output_zero_allocator) if use_flashinfer_trtllm_bypass: @@ -973,14 +1016,44 @@ class DeepseekV2MoE(nn.Module): if getattr(self, "is_hash", False) else {} ) - topk_output = self.topk( - hidden_states, - router_logits, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=dispatch_info, - **topk_kwargs, + if self.gate.e_score_correction_bias_vl is not None: + topk_output = vision_topk( + self, + router_logits, + input_ids_global, + num_token_non_padded=num_token_non_padded, + ) + else: + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=dispatch_info, + **topk_kwargs, + ) + # Issued after the router so the main chain stays on the main stream at replay. + routed_pre_quant_input = pre_quant_input + if should_quant_routed_input_mxfp8: + with torch.cuda.stream(self.routed_quant_stream): + x_q, x_sf = self.experts.quant_method.quantize_routed_input( + hidden_states, routed_hidden_size(self.experts) + ) + ready = self.routed_quant_stream.record_event() + routed_pre_quant_input = Mxfp8RoutedInputPreQuant(x_q, x_sf, ready) + # The mHC post-split consumes the reduced row without an RMSNorm. + use_fused_finalize_all_reduce = ( + self._fuse_finalize_all_reduce + and has_shared_output + and hidden_states.shape[-1] == 5120 + and not self._shared_expert_tp1 + and self.tp_size > 1 + and hidden_states.shape[0] <= _FUSED_FINALIZE_ALL_REDUCE_MAX_TOKENS + and not should_skip_post_experts_all_reduce(is_tp_path=True) + and should_use_fuse_finalize_all_reduce( + self.experts, hidden_states.shape[0], hidden_states.shape[-1] ) - deferred_finalize = ( + ) + deferred_finalize = use_fused_finalize_all_reduce or ( has_shared_output and not self._shared_expert_tp1 and topk_output.format == TopKOutputFormat.BYPASSED @@ -988,13 +1061,13 @@ class DeepseekV2MoE(nn.Module): ) if deferred_finalize: final_hidden_states = self.experts.forward_deferred_finalize( - hidden_states, topk_output + hidden_states, topk_output, pre_quant_input=routed_pre_quant_input ) elif use_flashinfer_trtllm_bypass: final_hidden_states = self.experts.forward_impl(hidden_states, topk_output) - elif pre_quant_input is not None: + elif routed_pre_quant_input is not None: final_hidden_states = self.experts( - hidden_states, topk_output, pre_quant_input=pre_quant_input + hidden_states, topk_output, pre_quant_input=routed_pre_quant_input ) else: final_hidden_states = self.experts(hidden_states, topk_output) @@ -1007,6 +1080,7 @@ class DeepseekV2MoE(nn.Module): final_hidden_states *= self.routed_scaling_factor # Shared expert on alt stream, issued AFTER the main (routed) branch. See note above. + # Only the quant-once fp8 pair is shared with it; the routed MXFP8 pre-quant is not. with torch.cuda.stream(self.alt_stream): shared_output = self._forward_shared_experts( hidden_states, @@ -1014,17 +1088,84 @@ class DeepseekV2MoE(nn.Module): pre_quant_input=pre_quant_input, ) + # The routed-input pre-quant was already joined inside the routed MoE apply. current_stream.wait_stream(self.alt_stream) + all_reduce_done = False if deferred_finalize: from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( finalize_flashinfer_trtllm_deferred_output, ) - final_hidden_states = finalize_flashinfer_trtllm_deferred_output( - final_hidden_states, - shared_output, - ) + deferred = final_hidden_states + if ( + use_fused_finalize_all_reduce + and deferred.gemm2_out.shape[1] == hidden_states.shape[-1] + ): + from sglang.kernels.ops.communication.all_reduce_fusion import ( + moe_finalize_all_reduce, + ) + from sglang.srt.layers.moe.mhc_post_fusion import ( + current_mhc_post_fusion, + ) + + mhc = current_mhc_post_fusion() + if mhc is not None: + from sglang.kernels.ops.communication.all_reduce_mhc import ( + moe_finalize_all_reduce_mhc, + ) + + # Join the coefficients before the fused epilogue reads them. + mhc.materialize_stats() + if mhc.stats_stream is not None: + current_stream.wait_stream(mhc.stats_stream) + args = ( + deferred.gemm2_out, + deferred.expanded_idx_to_permuted_idx, + deferred.expert_weights, + deferred.top_k, + shared_output, + mhc.residual, + mhc.post, + mhc.comb, + ) + if mhc.norm_weight is not None: + from sglang.kernels.ops.communication.all_reduce_mhc import ( + moe_finalize_all_reduce_mhc_quant, + ) + + final_hidden_states, mhc.output, mhc.normalized, q, sf = ( + moe_finalize_all_reduce_mhc_quant( + *args, + mhc.pre, + mhc.norm_weight, + mhc.norm_eps, + world_size=self.tp_size, + ) + ) + mhc.quantized = (q, sf) + else: + final_hidden_states, mhc.output = moe_finalize_all_reduce_mhc( + *args, world_size=self.tp_size + ) + else: + final_hidden_states = moe_finalize_all_reduce( + deferred.gemm2_out, + deferred.expanded_idx_to_permuted_idx, + deferred.expert_weights, + deferred.top_k, + shared_output, + world_size=self.tp_size, + hidden_dim=hidden_states.shape[-1], + # Routing metadata must be ready before it is consumed. + prefetch_metadata=False, + ) + all_reduce_done = True + else: + final_hidden_states = finalize_flashinfer_trtllm_deferred_output( + deferred, + shared_output, + ) else: final_hidden_states = maybe_fuse_routed_scale_and_shared_add( self.experts, @@ -1033,7 +1174,8 @@ class DeepseekV2MoE(nn.Module): self.routed_scaling_factor, ) - final_hidden_states = post_experts_all_reduce(final_hidden_states) + if not all_reduce_done: + final_hidden_states = post_experts_all_reduce(final_hidden_states) # TP1 shared experts are replicated, so add them after all-reduce to # avoid summing the same shared output once per TP rank. if self._shared_expert_tp1: @@ -1088,13 +1230,21 @@ class DeepseekV2MoE(nn.Module): if getattr(self, "is_hash", False) else {} ) - topk_output = self.topk( - hidden_states, - router_logits, - num_token_non_padded=num_token_non_padded, - expert_location_dispatch_info=dispatch_info, - **topk_kwargs, - ) + if self.gate.e_score_correction_bias_vl is not None: + topk_output = vision_topk( + self, + router_logits, + input_ids_global, + num_token_non_padded=num_token_non_padded, + ) + else: + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=num_token_non_padded, + expert_location_dispatch_info=dispatch_info, + **topk_kwargs, + ) else: pre_quant_input = None shared_output = None @@ -1108,7 +1258,6 @@ class DeepseekV2MoE(nn.Module): def _pre_combine_hook( dispatcher: BaseDispatcher, combine_input: CombineInput ): - nonlocal shared_output self.alt_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.alt_stream): @@ -1341,7 +1490,6 @@ class DeepseekV2MoE(nn.Module): def _post_dispatch_hook( dispatcher: BaseDispatcher, dispatch_output: DispatchOutput ): - combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = ( compute_overlap_args(dispatch_output, self.alt_stream) ) @@ -1359,7 +1507,6 @@ class DeepseekV2MoE(nn.Module): def _pre_combine_hook( dispatcher: BaseDispatcher, combine_input: CombineInput ): - nonlocal shared_output if ( @@ -1397,7 +1544,6 @@ class DeepseekV2MoE(nn.Module): def _post_dispatch_hook( dispatcher: BaseDispatcher, dispatch_output: DispatchOutput ): - combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = ( compute_overlap_args(dispatch_output, self.alt_stream) ) @@ -1595,6 +1741,48 @@ class DeepseekV2MoE(nn.Module): q, s = sglang_per_token_group_quant_fp8_row_padded(hidden_states, 128) return q, s + @cached_property + def _routed_mxfp8_prequant_static_enabled(self) -> bool: + return self._compute_routed_mxfp8_prequant_enabled()[0] + + def _compute_routed_mxfp8_prequant_enabled(self) -> Tuple[bool, str]: + from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatcher + + if not _is_cuda: + return False, "not CUDA" + if self.routed_quant_stream is None: + return False, "no routed_quant_stream" + if self._enable_a2a_moe or self._fuse_shared_experts_inside_sbo: + return False, "a2a MoE or SBO shared-expert fusion" + if not get_moe_runner_backend().is_flashinfer_mxfp4(): + return False, "MoE runner backend not flashinfer_mxfp4" + experts = self.experts + if not isinstance(experts, FusedMoE): + return False, "experts not FusedMoE" + quant_method = experts.quant_method + if not isinstance(quant_method, Mxfp4FlashinferTrtllmMoEMethod): + return False, "experts quant method not Mxfp4FlashinferTrtllmMoEMethod" + if quant_method.flashinfer_mxfp4_moe_precision != "default": + return False, "flashinfer_mxfp4_moe_precision not default (no MXFP8 quant)" + # The pre-quant must describe exactly the tensor apply() receives: the + # standard dispatcher passes hidden_states through, the fp4 all-gather does not. + if not isinstance(experts.dispatcher, StandardDispatcher): + return False, "dispatcher not StandardDispatcher" + if should_use_flashinfer_cutlass_moe_fp4_allgather(): + return False, "flashinfer cutlass fp4 all-gather dispatch" + return True, "ok" + + def _should_quant_routed_input_mxfp8(self, hidden_states: torch.Tensor) -> bool: + return ( + # Capture-only: graph-pool tensors need no record_stream. + torch.cuda.is_current_stream_capturing() + # The piecewise TC graph's MoE op drops pre_quant_input. + and not is_in_tc_piecewise_cuda_graph() + and hidden_states.shape[0] > 0 + and hidden_states.dtype == torch.bfloat16 + and self._routed_mxfp8_prequant_static_enabled + ) + def op_gate(self, state): if state.hidden_states_mlp_input.shape[0] > 0: # router_logits: (num_tokens, n_experts) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 94bea84c3..40a261c6a 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -5,6 +5,7 @@ import functools import logging import time from contextlib import contextmanager, nullcontext +from types import SimpleNamespace from typing import ( TYPE_CHECKING, Any, @@ -29,7 +30,13 @@ from sglang.kernels.ops.attention.dsv4 import ( fused_rope_inplace, sglang_per_token_group_quant_fp8_dsv4_wo_a, ) +from sglang.kernels.ops.attention.dsv4.wo_a_bf16 import ( + wo_a_bf16_gemv, + wo_a_bf16_small_batch, + wo_a_bf16_small_batch_mxfp8, +) from sglang.kernels.ops.attention.flash_mla_sm120 import SM120_DECODE_MAX_TOKENS +from sglang.kernels.ops.layernorm.mhc_post_split_h import mhc_post_split_h from sglang.kernels.ops.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, ) @@ -59,6 +66,10 @@ from sglang.srt.layers.attention.dsa.utils import ( is_dsa_enable_prefill_cp, ) from sglang.srt.layers.attention.dsv4.compressor import Compressor +from sglang.srt.layers.attention.dsv4.dsv41_sparse import ( + DeepseekV41Compressor, + DeepseekV41Indexer, +) from sglang.srt.layers.attention.dsv4.indexer import C4Indexer from sglang.srt.layers.communicator import get_attn_tp_context from sglang.srt.layers.communicator_dsa_cp import ( @@ -69,6 +80,7 @@ from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx from sglang.srt.layers.cp.utils import ( cp_gather_full_sequence_states, cp_materialize_global_token_order, + is_cp_active, ) from sglang.srt.layers.dp_attention import ( _tbo_event, @@ -90,28 +102,40 @@ from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, is_dp_gatherv_active, ) +from sglang.srt.layers.engram import Engram, EngramHasher, EngramLayout from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear -from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.moe.utils import ( is_shared_experts_fusion_disabled, uses_per_rank_fused_shared_slots, ) +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, ) +from sglang.srt.layers.quantization.mxfp8_input import Mxfp8SwizzledInput from sglang.srt.layers.rotary_embedding import get_rope_wrapper from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding +from sglang.srt.managers.mm_utils import ( + MultiModalityDataPaddingPatternMultimodalTokens, + embed_mm_inputs, +) +from sglang.srt.managers.schedule_batch import MM_PAD_SHIFT_VALUE, MultimodalInputs from sglang.srt.mem_cache.memory_pool import RadixAttention from sglang.srt.model_executor.cuda_graph_config import ( Backend, Phase, check_cuda_graph_backend, ) -from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + PPProxyTensors, +) from sglang.srt.model_executor.forward_context import ( get_attn_backend, get_token_to_kv_pool, @@ -151,6 +175,12 @@ from sglang.srt.models.deepseek_v2 import ( _is_npu, _is_xpu, ) +from sglang.srt.models.deepseek_v41_vit import Aligner, ViT +from sglang.srt.multimodal.deepseek_v41_image_processing import ( + GPU_PLAN_KEY, + image_token_types, + materialize_image_gpu, +) from sglang.srt.runtime_context import ( get_device, get_exec, @@ -233,6 +263,18 @@ def _get_mhc_ops() -> MhcOps: logger = logging.getLogger(__name__) _FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get() + + +def wo_a_fp8_gemm_enabled(quant_config: Optional[QuantizationConfig]) -> bool: + """The fp8 wo_a absorb GEMM (DeepGEMM fp8_einsum, aiter mxscale) takes 128x128 + block scales only; any other layout dequantizes wo_a to bf16 at load.""" + return ( + _FP8_WO_A_GEMM + and isinstance(quant_config, Fp8Config) + and quant_config.weight_block_size == [128, 128] + ) + + _NPU_BF16_WO_A_GEMM = _is_npu and envs.SGLANG_OPT_NPU_BF16_WO_A_GEMM.get() _MHC_POST_MULT_VALUE = 2.0 _HC_PRENORM_DEEPGEMM_MIN_TOKENS = 1024 @@ -404,22 +446,61 @@ if _is_hip: def _apply_wo_a_bf16_matmul( - o: torch.Tensor, wo_a: torch.Tensor, is_decode: bool -) -> torch.Tensor: - """wo_a (attn output -> o_proj low-rank) bf16 batched matmul. - - ``o`` is ``[T, G, D]`` (tokens, groups, head_dim) and ``wo_a`` is - ``[G, R, D]`` (groups, o_lora_rank, head_dim); the result is ``[T, G, R]``. - - Dispatch contract: on the decode path, when the reroute is enabled - (``_wo_a_aiter_batched_gemm_enabled``, computed once at import) and has not - been disabled by a prior runtime failure, call the pre-imported aiter - ``batched_gemm_bf16`` (``Y[i] = X[i] @ W[i]^T``). Otherwise -- prefill, any - gate off, or after a failure -- use the numerically-equivalent - ``torch.einsum("tgd,grd->tgr", ...)``. The first runtime kernel failure - disables the reroute for the process (logged once). - """ + o: torch.Tensor, + wo_a: torch.Tensor, + is_decode: bool, + is_target_verify: bool = False, + fuse_mxfp8_quant: bool = False, + is_prefill: bool = False, + fast_path: bool = False, +) -> torch.Tensor | Mxfp8SwizzledInput: + # o [T, G, D] @ wo_a [G, R, D] -> [T, G, R]; the fast paths below are gated + # on the exact validated TP4 shapes and write token-major output directly. global _wo_a_aiter_batched_gemm_disabled + if ( + fast_path + and _is_cuda + and ( + ( + is_decode + and o.shape[0] == 1 + and (get_platform().is_blackwell or get_platform().is_sm90) + ) + or ( + is_target_verify + 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) + and o.dtype == wo_a.dtype == torch.bfloat16 + and o.stride(2) == 1 + and o.stride(1) == 4096 + and o.stride(0) >= 8192 + and wo_a.is_contiguous() + ): + if is_decode and o.shape[0] == 1: + return wo_a_bf16_gemv(o, wo_a) + if 2 <= o.shape[0] <= 8: + if fuse_mxfp8_quant: + return Mxfp8SwizzledInput(*wo_a_bf16_small_batch_mxfp8(o, wo_a)) + return wo_a_bf16_small_batch(o, wo_a) + result = torch.empty( + (o.shape[0], wo_a.shape[0], wo_a.shape[1]), dtype=o.dtype, device=o.device + ) + # The strided destination keeps the einsum reduction while producing the + # layout wo_b consumes; draft warmup/capture can enter with grad enabled. + with torch.no_grad(): + torch.bmm( + o.transpose(0, 1), wo_a.transpose(1, 2), out=result.transpose(0, 1) + ) + return result if ( is_decode and _wo_a_aiter_batched_gemm_enabled @@ -560,6 +641,7 @@ def _apply_gguf_grouped_wo_a( if TYPE_CHECKING: from sglang.srt.layers.attention.deepseek_v4_backend import ( DeepseekV4AttnBackend, + LateLayerTail, ) from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import ( DeepseekV4HipRadixBackend, @@ -616,6 +698,7 @@ def deepseek_v4_attention_with_output( ) output[:real_num_tokens].view(ret.shape).copy_(ret) + output[real_num_tokens:].zero_() return @@ -624,7 +707,37 @@ bcg_deepseek_v4_attention_with_output = eager_on_graph(True)( ) +def deepseek_v4_low_ratio_sources(layer, x, q_lora, positions) -> None: + # The compressor and prefill indexer sync with the host, like the attention. + forward_batch = get_tc_piecewise_forward_context().forward_batch + real_num_tokens = forward_batch.global_num_token_non_padded_cpu + if real_num_tokens == 0: + return + get_attn_backend().forward_low_ratio_sources( + layer=layer, + x=x[:real_num_tokens], + q_lora=q_lora[:real_num_tokens], + positions=positions[:real_num_tokens], + forward_batch=forward_batch, + ) + + +bcg_deepseek_v4_low_ratio_sources = eager_on_graph(True)(deepseek_v4_low_ratio_sources) + + +def deepseek_v4_engram_hash_ids(hasher, input_ids: torch.Tensor) -> torch.Tensor: + # The hasher reads per-request rows, so it cannot run inside the CUDA graph. + forward_batch = get_tc_piecewise_forward_context().forward_batch + return hasher(input_ids, forward_batch) + + +bcg_deepseek_v4_engram_hash_ids = eager_on_graph(True)(deepseek_v4_engram_hash_ids) + + class MqaAttentionBase(nn.Module): + # Class-level default for subclasses that read it without running __init__. + wo_a_fp8: bool = False + def __init__( self, config: DeepSeekV4Config, @@ -642,6 +755,7 @@ class MqaAttentionBase(nn.Module): rope_original_seq_len: Optional[int] = None, ) -> None: super().__init__() + self.is_dsv41 = getattr(config, "model_type", None) == "deepseek_v41" self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() if attn_tp_rank is None or attn_tp_size is None: attn_tp_rank = get_parallel().attn_tp_rank @@ -664,6 +778,7 @@ class MqaAttentionBase(nn.Module): self.o_lora_rank = config.o_lora_rank self.eps = config.rms_norm_eps self.softmax_scale = self.head_dim**-0.5 + self.q_head_norm = config.q_head_norm self.compress_ratio: int = ( compress_ratio @@ -672,9 +787,13 @@ class MqaAttentionBase(nn.Module): ) assert self.compress_ratio in ( 0, + 1, + 2, 4, 128, - ), f"V4 compress_ratio: expected one of (0, 4, 128), got {self.compress_ratio}" + ), ( + f"compress_ratio: expected one of (0, 1, 2, 4, 128), got {self.compress_ratio}" + ) assert self.head_dim == config.head_dim assert config.num_key_value_heads == 1 @@ -682,7 +801,9 @@ class MqaAttentionBase(nn.Module): fuse: bool = ( envs.SGLANG_OPT_FUSE_WQA_WKV.get() if fuse_wqa_wkv is None else fuse_wqa_wkv ) - fp8: bool = _FP8_WO_A_GEMM if wo_a_fp8 is None else wo_a_fp8 + fp8: bool = ( + wo_a_fp8_gemm_enabled(quant_config) if wo_a_fp8 is None else wo_a_fp8 + ) reduce_results: bool = ( (self.attn_tp_size == get_parallel().tp_size and self.attn_tp_size > 1) if wo_b_reduce_results is None @@ -705,6 +826,7 @@ class MqaAttentionBase(nn.Module): wo_a_quant_config = None self.fuse_wqa_wkv = fuse + self.wo_a_fp8 = fp8 self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32)) self._attn_sink_local: Optional[torch.Tensor] = None @@ -904,6 +1026,8 @@ class MqaAttentionBase(nn.Module): class MQALayer(MqaAttentionBase): + is_dsv41: bool = False + def __init__( self, config: DeepSeekV4Config, @@ -922,7 +1046,7 @@ class MQALayer(MqaAttentionBase): ) active_rope_scaling = None - if self.compress_ratio in (4, 128): + if self.compress_ratio: active_rope_scaling = dict(self.rope_scaling or {}) active_rope_scaling["rope_type"] = "deepseek_yarn" self.rotary_emb = get_rope_wrapper( @@ -1007,6 +1131,23 @@ class MQALayer(MqaAttentionBase): fp4_cos=(self.cos_cache[:, 0, 0, :] if _is_hip else None), fp4_sin=(self.sin_cache[:, 0, 0, :] if _is_hip else None), ) + elif self.compress_ratio in (1, 2): + # The layers in between read both through the attention backend. + if self.layer_id in config.kv_source_layer_ids: + self.compressor = DeepseekV41Compressor( + hidden_size=config.hidden_size, + head_dim=self.head_dim, + compress_ratio=self.compress_ratio, + eps=config.rms_norm_eps, + ) + if self.layer_id in config.index_source_layer_ids: + self.indexer = DeepseekV41Indexer( + config, + layer_id=self.layer_id, + head_dim=self.head_dim, + quant_config=quant_config, + prefix=add_prefix("indexer", prefix), + ) self.attn_mqa = RadixAttention( self.n_local_heads, @@ -1055,16 +1196,83 @@ class MQALayer(MqaAttentionBase): inverse=inverse, ) + def accepts_mxfp8_swizzled_input(self) -> bool: + """Whether the first projection consumes a 128x4 MXFP8 activation tuple.""" + cached = getattr(self, "_accepts_mxfp8_swizzled_input", None) + if cached is not None: + return cached + if self.fuse_wqa_wkv: + linears = [getattr(self, "wqkv_a", None)] + else: + # Both projections read the same activation on this path. + linears = [getattr(self, "wq_a", None), getattr(self, "wkv", None)] + + def _takes_swizzled(linear) -> bool: + method = getattr(linear, "quant_method", None) + return bool( + linear is not None + and getattr(method, "mxfp8_dense_backend", None) + in ( + Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL, + Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS, + ) + and ( + getattr(method, "use_mxfp8", False) + or getattr(linear, "block_fp8_mxfp8_ready", False) + ) + ) + + ok = all(_takes_swizzled(linear) for linear in linears) + self._accepts_mxfp8_swizzled_input = ok + return ok + + def _normalize_q_lora( + self, q: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor | Mxfp8SwizzledInput]: + # The indexer needs the BF16 normalized row; wq_b needs the quantized one. + method = self.wq_b.quant_method + if ( + _is_cuda + and self.is_dsv41 + and get_platform().is_blackwell + and q.dtype == self.q_norm.weight.dtype == torch.bfloat16 + and q.ndim == 2 + and 0 < q.shape[0] <= 8 + and q.shape[1] == 1280 + and q.stride(1) == 1 + and getattr(method, "mxfp8_dense_backend", None) + == Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL + and ( + getattr(method, "use_mxfp8", False) + or getattr(self.wq_b, "block_fp8_mxfp8_ready", False) + ) + ): + from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled + from sglang.srt.runtime_context import get_exec + + if not ( + is_batch_invariant_mode_enabled() + or get_exec().deterministic.enable_deterministic_inference + ): + from sglang.kernels.ops.layernorm.mxfp8_epilogue import rmsnorm_mxfp8 + + y, quant, scale = rmsnorm_mxfp8( + q, self.q_norm.weight, self.q_norm.variance_epsilon + ) + return y, Mxfp8SwizzledInput(quant, scale) + q = self.q_norm(q) + return q, q + def _compute_q_a( self, x: torch.Tensor, qkv_a: Optional[torch.Tensor] = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, torch.Tensor | Mxfp8SwizzledInput]: if qkv_a is not None: q = qkv_a[..., : self.q_lora_rank] else: q, _ = self.wq_a(x) - return self.q_norm(q) + return self._normalize_q_lora(q) def _compute_q_b( self, @@ -1074,6 +1282,39 @@ class MQALayer(MqaAttentionBase): ) -> torch.Tensor: q, _ = self.wq_b(q) q = q.view(-1, self.n_local_heads, self.head_dim) + if not self.q_head_norm: + if ( + _is_cuda + and q_out is not None + 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 + and q.stride(1) == q_out.stride(1) == 512 + and q.stride(2) == q_out.stride(2) == 1 + ): + from sglang.kernels.ops.attention.dsv4.q_rope_store import q_rope_store + + q_rope_store(q, q_out, self.freqs_cis, positions) + return q_out + fused_rope_inplace( + q[..., -self.qk_rope_head_dim :], + None, + self.freqs_cis, + positions=positions, + ) + if q_out is None: + return q + q_out.copy_(q) + return q_out if q_out is None: q_out = torch.empty_like(q) # Fused warp-per-(token, head) rmsnorm-self + RoPE + write to q_out. @@ -1168,7 +1409,7 @@ class MQALayer(MqaAttentionBase): qkv_a, _ = self.wqkv_a(x_linear) qkv_a_ready = current_stream.record_event() - q_lora = self._compute_q_a(x_linear, qkv_a=qkv_a) + q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a) q_lora_ready = current_stream.record_event() if self.indexer is not None: @@ -1196,7 +1437,7 @@ class MQALayer(MqaAttentionBase): x, forward_batch, self.layer_id, self.compressor ) - q = self._compute_q_b(q_lora, positions, q_out) + q = self._compute_q_b(q_for_wqb, positions, q_out) current_stream.wait_stream(stream_kv) current_stream.wait_stream(stream_compressor) current_stream.wait_stream(stream_indexer) @@ -1204,6 +1445,70 @@ class MQALayer(MqaAttentionBase): return q + def _forward_prepare_low_ratio_multi_stream( + self, + x: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + attn_backend, + q_out: Optional[torch.Tensor] = None, + x_quant=None, + ) -> torch.Tensor: + # Both side streams are joined before returning, and nothing they read is + # released before the join. + assert self.alt_streams is not None + current_stream = torch.cuda.current_stream() + stream_kv = self.alt_streams[0] + stream_sources = self.alt_streams[-1] + x_linear = x_quant if x_quant is not None else x + + # NOTE: wait for x ready + if self.compressor is not None: + stream_sources.wait_stream(current_stream) + qkv_a: Optional[torch.Tensor] = None + if self.fuse_wqa_wkv: + qkv_a, _ = self.wqkv_a(x_linear) + + if self.compressor is not None: + with torch.cuda.stream(stream_sources): + attn_backend.forward_low_ratio_sources( + layer=self, + x=x, + q_lora=None, + positions=positions, + forward_batch=forward_batch, + run_indexer=False, + ) + + stream_kv.wait_stream(current_stream) + q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a) + # NOTE: wait for the q_lora ready + if self.indexer is not None: + stream_sources.wait_stream(current_stream) + + q = self._compute_q_b(q_for_wqb, positions, q_out) + if self.indexer is not None: + # Forked above, right after q_lora; recorded here, after the Q chain. + with torch.cuda.stream(stream_sources): + attn_backend.forward_low_ratio_sources( + layer=self, + x=x, + q_lora=q_lora, + positions=positions, + forward_batch=forward_batch, + run_compressor=False, + ) + + with torch.cuda.stream(stream_kv): + self._compute_kv_to_cache( + x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a + ) + + current_stream.wait_stream(stream_kv) + if self.compressor is not None or self.indexer is not None: + current_stream.wait_stream(stream_sources) + return q + def _forward_prepare_multi_stream_npu( self, x: torch.Tensor, @@ -1373,7 +1678,7 @@ class MQALayer(MqaAttentionBase): token_to_kv_pool = get_token_to_kv_pool() swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch) swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id) - swa_page_size = token_to_kv_pool.swa_kv_pool.page_size + swa_page_size = token_to_kv_pool.swa_page_size q = fused_qk_norm_rope_swa_store( q=q, @@ -1430,7 +1735,6 @@ class MQALayer(MqaAttentionBase): k_rope_out: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: x_linear = x_quant if x_quant is not None else x - if self.fuse_wqa_wkv: qkv_a, _ = self.wqkv_a(x_linear) q_lora = qkv_a[..., : self.q_lora_rank] @@ -1498,8 +1802,8 @@ class MQALayer(MqaAttentionBase): ) q, _ = self.wq_b(q_for_wqb) else: - q_lora = self.q_norm(q_lora) - q, _ = self.wq_b(q_lora) + q_lora, q_for_wqb = self._normalize_q_lora(q_lora) + q, _ = self.wq_b(q_for_wqb) kv = ( qkv_a[..., self.q_lora_rank :] @@ -1555,10 +1859,10 @@ class MQALayer(MqaAttentionBase): # kv stays the strided slice of qkv_a -- the group-quant # kernel takes the row stride as an argument. else: - swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id) swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch) + swa_cache = token_to_kv_pool.get_swa_raw_buffer(self.layer_id) swa_page_size, bf16_store = ( - token_to_kv_pool.swa_kv_pool.page_size, + token_to_kv_pool.swa_page_size, False, ) @@ -1651,15 +1955,13 @@ class MQALayer(MqaAttentionBase): if q_out is not None: q_out.copy_(q) else: - q_lora = self.q_norm(q_lora) - q = self._compute_q_b(q_lora, positions, q_out) + q_lora, q_for_wqb = self._normalize_q_lora(q_lora) + q = self._compute_q_b(q_for_wqb, positions, q_out) if unified: # unified_kv prefill: keep bf16 kv; the backend writes # the ring AFTER attention (2-source path). kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a) - elif use_cp: - # NSA CP: keep bf16 kv around for the cross-rank all-gather, then - # write to the FlashMLA cache after gather. + elif use_cp and not self.is_dsv41: kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a) kv = cp_materialize_global_token_order( kv.contiguous(), @@ -1671,6 +1973,33 @@ class MQALayer(MqaAttentionBase): swa_k=kv, forward_batch=forward_batch, ) + elif use_cp: + # every rank writes the whole chunk's window KV with the fused fp32 store + if qkv_a is not None: + kv = qkv_a[..., self.q_lora_rank :] + else: + kv, _ = self.wkv(x_linear) + kv = cp_materialize_global_token_order( + kv.contiguous(), + forward_batch, + torch.cuda.current_stream(), + ) + tail = attn_backend.forward_metadata.late_layer_tail + global_positions = ( + tail.pos_global + if tail is not None + else forward_batch.positions[: kv.shape[0]] + ) + get_token_to_kv_pool().set_swa_key_buffer_radix_fused_norm_rope( + layer_id=self.layer_id, + swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch), + kv=kv, + kv_weight=self.kv_norm.weight.data, + eps=self.eps, + freqs_cis=self.freqs_cis, + positions=global_positions, + ) + kv = None else: self._compute_kv_to_cache( x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a @@ -1679,46 +2008,64 @@ class MQALayer(MqaAttentionBase): del qkv_a - use_npu_cp_full_metadata = use_cp and _is_npu - if self.indexer is not None: - if use_npu_cp_full_metadata: - with attn_backend.use_dsv4_cp_full_metadata(forward_batch): - attn_backend.forward_indexer_compressor( - x, - forward_batch, - self.indexer.layer_id, - self.indexer.compressor, - ) - self.indexer( - x=x, - q_lora=q_lora, - forward_batch=forward_batch, - attn_backend=attn_backend, - skip_compressor=True, - ) + if self.compress_ratio in (1, 2) and ( + self.compressor is not None or self.indexer is not None + ): + if ( + forward_batch.forward_mode.is_extend() + and is_in_breakable_cuda_graph() + and not getattr(attn_backend, "low_ratio_prefill_graph", False) + ): + bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions) else: - self.indexer( + attn_backend.forward_low_ratio_sources( + layer=self, x=x, q_lora=q_lora, + positions=positions, forward_batch=forward_batch, - attn_backend=attn_backend, ) - if self.compressor is not None: - if use_npu_cp_full_metadata: - with attn_backend.use_dsv4_cp_full_metadata(forward_batch): + else: + use_npu_cp_full_metadata = use_cp and _is_npu + if self.indexer is not None: + if use_npu_cp_full_metadata: + with attn_backend.use_dsv4_cp_full_metadata(forward_batch): + attn_backend.forward_indexer_compressor( + x, + forward_batch, + self.indexer.layer_id, + self.indexer.compressor, + ) + self.indexer( + x=x, + q_lora=q_lora, + forward_batch=forward_batch, + attn_backend=attn_backend, + skip_compressor=True, + ) + else: + self.indexer( + x=x, + q_lora=q_lora, + forward_batch=forward_batch, + attn_backend=attn_backend, + ) + if self.compressor is not None: + if use_npu_cp_full_metadata: + with attn_backend.use_dsv4_cp_full_metadata(forward_batch): + attn_backend.forward_core_compressor( + x, + forward_batch, + self.layer_id, + self.compressor, + ) + else: attn_backend.forward_core_compressor( x, forward_batch, self.layer_id, self.compressor, ) - else: - attn_backend.forward_core_compressor( - x, - forward_batch, - self.layer_id, - self.compressor, - ) return q, kv @@ -1749,6 +2096,7 @@ class MQALayer(MqaAttentionBase): ) and not (self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)) and not (_is_hip and self.compressor is None) + and self.compress_ratio not in (1, 2) ) or ( _is_npu and envs.SGLANG_NPU_USE_MULTI_STREAM.get() @@ -1757,6 +2105,21 @@ class MQALayer(MqaAttentionBase): and not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed() ) + low_ratio_multi_stream = ( + _is_cuda + and get_platform().is_blackwell + and self.compress_ratio in (1, 2) + and self.alt_streams is not None + and ( + forward_batch.forward_mode.is_decode() + or ( + forward_batch.forward_mode.is_target_verify() + # Other MXFP8 backends may share mutable GEMM workspace. + and getattr(self.wq_b.quant_method, "mxfp8_dense_backend", None) + == Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL + ) + ) + ) from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_fp8, is_unified_kv_triton, @@ -1838,11 +2201,27 @@ class MQALayer(MqaAttentionBase): if kernel_num_heads != self.n_local_heads: # Backends without an exact-head specialization retain the existing # padded shape. attn_sink is sliced to this rank and padded to match. - # Only [0:n_local_heads] is written below. Uninitialized padded TP - # heads inject NaN into attention on gfx942 (fnuz), so zero-init - # there; other archs tolerate new_empty and skip the per-forward - # memset. - if _is_gfx942_supported: + if self.is_dsv41: + # V4.1 kernels read all padded heads, so the padding must be + # zero. The buffer is reused per layer; no consumer may keep it. + want = (x.shape[0], kernel_num_heads, self.head_dim) + meta = getattr(attn_backend, "forward_metadata", None) + q_padded = getattr(meta, "q_pad_buffer", None) + if ( + q_padded is None + or tuple(q_padded.shape) != want + or q_padded.dtype != x.dtype + ): + q_padded = x.new_zeros(*want) + if meta is not None: + try: + meta.q_pad_buffer = q_padded + except (AttributeError, TypeError): + pass + elif _is_gfx942_supported: + # Uninitialized padded TP heads inject NaN into attention on gfx942 + # (fnuz), so zero-init there; other archs tolerate new_empty and skip + # the per-forward memset. q_padded = x.new_zeros(x.shape[0], kernel_num_heads, self.head_dim) else: q_padded = x.new_empty(x.shape[0], kernel_num_heads, self.head_dim) @@ -1881,6 +2260,16 @@ class MQALayer(MqaAttentionBase): x_quant=x_quant, ) kv = None + elif low_ratio_multi_stream: + q = self._forward_prepare_low_ratio_multi_stream( + x, + positions, + forward_batch, + attn_backend, + q_out, + x_quant=x_quant, + ) + kv = None else: q, kv = self._forward_prepare( x, @@ -1952,7 +2341,7 @@ class MQALayer(MqaAttentionBase): ) o = o[:, tp_slice, :] if ( - _FP8_WO_A_GEMM + self.wo_a_fp8 and _wo_a_fp8_mxscale_fused_invrope is not None and not _is_npu ): @@ -2010,7 +2399,7 @@ class MQALayer(MqaAttentionBase): perm_x2=(0, 1, 2), perm_y=(1, 0, 2), ) - elif _FP8_WO_A_GEMM and _wo_a_fp8_mxscale is not None: + elif self.wo_a_fp8 and _wo_a_fp8_mxscale is not None: # ROCm gfx950: same fp8 absorb GEMM as the DeepGEMM path below, # but through aiter's e8m0 block-scale batched GEMM. The # activation is quantized per token-group inside the helper. @@ -2020,7 +2409,7 @@ class MQALayer(MqaAttentionBase): self.wo_a.weight.view(G, self.o_lora_rank, D), self.wo_a.weight_scale_inv.data, ) - elif _FP8_WO_A_GEMM: + elif self.wo_a_fp8: import deep_gemm from sglang.srt.layers import deep_gemm_wrapper @@ -2071,7 +2460,31 @@ class MQALayer(MqaAttentionBase): self.n_local_groups, self.o_lora_rank, -1 ) o = _apply_wo_a_bf16_matmul( - o, wo_a, is_decode=forward_batch.forward_mode.is_decode() + o, + 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(), + fast_path=self.is_dsv41, + fuse_mxfp8_quant=( + not get_forward().sp_active + and getattr( + getattr(self.wo_b, "quant_method", None), + "mxfp8_dense_backend", + None, + ) + == Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL + and ( + getattr( + getattr(self.wo_b, "quant_method", None), + "use_mxfp8", + False, + ) + or getattr( + self.wo_b, "block_fp8_mxfp8_ready", False + ) + ) + ), ) else: o = _apply_gguf_grouped_wo_a( @@ -2081,7 +2494,31 @@ class MQALayer(MqaAttentionBase): self.o_lora_rank, ) - o, _ = self.wo_b(o.flatten(1)) + from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion + + mhc = current_mhc_post_fusion() + o, _ = self.wo_b( + o if isinstance(o, Mxfp8SwizzledInput) else o.flatten(1), + skip_all_reduce=mhc is not None, + ) + if mhc is not None: + from sglang.kernels.ops.communication.all_reduce_mhc import ( + all_reduce_mhc_norm, + ) + + mhc.materialize_stats() + if mhc.stats_stream is not None: + torch.cuda.current_stream().wait_stream(mhc.stats_stream) + o, mhc.output, mhc.normalized = all_reduce_mhc_norm( + o, + mhc.residual, + mhc.post, + mhc.comb, + mhc.pre, + mhc.norm_weight, + mhc.norm_eps, + world_size=self.attn_tp_size, + ) if self.attn_tp_size > 1 and self.attn_tp_size < get_parallel().tp_size: o = attn_tp_all_reduce(o) @@ -2103,6 +2540,25 @@ class MQALayer(MqaAttentionBase): ) +@contextmanager +def _every_row_routed(forward_batch: ForwardBatch, num_rows: int): + # Under CP the real rows are not a prefix, so every row must be routed. + saved = ( + forward_batch.num_token_non_padded, + forward_batch.global_num_token_non_padded_cpu, + ) + if saved[0] is not None: + forward_batch.num_token_non_padded = torch.full_like(saved[0], num_rows) + forward_batch.global_num_token_non_padded_cpu = num_rows + try: + yield + finally: + ( + forward_batch.num_token_non_padded, + forward_batch.global_num_token_non_padded_cpu, + ) = saved + + class DeepseekV4DecoderLayer(nn.Module): def __init__( self, @@ -2114,8 +2570,12 @@ class DeepseekV4DecoderLayer(nn.Module): prefix: str = "", alt_streams: Optional[List[torch.cuda.Stream]] = None, compress_ratio_override: Optional[int] = None, + engram_layout: Optional[EngramLayout] = None, + hc_stats_stream: Optional[torch.cuda.Stream] = None, + moe_routed_quant_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() + self.hc_stats_stream = hc_stats_stream self.config = config self.hidden_size = config.hidden_size self.layer_id = layer_id @@ -2145,8 +2605,11 @@ class DeepseekV4DecoderLayer(nn.Module): prefix=add_prefix("mlp", prefix), layer_id=self.layer_id, alt_stream=moe_alt_stream, + routed_quant_stream=moe_routed_quant_stream, is_nextn=is_nextn, is_deepseek_v4=True, + vl_correction_bias=config.model_type == "deepseek_v41" + and config.vision_n_layers > 0, ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -2170,6 +2633,19 @@ class DeepseekV4DecoderLayer(nn.Module): self.use_fused_mhc_post_pre = ( is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu() ) + # The fused post+pre boundary bakes in the same-sublayer pre-mix. + self.hc_pre_from_prev_sublayer = config.hc_pre_from_prev_sublayer + if self.hc_pre_from_prev_sublayer: + self.use_fused_mhc_post_pre = False + self.engram = None + if engram_layout is not None and layer_id in engram_layout.layer_ids: + self.engram = Engram( + config, + layer_id, + engram_layout, + quant_config=quant_config, + prefix=add_prefix("engram", prefix), + ) self._input_layernorm_weight_bf16 = None self._post_attention_layernorm_weight_bf16 = None @@ -2201,6 +2677,46 @@ class DeepseekV4DecoderLayer(nn.Module): self.post_attention_layernorm.weight.data.bfloat16().contiguous() ) + from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled + + # The original FP32 parameters stay 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 getattr(self.config, "model_type", None) == "deepseek_v41" + and not is_batch_invariant_mode_enabled() + ): + from sglang.kernels.ops.layernorm.mhc import ( + split_tf32_hc_weight, + ) + from sglang.srt.layers.deep_gemm_wrapper.configurer import ( + ENABLE_JIT_DEEPGEMM, + ) + + if ENABLE_JIT_DEEPGEMM: + import deep_gemm + + if not callable(getattr(deep_gemm, "tf32_hc_prenorm_gemm", None)): + return + 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.mhc 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, x: torch.Tensor, @@ -2368,7 +2884,6 @@ class DeepseekV4DecoderLayer(nn.Module): post: torch.Tensor, comb: torch.Tensor, ): - if x.shape[0] == 0: return torch.empty( (0, self.hc_mult, x.shape[-1]), dtype=x.dtype, device=x.device @@ -2389,12 +2904,40 @@ class DeepseekV4DecoderLayer(nn.Module): if _is_xpu: return _get_mhc_ops().mhc_post(x, residual, post, comb) + if ( + _is_cuda + and get_platform().is_blackwell + and self.hc_pre_from_prev_sublayer + and self.hc_mult == 4 + and x.shape[1] == 5120 + and x.shape[0] <= 384 + and x.dtype == residual.dtype == torch.bfloat16 + and post.dtype == comb.dtype == torch.float32 + and all(t.is_contiguous() for t in (x, residual, post, comb)) + ): + return mhc_post_split_h(x, residual, post, comb) + if envs.SGLANG_OPT_USE_FLASHINFER_MHC.get(): from flashinfer.mhc import mhc_post return mhc_post(x, residual, post, comb) if envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get(): + if ( + self.hc_pre_from_prev_sublayer + and get_platform().is_sm90 + and x.is_cuda + and 1 <= x.shape[0] <= 64 + and x.shape[1] == 5120 + and residual.shape == (x.shape[0], 4, 5120) + and x.dtype == residual.dtype == torch.bfloat16 + and post.dtype == comb.dtype == torch.float32 + and post.numel() == x.shape[0] * 4 + and comb.shape == (x.shape[0], 4, 4) + and all(t.is_contiguous() for t in (x, residual, post, comb)) + ): + return mhc_post_split_h(x, residual, post, comb) + from sglang.kernels.ops.layernorm.mhc import mhc_post return mhc_post(x, residual, post, comb) @@ -2604,6 +3147,390 @@ class DeepseekV4DecoderLayer(nn.Module): # cross-layer fusion, and the final layer is completed in DeepseekV4Model. return hidden_states, residual, post, comb + def _hc_combine( + self, + x: torch.Tensor, + apply_pre: Optional[torch.Tensor], + norm: RMSNorm, + stats_stream: Optional[torch.cuda.Stream] = None, + quantized: Optional[list] = None, + normalized: Optional[torch.Tensor] = None, + precomputed: Optional[tuple] = None, + ) -> torch.Tensor: + from sglang.kernels.ops.layernorm.mhc import hc_combine + + quantize = quantized is not None + x_flat = x.flatten(1) + tiny = 0 < x.shape[0] <= 8 + if stats_stream is not None and not tiny: + stats_stream.wait_stream(torch.cuda.current_stream()) + + def combine_and_norm(): + if precomputed is not None: + assert quantized is not None + quantized.append(precomputed[1]) + return precomputed[0] + if normalized is not None: + assert not quantize + return normalized + if apply_pre is None: + return norm(x[:, 0, :].contiguous()) + from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled + + if ( + x.is_cuda + and get_platform().is_blackwell + 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 + and apply_pre.stride(1) == 1 + and not norm.cast_x_before_out_mul + and norm.variance_size_override is None + and not is_batch_invariant_mode_enabled() + ): + # The fused scale writer supports the small decode/verify tile only. + if quantize and x.shape[0] <= 8: + from sglang.kernels.ops.layernorm.hc_combine_norm import ( + hc_combine_norm_mxfp8, + ) + + y, y_q, y_sf = hc_combine_norm_mxfp8( + x_flat, apply_pre, norm.weight, norm.variance_epsilon + ) + quantized.append(Mxfp8SwizzledInput(y_q, y_sf)) + return y + from sglang.kernels.ops.layernorm.hc_combine_norm import hc_combine_norm + + return hc_combine_norm( + x_flat, apply_pre, norm.weight, norm.variance_epsilon + ) + return norm(hc_combine(x_flat, apply_pre, self.hc_mult, x.dtype)) + + y = combine_and_norm() + if stats_stream is not None and tiny: + stats_stream.wait_stream(torch.cuda.current_stream()) + return y + + def _hc_mix_stats( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + stats_stream: Optional[torch.cuda.Stream] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + from sglang.kernels.ops.layernorm.mhc import hc_mix_stats, hc_mix_stats_sinkhorn + + x_flat = x.flatten(1) + + if ( + x.is_cuda + and torch.version.cuda is not None + and ( + get_platform().is_blackwell + or (get_platform().is_sm90 and x.shape[0] == 1) + ) + and x.dtype == torch.bfloat16 + ): + # Fusing the split-K reduction with sinkhorn keeps it batch-invariant. + main_stream = torch.cuda.current_stream() + if stats_stream is not None: + x.record_stream(stats_stream) + with ( + torch.cuda.stream(stats_stream) + if stats_stream is not None + else nullcontext() + ): + 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.mhc 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.mhc 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: + # Allocated on the side stream, read on the main stream after the join. + for coefficient in (pre, post, comb): + coefficient.record_stream(main_stream) + return pre, post, comb + if x.is_cuda and torch.version.cuda is not None: + # cuBLAS/torch reductions can change order with num_tokens; this kernel + # keeps the mixing and RMS reductions batch-invariant. + mixes = hc_mix_stats(x_flat, hc_fn, self.rms_norm_eps).unsqueeze(1) + else: + x_flat = x_flat.float() + rsqrt = torch.rsqrt( + x_flat.square().mean(-1, keepdim=True) + self.rms_norm_eps + ) + mixes = (F.linear(x_flat, hc_fn) * rsqrt).unsqueeze(1) + pre, post, comb = _get_mhc_ops().hc_split_sinkhorn( + mixes, + hc_scale, + hc_base, + self.hc_mult, + self.hc_sinkhorn_iters, + self.hc_eps, + ) + return pre.squeeze(1), post.squeeze(1), comb.squeeze(1) + + def _hc_mix_and_combine( + self, + x, + hc_fn, + hc_scale, + hc_base, + apply_pre, + norm, + stats_stream=None, + quantized=None, + normalized=None, + precomputed=None, + ): + y = DeepseekV4DecoderLayer._hc_combine( + self, x, apply_pre, norm, stats_stream, quantized, normalized, precomputed + ) + return ( + y, + *DeepseekV4DecoderLayer._hc_mix_stats( + self, x, hc_fn, hc_scale, hc_base, stats_stream + ), + ) + + def _get_hc_stats_stream(self, hidden_states, forward_batch): + # Every branch joins this stream before hc_post reads the coefficients. + return ( + self.hc_stats_stream + if ( + forward_batch.forward_mode.is_decode() + or ( + forward_batch.forward_mode.is_target_verify() + and hidden_states.shape[0] > 0 + ) + ) + and (not get_platform().is_sm90 or hidden_states.shape[0] == 1) + else None + ) + + def forward_hc_pre_from_prev( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + forward_batch: ForwardBatch, + input_ids_global: torch.Tensor, + prev_pre: Optional[torch.Tensor], + precomputed_attn: Optional[tuple] = None, + next_norm: Optional[RMSNorm] = None, + next_input: Optional[list] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Layer forward where each sublayer consumes the previous sublayer's + pre-mix. Returns (hidden_states, ffn_pre).""" + from functools import partial + + stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch) + residual = hidden_states + attn_quantized: Optional[list] = ( + [] if self.self_attn.accepts_mxfp8_swizzled_input() else None + ) + attn_stats = partial( + self._hc_mix_stats, + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + stats_stream, + ) + x = self._hc_combine( + hidden_states, + apply_pre=prev_pre, + norm=self.input_layernorm, + stats_stream=stats_stream, + quantized=attn_quantized, + precomputed=precomputed_attn, + ) + attn_mhc = None + if ( + self.config.model_type == "deepseek_v41" + and x.is_cuda + and get_platform().is_blackwell + and 0 < x.shape[0] <= 8 + and x.shape[1] == 5120 + and self.hc_mult == 4 + and x.dtype == residual.dtype == torch.bfloat16 + and residual.is_contiguous() + and get_parallel().attn_dp_size == 1 + and get_parallel().tp_size == self.self_attn.attn_tp_size == 4 + and self.self_attn.wo_b.reduce_results + and not get_forward().sp_active + and not self.dsa_enable_prefill_cp + and not self.post_attention_layernorm.cast_x_before_out_mul + and self.post_attention_layernorm.variance_size_override is None + and self.post_attention_layernorm.weight.dtype == torch.bfloat16 + ): + from sglang.kernels.ops.communication.all_reduce_fusion import ( + get_registered_comm, + ) + from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled + from sglang.srt.layers.moe.mhc_post_fusion import ( + MhcPostFusion, + use_mhc_post_fusion, + ) + + if ( + not is_batch_invariant_mode_enabled() + and get_registered_comm(self.self_attn.attn_tp_size) is not None + ): + attn_mhc = MhcPostFusion( + residual, + None, + None, + stats_stream, + record_stats=attn_stats, + norm_weight=self.post_attention_layernorm.weight, + norm_eps=self.post_attention_layernorm.variance_epsilon, + ) + context = ( + use_mhc_post_fusion(attn_mhc) if attn_mhc is not None else nullcontext() + ) + with context, self.self_attn.maybe_use_decode_attn_tp(forward_batch): + x = self.self_attn( + x=x, + positions=positions, + forward_batch=forward_batch, + x_quant=attn_quantized[0] if attn_quantized else None, + ) + if attn_mhc is not None: + attn_mhc.materialize_stats() + attn_pre = attn_mhc.pre + hidden_states = attn_mhc.output + else: + attn_pre, attn_post, attn_comb = attn_stats() + if stats_stream is not None: + torch.cuda.current_stream().wait_stream(stats_stream) + hidden_states = self.hc_post(x, residual, attn_post, attn_comb) + + residual = hidden_states + ffn_stats = partial( + self._hc_mix_stats, + hidden_states, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + stats_stream, + ) + x = self._hc_combine( + hidden_states, + apply_pre=attn_pre, + norm=self.post_attention_layernorm, + stats_stream=stats_stream, + normalized=attn_mhc.normalized if attn_mhc is not None else None, + ) + mhc = None + if ( + self.config.model_type == "deepseek_v41" + and x.is_cuda + and get_platform().is_blackwell + and 0 < x.shape[0] <= 8 + and x.shape[1] == 5120 + and self.hc_mult == 4 + and x.dtype == residual.dtype == torch.bfloat16 + and residual.is_contiguous() + and get_parallel().attn_dp_size == 1 + and get_moe_a2a_backend().is_none() + and not self.dsa_enable_prefill_cp + and not self.mlp._shared_expert_tp1 + and self.mlp.tp_size == 4 + ): + from sglang.srt.layers.moe.mhc_post_fusion import ( + MhcPostFusion, + use_mhc_post_fusion, + ) + + mhc = MhcPostFusion( + residual, None, None, stats_stream, record_stats=ffn_stats + ) + if next_norm is not None: + mhc.norm_weight = next_norm.weight + mhc.norm_eps = next_norm.variance_epsilon + context = use_mhc_post_fusion(mhc) + else: + context = nullcontext() + with context: + x = self._run_moe_ffn_dp_sync( + x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global + ) + if mhc is not None: + mhc.materialize_stats() + ffn_pre, ffn_post, ffn_comb = mhc.pre, mhc.post, mhc.comb + else: + ffn_pre, ffn_post, ffn_comb = ffn_stats() + if mhc is not None and mhc.output is not None: + hidden_states = mhc.output + if next_input is not None and mhc.quantized is not None: + next_input.append((mhc.normalized, Mxfp8SwizzledInput(*mhc.quantized))) + else: + if stats_stream is not None: + torch.cuda.current_stream().wait_stream(stats_stream) + hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb) + return hidden_states, ffn_pre + def _run_moe_ffn_dp_sync( self, hidden_states: torch.Tensor, @@ -2706,14 +3633,30 @@ class DeepseekV4DecoderLayer(nn.Module): # Skip the MoE-internal post-experts all_reduce when we will do the # reduce via reduce_scatterv/reduce_scatter at the combine below # (else double-reduce). - with get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter): - hidden_states = self.mlp( - hidden_states, - forward_batch, - input_ids=input_ids, - input_ids_global=input_ids_global, - skip_shared_experts=_do_shared_local, - ) + gathered_rows = ( + _every_row_routed(forward_batch, hidden_states.shape[0]) + if _use_cp and get_moe_a2a_backend().is_none() + else nullcontext() + ) + # The MoE sees DP-gathered rows, so this rank's local count cannot mask them. + # The standard dispatcher masks padding in the gathered buffer. + saved_num_token_non_padded = forward_batch.num_token_non_padded + if _use_tp_moe_gather: + forward_batch.num_token_non_padded = None + try: + with ( + get_forward().scoped(mlp_reduce_scatter=mlp_reduce_scatter), + gathered_rows, + ): + hidden_states = self.mlp( + hidden_states, + forward_batch, + input_ids=input_ids, + input_ids_global=input_ids_global, + skip_shared_experts=_do_shared_local, + ) + finally: + forward_batch.num_token_non_padded = saved_num_token_non_padded if _use_cp and get_moe_a2a_backend().is_none(): hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) elif _use_tp_moe_gather: @@ -3021,6 +3964,18 @@ class DeepseekV4DecoderLayer(nn.Module): state.hidden_states_mlp_output = hidden +def _scatter_tail_rows( + tail: LateLayerTail, rows: torch.Tensor, num_tokens: int +) -> torch.Tensor: + # Rows outside the tail are never read (see _check_late_layer_tail_readers). + full = rows.new_empty((num_tokens, rows.shape[1])) + if tail.contiguous_start is not None: + full[tail.contiguous_start :].copy_(rows) + else: + full[tail.token_indices] = rows[: tail.token_indices.shape[0]] + return full + + class DeepseekV4Model(nn.Module): fall_back_to_pt_during_load = False @@ -3031,6 +3986,7 @@ class DeepseekV4Model(nn.Module): prefix: str = "", ) -> None: super().__init__() + self.config = config self.pp_group = get_pp_group() self.hidden_size = config.hidden_size if self.pp_group.is_first_rank: @@ -3067,6 +4023,26 @@ class DeepseekV4Model(nn.Module): if use_stream_pool else None ) + # Routed-MoE input pre-quant, separate from the attention/indexer and + # shared-expert streams. + self.moe_routed_quant_stream = ( + device_module.Stream() + if _is_cuda + and config.hc_pre_from_prev_sublayer + and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() + else None + ) + # One shared stream for all layers; every sublayer joins it before + # reusing its residual. + self.hc_stats_stream = ( + device_module.Stream() + if _is_cuda + and (get_platform().is_blackwell or get_platform().is_sm90) + and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() + and config.hc_pre_from_prev_sublayer + else None + ) + self.engram_layout = EngramLayout.from_config(config) self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: DeepseekV4DecoderLayer( @@ -3075,6 +4051,9 @@ class DeepseekV4Model(nn.Module): quant_config=quant_config, prefix=prefix, alt_streams=self.alt_streams, + engram_layout=self.engram_layout, + hc_stats_stream=self.hc_stats_stream, + moe_routed_quant_stream=self.moe_routed_quant_stream, ), pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, @@ -3088,12 +4067,26 @@ class DeepseekV4Model(nn.Module): self.hc_eps = config.hc_eps self.hc_mult = hc_mult = config.hc_mult self.norm_eps = config.rms_norm_eps - if self.pp_group.is_last_rank: + self.hc_pre_from_prev_sublayer = config.hc_pre_from_prev_sublayer + self.hc_head_fn = self.hc_head_base = self.hc_head_scale = None + if self.pp_group.is_last_rank and not self.hc_pre_from_prev_sublayer: ( self.hc_head_fn, self.hc_head_base, self.hc_head_scale, ) = make_hc_head_params(hc_mult, config.hidden_size) + self.engram_hasher = None + if self.engram_layout is not None: + self.engram_hasher = EngramHasher.from_config( + config, + self.engram_layout, + image_token_id=( + config.image_token_id + if config.model_type == "deepseek_v41" + and config.vision_n_layers > 0 + else None + ), + ) self.use_fused_mhc_post_pre = ( is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu() @@ -3101,6 +4094,22 @@ class DeepseekV4Model(nn.Module): self.dspark_layers_to_capture: Optional[List[int]] = None + # Decoder SWA bounded replay: layers past the last kv_source layer run over + # each request's last SWA_WINDOW extend tokens only. + self.late_layer_start: Optional[int] = None + if get_exec().features.enable_decoder_swa_bounded_replay: + assert config.kv_source_layer_ids, ( + "decoder SWA bounded replay needs kv_source_layer_ids" + ) + self.late_layer_start = max(config.kv_source_layer_ids) + 1 + late_ratios = set( + config.compress_ratios[self.late_layer_start : config.num_hidden_layers] + ) + assert late_ratios <= { + 0, + 1, + }, f"late layers must not compress on their own, got ratios {late_ratios}" + def get_input_embeddings(self) -> nn.Module: return self.embed_tokens @@ -3140,6 +4149,168 @@ class DeepseekV4Model(nn.Module): hc_eps=self.hc_eps, ) + def _check_late_layer_tail_readers(self, forward_batch: ForwardBatch) -> None: + # Rows outside the tail are never computed past the last kv_source layer. + if ( + forward_batch.capture_hidden_mode == CaptureHiddenMode.FULL + and self.dspark_layers_to_capture is None + ): + raise ValueError( + "decoder SWA bounded replay cannot capture hidden states of all " + "prompt tokens" + ) + if forward_batch.return_logprob and any( + start < n + for start, n in zip( + forward_batch.extend_logprob_start_lens_cpu, + forward_batch.extend_seq_lens_cpu, + ) + ): + raise ValueError( + "decoder SWA bounded replay cannot return logprobs of prompt tokens; " + "set logprob_start_len to the prompt length" + ) + + def _forward_layers_hc_pre_from_prev( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + input_ids: torch.Tensor, + input_ids_global: torch.Tensor, + capture_dspark: bool, + dspark_aux_hidden_states: List[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, Optional[LateLayerTail]]: + assert self.pp_group.world_size == 1, "pre-mix hand-off across PP is not wired" + hash_ids = None + cp_extend = ( + is_cp_active(forward_batch) and forward_batch.forward_mode.is_extend() + ) + if self.engram_hasher is not None: + if cp_extend: + # n-gram hashing needs each token's predecessors: hash the whole prompt + total = int(forward_batch.attn_cp_metadata.total_seq_lens) + hash_ids = self.engram_hasher( + forward_batch.input_ids[:total], forward_batch + ) + parallel = get_parallel() + hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size] + pad_rows = hidden_states.shape[0] - hash_ids.shape[0] + if pad_rows > 0: + hash_ids = torch.cat( + [hash_ids, hash_ids.new_zeros(pad_rows, *hash_ids.shape[1:])] + ) + elif ( + forward_batch.forward_mode.is_extend() and is_in_breakable_cuda_graph() + ): + hash_ids = bcg_deepseek_v4_engram_hash_ids( + self.engram_hasher, input_ids + ) + else: + hash_ids = self.engram_hasher(input_ids, forward_batch) + tail = None + if ( + self.late_layer_start is not None + and forward_batch.forward_mode.is_extend_without_speculative() + ): + self._check_late_layer_tail_readers(forward_batch) + attn_backend = get_attn_backend() + tail = attn_backend.tail_forward_metadata.late_layer_tail + saved_full = None + prev_pre = None + precomputed_attn = None + for i in range(self.start_layer, self.end_layer): + if tail is not None and i == self.late_layer_start: + # Decode reaches back at most SWA_WINDOW positions. + saved_full = attn_backend.enter_late_layer_tail(forward_batch) + hidden_states, prev_pre, input_ids, input_ids_global = ( + tail.rows(hidden_states), + tail.rows(prev_pre), + tail.rows(input_ids), + tail.rows(input_ids_global), + ) + positions = tail.positions + if hash_ids is not None: + hash_ids = tail.rows(hash_ids) + engram = self.layers[i].engram + if engram is not None: + precomputed_attn = None + before_engram = hidden_states + hidden_states = engram( + hidden_states, + hash_ids[:, engram.layer_hash_index], + forward_batch, + cp_all_tokens=cp_extend, + ) + if ( + self.config.model_type == "deepseek_v41" + and self.config.vision_n_layers > 0 + ): + hidden_states = torch.where( + (input_ids == self.config.image_token_id)[:, None, None], + before_engram, + hidden_states, + ) + if capture_dspark and i in self.dspark_layers_to_capture: + # The draft head reads the attention input of its target layers. + aux = hidden_states + if tail is not None and i < self.late_layer_start: + aux = tail.rows(aux) + dspark_aux_hidden_states.append(aux.mean(dim=1)) + ctx = ( + nullcontext() + if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) + else get_global_expert_distribution_recorder().with_current_layer(i) + ) + next_norm = None + next_input = [] + if ( + self.config.model_type == "deepseek_v41" + and i + 1 < self.end_layer + and tail is None + and hidden_states.is_cuda + and get_platform().is_blackwell + and 0 < hidden_states.shape[0] <= 8 + and ( + forward_batch.forward_mode.is_decode() + or forward_batch.forward_mode.is_target_verify() + ) + and not get_forward().sp_active + and self.layers[i + 1].engram is None + and self.layers[i + 1].self_attn.accepts_mxfp8_swizzled_input() + ): + from sglang.srt.batch_invariant_ops import ( + is_batch_invariant_mode_enabled, + ) + + norm = self.layers[i + 1].input_layernorm + if ( + not norm.cast_x_before_out_mul + and norm.variance_size_override is None + and norm.weight.dtype == torch.bfloat16 + and norm.weight.shape == (5120,) + and norm.weight.is_contiguous() + and not is_batch_invariant_mode_enabled() + ): + next_norm = norm + with ctx: + hidden_states, prev_pre = self.layers[i].forward_hc_pre_from_prev( + positions=positions, + hidden_states=hidden_states, + input_ids=input_ids, + forward_batch=forward_batch, + input_ids_global=input_ids_global, + prev_pre=prev_pre, + precomputed_attn=precomputed_attn, + next_norm=next_norm, + next_input=next_input, + ) + precomputed_attn = next_input[0] if next_input else None + if saved_full is not None: + attn_backend.exit_late_layer_tail(saved_full, forward_batch) + return hidden_states, prev_pre, tail + return hidden_states, prev_pre, None + def _can_run_tbo(self, forward_batch: ForwardBatch) -> bool: """DSV4 prefill-only two-batch-overlap gate. @@ -3330,8 +4501,20 @@ class DeepseekV4Model(nn.Module): forward_batch, positions, ) - - if run_tbo: + last_pre = None + tail = None + if self.hc_pre_from_prev_sublayer: + assert not run_tbo, "two-batch overlap is not wired for this hc scheme" + hidden_states, last_pre, tail = self._forward_layers_hc_pre_from_prev( + positions, + hidden_states, + forward_batch, + input_ids, + input_ids_global, + capture_dspark, + dspark_aux_hidden_states, + ) + elif run_tbo: # Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is # disabled here (each layer self-contained), so no trailing hc_post. hidden_states = self._forward_layers_tbo( @@ -3381,11 +4564,28 @@ class DeepseekV4Model(nn.Module): pre_hc_head = hidden_states.flatten(1) - hidden_states = self.hc_head( - hidden_states, self.hc_head_fn, self.hc_head_scale, self.hc_head_base - ) + if self.hc_pre_from_prev_sublayer: + from sglang.kernels.ops.layernorm.mhc import hc_combine + + hidden_states = hc_combine( + pre_hc_head.float(), last_pre, self.hc_mult, hidden_states.dtype + ) + else: + hidden_states = self.hc_head( + hidden_states, self.hc_head_fn, self.hc_head_scale, self.hc_head_base + ) hidden_states = self.norm(hidden_states) + if tail is not None and not capture_dspark: + # The logits processor indexes rows by the full extend layout. + num_tokens = input_ids.shape[0] + hidden_states = _scatter_tail_rows( + tail=tail, rows=hidden_states, num_tokens=num_tokens + ) + pre_hc_head = _scatter_tail_rows( + tail=tail, rows=pre_hc_head, num_tokens=num_tokens + ) + if capture_dspark: return (hidden_states, pre_hc_head), dspark_aux_hidden_states @@ -3393,6 +4593,8 @@ class DeepseekV4Model(nn.Module): class DeepseekV4ForCausalLM(nn.Module): + supports_cuda_vmm_feature_transport = True + def __init__( self, config: DeepSeekV4Config, @@ -3412,7 +4614,25 @@ class DeepseekV4ForCausalLM(nn.Module): self.config = config self.tp_size = get_parallel().tp_size self.quant_config = quant_config + self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config) self.determine_num_fused_shared_experts() + self.vision = None + if config.model_type == "deepseek_v41" and config.vision_n_layers > 0: + if ( + get_parallel().attn_cp_size != 1 + or get_pp_group().world_size != 1 + or not get_moe_a2a_backend().is_none() + ): + raise ValueError( + "V4.1 vision currently supports TP/EP/DP without CP, PP or MoE A2A" + ) + + args = SimpleNamespace(**vars(config), dim=config.hidden_size) + self.vision = ViT(args) + self.aligner = Aligner(args) + self.image_start = nn.Parameter(torch.empty(config.hidden_size)) + self.image_end = nn.Parameter(torch.empty(config.hidden_size)) + self.image_newline = nn.Parameter(torch.empty(config.hidden_size)) self.model = DeepseekV4Model( config, quant_config, prefix=add_prefix("model", prefix) ) @@ -3453,10 +4673,116 @@ class DeepseekV4ForCausalLM(nn.Module): # its barrier must only run on the first (startup) load. self._mhc_prewarmed_at_load = False + @torch.inference_mode() + def wants_prefill_autotune(self) -> bool: + return getattr(self.config, "model_type", None) == "deepseek_v41" + + def autotune_prefill_kernels(self, num_tokens: int, *, dtype: torch.dtype) -> int: + """Tune resident MXFP8 linears for every M bucket up to ``num_tokens``. + The quant method is called directly, so no TP collectives run and no + request/KV/draft state is touched; the runner owns the 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 + ): + # No swizzled MXFP8 scale buffer: these kept the block-FP8 fallback. + continue + backend = method.mxfp8_dense_backend + if backend is None or not backend.is_flashinfer_cutedsl(): + continue + if method.block_fp8_as_mxfp8: + # Small shapes and deterministic execution keep 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 + def pad_input_ids(self, input_ids, mm_inputs): + return MultiModalityDataPaddingPatternMultimodalTokens().pad_input_tokens( + input_ids, mm_inputs + ) + + def get_image_feature(self, items): + """Return complete spans for the shared MM cache and chunk scheduler.""" + + spans = [] + device, dtype = self.image_start.device, self.image_start.dtype + for item in items: + item.reconstruct(device.index, ipc_consumer_count=self.tp_size) + h, w = int(item.n_vit_h), int(item.n_vit_w) + pixels = torch.as_tensor(item.feature, device=device) + plan = item.model_specific_data.get(GPU_PLAN_KEY) + patches = ( + materialize_image_gpu(pixels, plan).to(dtype) + if plan is not None + else pixels.to(dtype) + ) + features = self.aligner(self.vision(patches, h, w), h, w) + r = self.config.vision_downsample_ratio + types = image_token_types((h + r - 1) // r, (w + r - 1) // r).to(device) + span = torch.empty( + (len(types), self.config.hidden_size), device=device, dtype=dtype + ) + span[types == 0] = self.image_start + span[types == 1] = features.to(dtype) + span[types == 2] = self.image_newline + span[types == 3] = self.image_end + spans.append(span) + return spans + + def _prepare_mm_embeddings(self, input_ids, forward_batch): + # Keep scheduler hash IDs intact: the shared embedder clamps its input in place. + input_embeds, _ = embed_mm_inputs( + mm_inputs_list=[ + item if item is not None else MultimodalInputs(mm_items=[]) + for item in forward_batch.mm_inputs + ], + extend_prefix_lens=forward_batch.extend_prefix_lens_cpu, + extend_seq_lens=forward_batch.extend_seq_lens_cpu, + input_ids=input_ids.clone(), + input_embedding=self.get_input_embeddings(), + multimodal_model=self, + ) + forward_batch.mm_input_embeds = input_embeds + return input_embeds + def get_input_embeddings(self) -> nn.Module: return self.model.get_input_embeddings() @@ -3513,6 +4839,26 @@ class DeepseekV4ForCausalLM(nn.Module): input_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: + if ( + self.vision is not None + and not forward_batch.forward_mode.is_decode() + and not forward_batch.forward_mode.is_target_verify() + and forward_batch.mm_inputs is not None + and any(x is not None for x in forward_batch.mm_inputs) + ): + if input_embeds is not None: + raise ValueError("Cannot combine input_embeds and image inputs") + input_embeds = self._prepare_mm_embeddings(input_ids, forward_batch) + if self.vision is not None and not ( + forward_batch.forward_mode.is_decode_or_idle() + or forward_batch.forward_mode.is_target_verify() + ): + # Decode/verify IDs are already vocabulary IDs; remap prompt image + # hashes for Engram and routing. + input_ids = input_ids.masked_fill( + input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id + ) + with get_attn_tp_context().maybe_input_scattered(forward_batch): hidden_states = self.model.forward( input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors @@ -3525,16 +4871,33 @@ class DeepseekV4ForCausalLM(nn.Module): hidden_states, aux_hidden_states = hidden_states hidden_states, pre_hc_head = hidden_states - return self.logits_processor( + logits_metadata = forward_batch + tail = None + if ( + self.capture_aux_hidden_states + and self.model.late_layer_start is not None + and forward_batch.forward_mode.is_extend_without_speculative() + ): + tail = get_attn_backend().tail_forward_metadata.late_layer_tail + input_ids = tail.rows(input_ids) + logits_metadata = LogitsMetadata.from_forward_batch(forward_batch) + logits_metadata.extend_seq_lens = tail.extend_seq_lens + logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu + logits_metadata.extend_logprob_start_lens_cpu = tail.extend_seq_lens_cpu + + output = self.logits_processor( input_ids, hidden_states, self.lm_head, - forward_batch, + logits_metadata, aux_hidden_states, hidden_states_before_norm=( None if aux_hidden_states is not None else pre_hc_head ), ) + if tail is not None: + output.hidden_states_token_indices = tail.token_indices + return output def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None: from sglang.srt.layers import deep_gemm_wrapper @@ -3593,7 +4956,7 @@ class DeepseekV4ForCausalLM(nn.Module): attn.wo_a.weight_scale_inv.format_ue8m0 = False def post_load_weights(self, is_nextn=False, weight_names=None): - if _FP8_WO_A_GEMM: + if self.wo_a_fp8: self._setup_fp8_wo_a_scales(is_nextn) if is_nextn: @@ -3619,6 +4982,12 @@ class DeepseekV4ForCausalLM(nn.Module): is_nextn: bool = False, num_hidden_layers: Optional[int] = None, ) -> str: + if name.startswith("vision."): + return name.replace(".attn.wqkv.", ".attn.qkv_proj.").replace( + ".attn.wo.", ".attn.proj." + ) + if name.startswith(("aligner.", "image_")): + return name if name.startswith("embed."): return "model.embed_tokens." + name.removeprefix("embed.") if name.startswith("head."): @@ -3665,6 +5034,8 @@ class DeepseekV4ForCausalLM(nn.Module): if "self_attn" in name and name.endswith(".scale"): name = name.removesuffix(".scale") + ".weight_scale_inv" + if ".engram.wkv." in name and name.endswith(".scale"): + name = name.removesuffix(".scale") + ".weight_scale_inv" name = name.replace(".gate.tid2eid", ".topk.tid2eid") name = name.replace(".gate.bias", ".gate.e_score_correction_bias") @@ -3763,7 +5134,7 @@ class DeepseekV4ForCausalLM(nn.Module): # Must mirror MQALayer.__init__'s `quantize_wo_a`: dequantizing wo_a here # while the layer allocated an FP8 parameter (or vice versa) fails the # weight loader's dtype check. - if not (_FP8_WO_A_GEMM or use_npu_arch35_mxfp8_wo_a(self.quant_config)): + if not (self.wo_a_fp8 or use_npu_arch35_mxfp8_wo_a(self.quant_config)): weights = _prepare_deepseek_v4_weights(weights, self.quant_config) stacked_params_mapping = DEEPSEEK_V4_STACKED_PARAMS_MAPPING @@ -3785,6 +5156,9 @@ class DeepseekV4ForCausalLM(nn.Module): fuse_wqa_wkv = envs.SGLANG_OPT_FUSE_WQA_WKV.get() cache_wqkv_a_weight: dict[str, dict[str, torch.Tensor]] = {} + skipped_by_group: dict[str, int] = {} + # V4 checkpoints must load every compressor / indexer tensor. + is_dsv41 = getattr(self.config, "model_type", None) == "deepseek_v41" def auto_weight_loader(module): return getattr(module, "weight_loader", default_weight_loader) @@ -3813,7 +5187,7 @@ class DeepseekV4ForCausalLM(nn.Module): weight_names = [] for name, loaded_weight in weights: if ( - _FP8_WO_A_GEMM + self.wo_a_fp8 and name.endswith(".wo_a.weight") and loaded_weight.dtype != torch.float8_e4m3fn ): @@ -3833,6 +5207,24 @@ class DeepseekV4ForCausalLM(nn.Module): num_hidden_layers=self.config.num_hidden_layers, ) + # V4.1 checkpoint tensors with no module in the text model yet. + skip_group = None + if not is_dsv41: + pass + elif self.vision is None and name.startswith( + ("vision.", "aligner.", "image_") + ): + skip_group = "vision" + elif self.vision is None and name.endswith( + ".gate.e_score_correction_bias_vl" + ): + skip_group = "gate.bias_vl" + if skip_group is not None: + skipped_by_group[skip_group] = ( + skipped_by_group.get(skip_group, 0) + 1 + ) + continue + layer_id = get_layer_id(name) if ( layer_id is not None @@ -3961,7 +5353,13 @@ class DeepseekV4ForCausalLM(nn.Module): or name == "lm_head.weight" ) and not self.pp_group.is_last_rank: continue - elif COMPRESSOR_PART in name and ".wkv_gate." not in name: + elif ( + COMPRESSOR_PART in name + and ".wkv_gate." not in name + and (name.rsplit(".", 2)[0] + ".wkv_gate.weight") + in params_dict + ): + # Split-projection modules load per parameter instead. is_kv = name.endswith(".wkv.weight") is_wgate = name.endswith(".wgate.weight") assert is_kv != is_wgate @@ -3993,15 +5391,20 @@ class DeepseekV4ForCausalLM(nn.Module): ) loaded_params.add(param_name) cache_compressor_weight.pop(key) - elif fuse_wqa_wkv and ( - name.endswith(".wq_a.weight") - or name.endswith(".wq_a.weight_scale_inv") - or name.endswith(".wkv.weight") - or name.endswith(".wkv.weight_scale_inv") - or name.endswith(".wq_a.qweight") - or name.endswith(".wkv.qweight") - or name.endswith(".wq_a.qweight_type") - or name.endswith(".wkv.qweight_type") + elif ( + fuse_wqa_wkv + and ".compressor." not in name + and ".engram." not in name + and ( + name.endswith(".wq_a.weight") + or name.endswith(".wq_a.weight_scale_inv") + or name.endswith(".wkv.weight") + or name.endswith(".wkv.weight_scale_inv") + or name.endswith(".wq_a.qweight") + or name.endswith(".wkv.qweight") + or name.endswith(".wq_a.qweight_type") + or name.endswith(".wkv.qweight_type") + ) ): is_q = ".wq_a." in name param_name = name.replace( @@ -4066,6 +5469,12 @@ class DeepseekV4ForCausalLM(nn.Module): assert len(cache_compressor_weight) == 0 assert len(cache_wqkv_a_weight) == 0, cache_wqkv_a_weight.keys() + if skipped_by_group: + log_info_on_rank0( + logger, + "Skipped checkpoint tensors not wired yet: " + + ", ".join(f"{k}={v}" for k, v in sorted(skipped_by_group.items())), + ) unloaded_params = params_dict.keys() - loaded_params skipped_checking_patterns = [ @@ -4096,6 +5505,9 @@ class DeepseekV4ForCausalLM(nn.Module): self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) if not is_nextn: + for i, layer in enumerate(self.model.layers): + if getattr(layer, "engram", None) is not None: + layer.engram.embed.finish_load(label=f"layer {i}") self._prewarm_mhc_kernels() def get_embed_and_head(self): @@ -4134,8 +5546,11 @@ def _dequant_fp8(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: torch.float32, ), f"expected fp8_e8m0fnu or float32, got {scale.dtype}" + # Block size is per-checkpoint: V4 128x128, V4.1 32x32. + bn = weight.shape[0] // scale.shape[0] + bk = weight.shape[1] // scale.shape[1] weight_f32 = rearrange( - weight.float(), "(sn bn) (sk bk) -> sn bn sk bk", bn=128, bk=128 + weight.float(), "(sn bn) (sk bk) -> sn bn sk bk", bn=bn, bk=bk ) result = rearrange( weight_f32 * scale.float()[:, None, :, None], "sn bn sk bk -> (sn bn) (sk bk)" diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 4994a9254..baebc2de6 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy import logging from typing import Iterable, List, Optional, Tuple @@ -17,6 +18,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_draft_model import ( CommitKvProj, ) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config +from sglang.srt.distributed.device_communicators.vocab_gather import make_vocab_gather from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import is_dp_attention_enabled from sglang.srt.layers.layernorm import RMSNorm @@ -39,6 +41,7 @@ from sglang.srt.models.deepseek_v4 import ( DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, MqaAttentionBase, + _apply_wo_a_bf16_matmul, _dequant_fp8_wo_a_streaming, hc_head_torch, make_hc_head_params, @@ -189,6 +192,22 @@ class DSparkAttention(MqaAttentionBase): q = self.q_norm(q) q, _ = self.wq_b(q) q = q.view(-1, self.n_local_heads, self.head_dim) + if not self.q_head_norm: + if self._use_fast_kernel and not _is_npu: + fused_rope_inplace( + q[..., -self.rope_head_dim :], + None, + self.freqs_cis, + positions=positions, + ) + else: + apply_rotary_emb( + q[..., -self.rope_head_dim :], self.freqs_cis[positions] + ) + if q_out is None: + return q + q_out.copy_(q) + return q_out if self._use_fast_kernel: if q_out is None: q_out = torch.empty_like(q) @@ -236,7 +255,6 @@ class DSparkAttention(MqaAttentionBase): hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - if _is_npu and forward_batch.forward_mode.is_idle(): return torch.zeros_like(hidden_states) @@ -336,7 +354,13 @@ class DSparkAttention(MqaAttentionBase): ) wo_a = self.wo_a.weight.view(self.n_local_groups, self.o_lora_rank, -1) if self._use_fast_kernel: - o = torch.einsum("bgd,grd->bgr", o, wo_a) + o = _apply_wo_a_bf16_matmul( + o, + wo_a, + is_decode=forward_batch.forward_mode.is_decode(), + is_target_verify=forward_batch.forward_mode.is_target_verify(), + fast_path=self.is_dsv41, + ) else: o = torch.einsum("bgd,grd->bgr", o.float(), wo_a.float()).to(q.dtype) out, _ = self.wo_b(o.reshape(o.shape[0], o.shape[1] * o.shape[2])) @@ -363,10 +387,13 @@ class MarkovW2ShardGeometry(msgspec.Struct, frozen=True): class DSparkV4MarkovHead(nn.Module): markov_head_type = "vanilla" - def __init__(self, *, vocab_size: int, markov_rank: int) -> None: + def __init__( + self, *, vocab_size: int, markov_rank: int, is_dsv41: bool = False + ) -> None: super().__init__() self.vocab_size = int(vocab_size) self.markov_rank = int(markov_rank) + self._is_dsv41 = bool(is_dsv41) if self.markov_rank <= 0: raise ValueError( f"DSparkV4MarkovHead requires markov_rank > 0, got {self.markov_rank}." @@ -418,6 +445,15 @@ class DSparkV4MarkovHead(nn.Module): "Disable SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD." ) self._shard_group = shard_group + self._vocab_gather = make_vocab_gather( + shard_group, + local_width=per_partition, + prefer_nvlink=self._is_dsv41 + and envs.SGLANG_DSPARK_NVLINK_VOCAB_GATHER.get(), + ) + if shard_group.rank == 0: + cls_name = type(self._vocab_gather).__name__ + logger.info("DSpark markov_w2 vocab gather: %s", cls_name) self._tp_shard = MarkovW2ShardGeometry( tp_size=tp_size, org_vocab_start=int(lm_head.shard_indices.org_vocab_start_index), @@ -469,13 +505,40 @@ class DSparkV4MarkovHead(nn.Module): else: bias = F.linear(latent.float(), weight_local) step_local = BuildStepLocal.execute(bias=bias, base_local=base_local) - if shard.tp_size > 1: - assert self._shard_group is not None - full = self._shard_group.all_gather(step_local, dim=-1) - else: - full = step_local + full = self._vocab_gather(step_local) return full[..., : self.vocab_size] + @property + def supports_sharded_greedy(self) -> bool: + return ( + self._is_dsv41 and self._tp_shard is not None and self._opt_markov_w2_bf16 + ) + + def sample_block_greedy_fused(self, base_logits, *, first_prev_tokens): + if not self.supports_sharded_greedy or not base_logits.is_cuda: + return None + from sglang.kernels.ops.speculative.dspark.sharded_greedy import ( + sharded_greedy_step, + ) + + shard = self._tp_shard + weight = self.markov_w2.weight[shard.org_vocab_start : shard.org_vocab_end] + prev = first_prev_tokens.long() + tokens = [] + for step in range(base_logits.shape[1]): + latent = self.get_prev_embeddings(prev) + # Preserve the same BF16 GEMM rounding before the FP32 logits add. + bias = F.linear(latent.to(weight.dtype), weight) + prev = sharded_greedy_step( + bias, + base_logits[:, step], + group=self._shard_group, + vocab_start=shard.org_vocab_start, + gather=self._vocab_gather.gather_stacked, + ) + tokens.append(prev) + return torch.stack(tokens, dim=1) + def forward(self, token_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: embed = self.get_prev_embeddings(token_ids) logits = self.project_bias(embed) @@ -527,6 +590,22 @@ def build_dspark_v4_confidence_head( ) +def _dspark_stage_config(config: DeepSeekV4Config) -> DeepSeekV4Config: + n_routed = int(getattr(config, "dspark_n_routed_experts", 0) or 0) + n_active = int(getattr(config, "dspark_num_experts_per_tok", 0) or 0) + has_vision = int(getattr(config, "vision_n_layers", 0) or 0) > 0 + if not (n_routed or n_active or has_vision): + return config + stage_config = copy.copy(config) + if n_routed: + stage_config.n_routed_experts = n_routed + if n_active: + stage_config.num_experts_per_tok = n_active + if has_vision: + stage_config.vision_n_layers = 0 + return stage_config + + class DSparkV4Stage(DeepseekV4DecoderLayer): def __init__( self, @@ -538,14 +617,18 @@ class DSparkV4Stage(DeepseekV4DecoderLayer): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_streams: Optional[List[torch.cuda.Stream]] = None, + hc_stats_stream: Optional[torch.cuda.Stream] = None, + moe_routed_quant_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__( - config=config, + config=_dspark_stage_config(config), layer_id=layer_id, quant_config=quant_config, prefix=prefix, is_nextn=True, alt_streams=alt_streams, + hc_stats_stream=hc_stats_stream, + moe_routed_quant_stream=moe_routed_quant_stream, ) self.stage_id = stage_id self.dim = config.hidden_size @@ -566,11 +649,16 @@ class DSparkV4Stage(DeepseekV4DecoderLayer): if stage_id == num_stages - 1: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - ( - self.hc_head_fn, - self.hc_head_base, - self.hc_head_scale, - ) = make_hc_head_params(config.hc_mult, config.hidden_size) + if self.hc_pre_from_prev_sublayer: + # V4.1 collapses the head with the last FFN's pre-mix; the + # checkpoint carries no hc_head_* tensors for the stages. + self.hc_head_fn = self.hc_head_base = self.hc_head_scale = None + else: + ( + self.hc_head_fn, + self.hc_head_base, + self.hc_head_scale, + ) = make_hc_head_params(config.hc_mult, config.hidden_size) def _build_self_attn( self, @@ -615,7 +703,12 @@ class DSparkV4Stage(DeepseekV4DecoderLayer): positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, - ) -> torch.Tensor: + prev_pre: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if self.hc_pre_from_prev_sublayer: + return self._forward_hc_pre_from_prev( + positions, hidden_states, forward_batch, prev_pre + ) residual = hidden_states x, post, comb = self._hc_pre_block( hidden_states, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base @@ -632,7 +725,49 @@ class DSparkV4Stage(DeepseekV4DecoderLayer): x = self.post_attention_layernorm(x) x = self._run_ffn(x, forward_batch) x = self._hc_post_block(x, residual, post, comb) - return x + return x, None + + def _forward_hc_pre_from_prev( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + prev_pre: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch) + residual = hidden_states + x = self._hc_combine( + hidden_states, prev_pre, self.input_layernorm, stats_stream + ) + with self.self_attn.maybe_use_decode_attn_tp(forward_batch): + x = self.self_attn(positions, x, forward_batch) + attn_pre, attn_post, attn_comb = self._hc_mix_stats( + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + stats_stream, + ) + if stats_stream is not None: + torch.cuda.current_stream().wait_stream(stats_stream) + hidden_states = self.hc_post(x, residual, attn_post, attn_comb) + + residual = hidden_states + x = self._hc_combine( + hidden_states, attn_pre, self.post_attention_layernorm, stats_stream + ) + x = self._run_ffn(x, forward_batch) + ffn_pre, ffn_post, ffn_comb = self._hc_mix_stats( + hidden_states, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + stats_stream, + ) + if stats_stream is not None: + torch.cuda.current_stream().wait_stream(stats_stream) + hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb) + return hidden_states, ffn_pre def _run_ffn(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor: shape = x.shape @@ -714,6 +849,21 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): self.alt_streams: Optional[List[torch.cuda.Stream]] = ( [torch.cuda.Stream()] if use_multi_stream else None ) + self.moe_routed_quant_stream = ( + torch.cuda.Stream() + if use_multi_stream + and torch.version.cuda is not None + and getattr(config, "hc_pre_from_prev_sublayer", False) + else None + ) + self.hc_stats_stream = ( + torch.cuda.Stream() + if use_multi_stream + and torch.version.cuda is not None + and get_platform().is_blackwell + and getattr(config, "hc_pre_from_prev_sublayer", False) + else None + ) self.stages = nn.ModuleList( [ DSparkV4Stage( @@ -725,6 +875,8 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): quant_config=quant_config, prefix=add_prefix(f"stages.{stage_id}", prefix), alt_streams=self.alt_streams, + hc_stats_stream=self.hc_stats_stream, + moe_routed_quant_stream=self.moe_routed_quant_stream, ) for stage_id in range(self.num_stages) ] @@ -732,6 +884,7 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): self.markov_head = DSparkV4MarkovHead( vocab_size=int(config.vocab_size), markov_rank=int(dspark_config.markov_rank), + is_dsv41=getattr(config, "model_type", None) == "deepseek_v41", ) self.confidence_head = build_dspark_v4_confidence_head( config=config, markov_rank=int(dspark_config.markov_rank) @@ -739,6 +892,9 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): self.hc_mult = int(config.hc_mult) self.norm_eps = float(config.rms_norm_eps) self.hc_eps = float(config.hc_eps) + self.hc_pre_from_prev_sublayer = bool( + getattr(config, "hc_pre_from_prev_sublayer", False) + ) if self.uses_own_vocab_modules: self.embed_tokens = VocabParallelEmbedding( @@ -791,6 +947,12 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): kvs = CommitKvProj.execute( main_x=main_x, wkv_linears=[stage.self_attn.wkv for stage in self.stages], + # The FlashMLA writer reads an explicit KV row stride, so views are fine. + allow_strided_output=( + get_platform().is_blackwell + and not is_unified_kv_triton() + and not pool.uniform_fp8 + ), ) # Under unified_kv the swa_kv_pool is None; the caller passes a unified # ring loc (state_slot * ring + pos % ring, -1 for uncommitted) so the @@ -836,12 +998,20 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): if input_embeds is None: input_embeds = self.forward_embed(input_ids) x = input_embeds + pre = None for stage in self.stages: - x = stage(positions, x, forward_batch) + x, pre = stage(positions, x, forward_batch, pre) + if self.hc_pre_from_prev_sublayer: + from sglang.kernels.ops.layernorm.mhc import hc_combine + + x = hc_combine(x.flatten(1).float(), pre, self.hc_mult, x.dtype) return LogitsProcessorOutput(next_token_logits=None, hidden_states=x) def collapse_hc_head(self, x: torch.Tensor) -> torch.Tensor: + if self.hc_pre_from_prev_sublayer: + assert x.dim() == 2, "V4.1 draft hidden states leave forward() collapsed" + return x last = self.stages[-1] return hc_head_torch( x, @@ -853,7 +1023,6 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): ) def compute_base_logits(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - x_post_hc = self.collapse_hc_head(x) return self._logits_from_x_post_hc(x_post_hc), x_post_hc @@ -1008,6 +1177,8 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): stage_id, rest = parts[1], parts[2] if rest.startswith("markov_head."): + rest = rest.replace("markov_head.embed.", "markov_head.markov_w1.", 1) + rest = rest.replace("markov_head.head.", "markov_head.markov_w2.", 1) return f"markov_head.{rest[len('markov_head.') :]}" if rest.startswith("confidence_head."): @@ -1024,6 +1195,8 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): mapped_rest = mapped_rest.replace(".w2.", ".down_proj.") mapped_rest = mapped_rest.replace(".w3.", ".up_proj.") mapped_rest = mapped_rest.replace(".gate.tid2eid", ".topk.tid2eid") + if mapped_rest.endswith(".gate.bias_vl"): + return None mapped_rest = mapped_rest.replace(".gate.bias", ".gate.e_score_correction_bias") mapped_rest = mapped_rest.replace(".scale", ".weight_scale_inv") return f"stages.{stage_id}.{mapped_rest}" diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index 469f74be1..f85401e1e 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -35,6 +35,7 @@ from sglang.srt.models.deepseek_v4 import ( DeepseekV4DecoderLayer, DeepseekV4ForCausalLM, _is_npu, + wo_a_fp8_gemm_enabled, ) from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix @@ -220,6 +221,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): self.tp_size = get_parallel().tp_size self.pp_group = get_pp_group() self.quant_config = quant_config + self.wo_a_fp8 = wo_a_fp8_gemm_enabled(quant_config) self.determine_num_fused_shared_experts() self.model = DeepseekV4ModelNextN( diff --git a/python/sglang/srt/multimodal/dsv41/__init__.py b/python/sglang/srt/multimodal/dsv41/__init__.py new file mode 100644 index 000000000..68cd951a3 --- /dev/null +++ b/python/sglang/srt/multimodal/dsv41/__init__.py @@ -0,0 +1 @@ +"""Vision tower, image preprocessing and VL expert routing for DeepSeek-V4.1.""" diff --git a/python/sglang/srt/multimodal/dsv41/vl_routing.py b/python/sglang/srt/multimodal/dsv41/vl_routing.py new file mode 100644 index 000000000..eaedc57f2 --- /dev/null +++ b/python/sglang/srt/multimodal/dsv41/vl_routing.py @@ -0,0 +1,106 @@ +import torch +import torch.nn.functional as F + +from sglang.srt.layers.moe.topk import ( + _RENORMALIZE_SUM_EPSILON, + StandardTopKOutput, + StandardTopKOutputPacked, + _mask_topk_ids_padded_region, + _zero_topk_weights_padded_region, +) +from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots +from sglang.srt.utils import is_cuda + + +def _scale_fused_shared_weights(weights, num_fused_shared_experts, scaling_factor): + # Standard EP replicates the fused shared expert on every rank and all-reduces, + # so the shared columns carry a 1/ep_size factor. + if num_fused_shared_experts and scaling_factor is not None: + weights[:, -num_fused_shared_experts:] *= scaling_factor + return weights + + +def vision_topk(moe, logits, input_ids, num_token_non_padded=None): + config = moe.topk.topk_config + num_fused_shared_experts = config.num_fused_shared_experts + if num_fused_shared_experts: + # This path bypasses _post_process_topk_ids, which appends the per-rank slots. + assert not has_per_rank_fused_shared_slots(num_fused_shared_experts), ( + "VL routing does not support per-rank fused shared slots" + ) + if is_cuda(): + from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate + from sglang.srt.layers.moe.utils import get_moe_runner_backend + + # Same admission as _fused_gate_emits_packed_ids: the shared-expert slots + # rescaled below would rewrite weights after the router. + packed_topk = None + if ( + num_fused_shared_experts == 0 + and get_moe_runner_backend().is_flashinfer_mxfp4() + ): + packed_topk = torch.empty( + (logits.shape[0], config.top_k), dtype=torch.int32, device=logits.device + ) + + weights, indices = moe_fused_gate( + logits, + moe.gate.e_score_correction_bias, + topk=config.top_k, + scoring_func="sqrtsoftplus", + num_fused_shared_experts=num_fused_shared_experts, + bias_alt=moe.gate.e_score_correction_bias_vl, + input_ids=input_ids, + bias_alt_token_id=moe.config.image_token_id, + renormalize=config.renormalize and config.top_k > 1, + renormalize_epsilon=_RENORMALIZE_SUM_EPSILON, + routed_scaling_factor=config.routed_scaling_factor, + apply_routed_scaling_factor_on_output=config.apply_routed_scaling_factor_on_output, + num_token_non_padded=num_token_non_padded, + packed_out=packed_topk, + sqrtsoftplus_log1p=True, + ) + weights = _scale_fused_shared_weights( + weights, + num_fused_shared_experts, + config.fused_shared_experts_scaling_factor, + ) + if packed_topk is not None: + return StandardTopKOutputPacked(weights, indices, logits, packed_topk) + return StandardTopKOutput(weights, indices, logits) + scores = F.softplus(logits.float()).sqrt() + if input_ids is None: + bias = moe.gate.e_score_correction_bias + else: + bias = torch.where( + (input_ids == moe.config.image_token_id)[:, None], + moe.gate.e_score_correction_bias_vl, + moe.gate.e_score_correction_bias, + ) + # The shared slots appended below use the same layout as biased_grouped_topk_gpu. + topk_routed = config.top_k - num_fused_shared_experts + indices = (scores + bias).topk(topk_routed, dim=-1).indices + weights = scores.gather(-1, indices) + routed_sum = weights.sum(-1, keepdim=True, dtype=torch.float32) + if num_fused_shared_experts: + shared_ids = logits.shape[-1] + torch.arange( + num_fused_shared_experts, device=indices.device, dtype=indices.dtype + ) + indices = torch.cat( + [indices, shared_ids.expand(indices.shape[0], -1)], + dim=-1, + ) + weights = F.pad(weights, (0, num_fused_shared_experts)) + weights[:, topk_routed:] = routed_sum / config.routed_scaling_factor + if config.renormalize and config.top_k > 1: + weights = weights / (routed_sum + _RENORMALIZE_SUM_EPSILON) + if config.apply_routed_scaling_factor_on_output: + weights = weights * config.routed_scaling_factor + weights = _scale_fused_shared_weights( + weights, num_fused_shared_experts, config.fused_shared_experts_scaling_factor + ) + weights, indices = weights.float(), indices.int() + if num_token_non_padded is not None: + _mask_topk_ids_padded_region(indices, num_token_non_padded) + _zero_topk_weights_padded_region(weights, num_token_non_padded) + return StandardTopKOutput(weights, indices, logits) diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 8390480af..9427ed944 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -49,6 +49,8 @@ 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. + 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_draft_sampler.py b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py index f086b91ba..f24e6d6e5 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py @@ -57,9 +57,6 @@ class DsparkDraftSampler: self.sample_from_anchor = bool(model.sample_from_anchor) self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1 max_bs = int(max_bs) - # Resolved once: this sampler runs inside cuda-graph capture, so the - # branch below is baked into the captured graph anyway. - self._fused_greedy = envs.SGLANG_DSPARK_OPT_FUSED_GREEDY_MARKOV.get() if out is not None: assert out.shape == (max_bs * self.gamma,) and out.dtype == torch.int64 self.out = out @@ -129,11 +126,11 @@ class DsparkDraftSampler: # Gated/RNN subclasses return None (hidden-state-dependent bias); fall # through to the block sampler below. draft_tokens = None - if ( - not self.folded_sampling - and self._fused_greedy + fused_greedy = getattr(self.markov_head, "supports_sharded_greedy", False) or ( + envs.SGLANG_DSPARK_OPT_FUSED_GREEDY_MARKOV.get() and isinstance(self.markov_head, VanillaMarkov) - ): + ) + if not self.folded_sampling and fused_greedy: draft_tokens = self.markov_head.sample_block_greedy_fused( base_logits, first_prev_tokens=anchor ) @@ -198,6 +195,9 @@ def _resolve_folded_sampling( return False if mode == DsparkFoldedSampling.FORCE: return True + # The V4.1 TP head reduces compact argmax summaries in the greedy graph. + if getattr(model.markov_head, "supports_sharded_greedy", False): + return False vocab = int(model.lm_head.org_vocab_size) noise_bytes = max_bs * vocab * 4 logits_bytes = max_bs * gamma * vocab * _base_logits_dtype(model).itemsize diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index c283ecff3..0499fc478 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -121,6 +121,24 @@ class TargetHiddenKvInjector: state_slot=state_slot, final_pos=final_pos, ) + elif ( + cache_loc.is_cuda + and cache_loc.is_contiguous() + and commit_lens is not None + and cache_loc_2d is not None + and commit_lens.is_contiguous() + and cache_loc.numel() == cache_loc_2d.numel() + ): + from sglang.kernels.ops.speculative.dspark.commit_swa import ( + committed_swa_locations, + ) + + swa_loc = committed_swa_locations( + cache_loc, + pool.full_to_swa_index_mapping, + commit_lens, + cache_loc_2d.shape[1], + ) else: swa_loc = pool.translate_loc_from_full_to_swa(cache_loc).to(torch.int32) if commit_lens is not None and cache_loc_2d is not None: diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index f118c2d60..dd45c0d90 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -78,6 +78,30 @@ 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 hold one unprocessed result, so reserve its full width; + the runner adds the current verify width. Aborted/embedding/multimodal requests + return None: their visible token IDs may not track cache positions.""" + 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, @@ -92,6 +116,15 @@ class TargetVerifyExecutor: simulate_acc_len: float = 0.0, ) -> None: self.target_worker = target_worker + # candidate_max_seq_len_upper_bound only feeds the V4.1 candidate graphs. + self._target_is_dsv41 = ( + getattr( + target_worker.model_runner.model_config.hf_text_config, + "model_type", + None, + ) + == "deepseek_v41" + ) self.gamma = int(gamma) self.verify_num_draft_tokens = verify_num_draft_tokens self.model_runner = model_runner @@ -135,6 +168,7 @@ class TargetVerifyExecutor: gamma=self.gamma, verify_num_draft_tokens=self.verify_num_draft_tokens, cutoff_layout=layout, + fused_argmax=self._target_is_dsv41, ) if self._simulate_acc_len > 0: correct_len = self._simulated_correct_len( @@ -296,6 +330,10 @@ class TargetVerifyExecutor: seq_lens_cpu_backup, seq_lens_sum_backup, ) -> TargetVerifyResult: + if verify_input.live_seq_lens_cpu is None and self._target_is_dsv41: + 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 ) @@ -479,6 +517,7 @@ class CommitInjectCtx(msgspec.Struct): block_pos_offsets: torch.Tensor resolve_pool: object resolve_req_to_token: object + kv_injector: Optional[TargetHiddenKvInjector] = None class AcceptOuts(msgspec.Struct): @@ -499,9 +538,11 @@ class DsparkVerifyEpilogue: device, tp_sync: SpecTpSync, commit_ctx: Optional[CommitInjectCtx] = None, + fused_argmax: bool = False, ) -> None: self.max_bs = int(max_bs) self.stride = int(verify_num_draft_tokens) + self._fused_argmax = bool(fused_argmax) self.gamma = self.stride - 1 self.commit_ctx = commit_ctx self._tp_sync = tp_sync @@ -530,9 +571,13 @@ class DsparkVerifyEpilogue: ) self.strided_logits: Optional[torch.Tensor] = None self.strided_hidden: Optional[torch.Tensor] = None + self._static_step_state: Optional[tuple[int, bool]] = None def capture_hook(self, runner, out, forward_batch, num_tokens) -> None: - if runner.model_runner.is_draft_worker or not runner.ragged_verify_mode: + if ( + runner.model_runner.is_draft_worker + or not forward_batch.forward_mode.is_target_verify() + ): return if ( not isinstance(out, LogitsProcessorOutput) @@ -540,6 +585,9 @@ class DsparkVerifyEpilogue: or out.hidden_states is None ): return + if not runner.ragged_verify_mode: + self._static_epilogue(out, forward_batch) + return self( compact_logits=out.next_token_logits, compact_hidden=out.hidden_states, @@ -550,6 +598,7 @@ class DsparkVerifyEpilogue: ) def begin_step(self, verify_lens, armed: bool) -> None: + self._static_step_state = None if verify_lens is None: self.verify_lens_buf.zero_() else: @@ -559,6 +608,49 @@ class DsparkVerifyEpilogue: self.verify_lens_buf[bs:].zero_() self.inject_gate_buf.fill_(1 if armed else 0) + def begin_static_step(self, bs: int, armed: bool) -> None: + state = (bs, armed) + if self._static_step_state == state: + return + self.verify_lens_buf[:bs].fill_(self.stride) + self.verify_lens_buf[bs:].zero_() + self.inject_gate_buf.fill_(int(armed)) + self._static_step_state = state + + def _static_epilogue(self, out, forward_batch) -> None: + bs = forward_batch.batch_size + verify_lens = self.verify_lens_buf[:bs] + candidates = forward_batch.input_ids.view(bs, self.stride) + commit_lens = self._accept( + candidates=candidates, + logits=out.next_token_logits, + draft_tokens=candidates[:, 1:].contiguous(), + seq_lens=forward_batch.seq_lens, + ) + if not self.folds_commit: + return + # Same staged locations as target verify; padded and fallback rows skip KV. + gated_commit_lens = ( + torch.minimum(commit_lens, verify_lens.to(torch.int32)) + * self.inject_gate_buf + ) + cache_loc = forward_batch.out_cache_loc + state_slot = None + if is_unified_kv_triton(): + state_slot = ( + forward_batch.req_pool_indices.view(-1, 1) + .expand(bs, self.stride) + .reshape(-1) + ) + self.commit_ctx.kv_injector.inject_target_hidden( + target_hidden=out.hidden_states, + cache_loc=cache_loc, + cache_loc_2d=cache_loc.view(bs, self.stride), + positions=forward_batch.positions, + commit_lens=gated_commit_lens, + state_slot=state_slot, + ) + def read_accept(self, bs: int) -> AcceptOuts: return AcceptOuts( correct_len=self.correct_len_buf[:bs], @@ -610,7 +702,23 @@ class DsparkVerifyEpilogue: self.strided_hidden = self._ensure_out(self.strided_hidden, compact_hidden) verify_lens = self.verify_lens_buf[:bs] self._scatter(compact_logits, compact_hidden, verify_lens, bs) - commit_lens = self._accept(input_ids, seq_lens, verify_lens, bs) + candidates = torch.zeros( + (bs * self.stride, 1), dtype=input_ids.dtype, device=input_ids.device + ) + scatter_compact_to_strided_into( + compact=input_ids.view(-1, 1), + verify_lens=verify_lens, + out=candidates, + stride=self.stride, + fill_value=0, + ) + commit_lens = self._accept( + candidates=candidates.view(bs, self.stride), + logits=self.strided_logits[: bs * self.stride], + draft_tokens=self.draft_tokens_buf[: bs * self.gamma].view(bs, self.gamma), + seq_lens=seq_lens, + cutoff_verify_lens=verify_lens, + ) if self.folds_commit: self._commit_inject( commit_lens, verify_lens, seq_lens, req_pool_indices, bs @@ -632,22 +740,16 @@ class DsparkVerifyEpilogue: fill_value=0.0, ) - def _accept(self, input_ids, seq_lens, verify_lens, bs: int) -> torch.Tensor: - candidates = torch.zeros( - (bs * self.stride, 1), dtype=input_ids.dtype, device=input_ids.device - ) - scatter_compact_to_strided_into( - compact=input_ids.view(-1, 1), - verify_lens=verify_lens, - out=candidates, - stride=self.stride, - fill_value=0, - ) + def _accept( + self, *, candidates, logits, draft_tokens, seq_lens, cutoff_verify_lens=None + ) -> torch.Tensor: + bs = candidates.shape[0] correct_len, bonus, cap_trim_lens = accept_greedy_triton( - candidates=candidates.view(bs, self.stride), - target_logits=self.strided_logits[: bs * self.stride], + candidates=candidates, + target_logits=logits, verify_num_draft_tokens=self.stride, - cutoff_verify_lens=verify_lens, + cutoff_verify_lens=cutoff_verify_lens, + fused_argmax=self._fused_argmax, ) self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, correct_len) self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, bonus) @@ -658,7 +760,7 @@ class DsparkVerifyEpilogue: prefix_lens=seq_lens[:bs], ) out_tokens = BuildOutTokens.execute( - draft_tokens=self.draft_tokens_buf[: bs * self.gamma].view(bs, self.gamma), + draft_tokens=draft_tokens, correct_len=correct_len, bonus=bonus, verify_num_draft_tokens=self.stride, @@ -719,6 +821,7 @@ def accept_draft_tokens( gamma: int, verify_num_draft_tokens: int, cutoff_layout: Optional[RaggedVerifyLayout] = None, + fused_argmax: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: greedy_mask = draft_block.greedy_mask cutoff_verify_lens = None if cutoff_layout is None else cutoff_layout.verify_lens @@ -729,6 +832,7 @@ def accept_draft_tokens( target_logits=target_logits, verify_num_draft_tokens=verify_num_draft_tokens, cutoff_verify_lens=cutoff_verify_lens, + fused_argmax=fused_argmax, ) bs, gamma_rows, vocab = draft_block.corrected_logits.shape draft_probs = SoftmaxTemp.execute( @@ -753,6 +857,7 @@ def accept_draft_tokens( target_logits=target_logits, verify_num_draft_tokens=verify_num_draft_tokens, cutoff_verify_lens=cutoff_verify_lens, + fused_argmax=fused_argmax, ) sampling_len, sampling_bonus, sampling_trim = AcceptSampling.execute( candidates=candidates, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 5b8c2e734..31a7bb82d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -304,8 +304,23 @@ class DSparkWorkerV2(BaseSpecWorker): dp_moe_sync=self._draft_is_moe and get_parallel().enable_dp_attention, ) self._verify_epilogue = None + target_is_dsv41 = ( + getattr( + self.target_worker.model_runner.model_config.hf_text_config, + "model_type", + None, + ) + == "deepseek_v41" + ) + static_epilogue_supported = ( + target_is_dsv41 + and self._verify_planner.mode_value == "static" + and self._draft_is_moe + and not get_parallel().enable_dp_attention + and self.ps.pp_size == 1 + ) if ( - self._verify_planner.is_compact_mode + (self._verify_planner.is_compact_mode or static_epilogue_supported) and self._decode_graph_allowed and is_cuda() ): @@ -314,6 +329,7 @@ class DSparkWorkerV2(BaseSpecWorker): verify_num_draft_tokens=self.verify_num_draft_tokens, device=self.device, tp_sync=self._tp_sync, + fused_argmax=target_is_dsv41, commit_ctx=CommitInjectCtx( draft_model=self.draft_model, block_pos_offsets=self._block_pos_offsets, @@ -321,6 +337,7 @@ class DSparkWorkerV2(BaseSpecWorker): resolve_req_to_token=lambda: ( self.model_runner.req_to_token_pool.req_to_token ), + kv_injector=self._kv_injector, ), ) self.model_runner.capture_tail_hooks.append( @@ -512,7 +529,11 @@ class DSparkWorkerV2(BaseSpecWorker): batch: ScheduleBatch, on_publish=None, grammar_barrier=None, + *, + pp_proxy_tensors=None, ) -> GenerationBatchResult: + # The non-overlap scheduler passes this keyword even when PP=1. + assert pp_proxy_tensors is None, "DSpark does not support pipeline parallelism" if batch.forward_mode.is_extend() or batch.is_extend_in_batch: self._verify_planner.note_non_decode_step() self._observers.note_prefill_step() @@ -590,9 +611,17 @@ class DSparkWorkerV2(BaseSpecWorker): final_pos = torch.repeat_interleave( (draft_seq_lens + ctx_lens - 1).to(torch.int64), repeats ) + cache_loc = batch.out_cache_loc + token_indices = logits_output.hidden_states_token_indices + if token_indices is not None: + cache_loc = cache_loc[token_indices] + positions = positions[token_indices] + if state_slot is not None: + state_slot = state_slot[token_indices] + final_pos = final_pos[token_indices] self._kv_injector.inject_target_hidden( target_hidden=logits_output.hidden_states, - cache_loc=batch.out_cache_loc, + cache_loc=cache_loc, positions=positions, state_slot=state_slot, final_pos=final_pos, @@ -600,6 +629,7 @@ class DSparkWorkerV2(BaseSpecWorker): ) # Avoid copying large hidden-state buffers to CPU in overlap scheduling. logits_output.hidden_states = None + logits_output.hidden_states_token_indices = None batch_output.next_draft_input = make_next_draft_input( bonus_tokens=next_token_ids, @@ -780,6 +810,11 @@ class DSparkWorkerV2(BaseSpecWorker): inject_gate=fold_eligible, ) else: + if ( + self._verify_epilogue is not None + and self._verify_planner.mode_value == "static" + ): + self._verify_epilogue.begin_static_step(bs, fold_eligible) target_verify = self._verify_executor.run_non_compact( batch=batch, draft_input=draft_input, @@ -804,7 +839,11 @@ class DSparkWorkerV2(BaseSpecWorker): grammar_mask.apply(logits_output.next_token_logits) epilogue = self._verify_executor.verify_epilogue - folded_accept = fold_eligible and run_compact and can_run_cuda_graph + folded_accept = ( + fold_eligible + and can_run_cuda_graph + and (run_compact or self._verify_planner.mode_value == "static") + ) accept = self._verify_executor.accept_and_finalize( folded_accept=folded_accept, bs=bs, @@ -817,6 +856,11 @@ class DSparkWorkerV2(BaseSpecWorker): prefix_lens=prefix_lens, draft_tokens=draft_tokens, ) + self.model_runner.ngram_embedding_manager.update_after_verify( + verify_ids_2d=verify_ids_2d, + req_pool_indices=batch.req_pool_indices, + commit_lens=accept.commit_lens, + ) if batch.return_logprob: compute_spec_logprobs( batch, diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index c0ab4d66c..6cd798913 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -91,6 +91,7 @@ from sglang.srt.configs import ( XllmConfig, ) from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config +from sglang.srt.configs.deepseek_v41 import DEEPSEEK_V41_CONFIG_CLASSES from sglang.srt.configs.internvl import InternVLChatConfig from sglang.srt.utils import get_bool_env_var, logger, lru_cache_frozenset from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri @@ -192,9 +193,28 @@ try: class _DeepseekV4ConfigAlias(_HFDeepseekV3Config): model_type = "deepseek_v4" + hc_pre_from_prev_sublayer = False + # V4 normalizes each attention query head (weightless rmsnorm) before RoPE. + q_head_norm = True + kv_source_layer_ids = () + index_source_layer_ids = () + candidate_source_layer_id = -1 + candidate_topk_blocks = 0 + candidate_block_size = 0 + engram_layer_ids = () + engram_num_embeddings = () + engram_max_ngram_size = 1 + engram_vocab_size = 0 + engram_n_heads = 0 + engram_head_dim = 0 + engram_pad_token_id = 2 + engram_compressed_vocab_size = 0 _CONFIG_REGISTRY["deepseek_v32"] = _DeepseekV32ConfigAlias _CONFIG_REGISTRY["deepseek_v4"] = _DeepseekV4ConfigAlias + _CONFIG_REGISTRY.update( + {cls.model_type: cls for cls in DEEPSEEK_V41_CONFIG_CLASSES} + ) # For kimi_k25_eagle3 class _KimiK2ConfigAlias(_HFDeepseekV3Config): diff --git a/python/sglang/srt/utils/hf_transformers/config.py b/python/sglang/srt/utils/hf_transformers/config.py index 1dc8be337..73f963d13 100644 --- a/python/sglang/srt/utils/hf_transformers/config.py +++ b/python/sglang/srt/utils/hf_transformers/config.py @@ -19,6 +19,10 @@ from typing import Optional from transformers import PretrainedConfig from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES +from sglang.srt.configs.deepseek_v41 import ( + DeepseekV41Config, + normalize_deepseek_v41_config, +) from sglang.srt.configs.model_config_parser_registry import ( ModelConfigParserBase, get_model_config_parser, @@ -179,6 +183,8 @@ class HfModelConfigParser(ModelConfigParserBase): _set_architectures(config, "DeepseekOCRForCausalLM") config = DeepseekVLV2Config.from_pretrained(model, revision=revision) _apply_deepseek_ocr_overrides(config, model) + elif isinstance(config, DeepseekV41Config): + config._name_or_path = model elif config.model_type in _CONFIG_REGISTRY: model_type = config.model_type if model_type == "deepseek_vl_v2" and is_ocr: @@ -315,6 +321,8 @@ def get_config( ) if model_override_args: + if isinstance(config, DeepseekV41Config): + model_override_args = normalize_deepseek_v41_config(model_override_args) # A plain update() setattrs a dict-valued override straight onto the # config, so '{"text_config": {...}}' on a VLM would replace the whole # sub-config with a dict and break attribute access downstream. diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 38ff80f81..a518097aa 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -282,6 +282,10 @@ class TinyDSV4ModelConfig: index_topk=DSV4_INDEX_TOPK, num_hidden_layers=len(compression_ratios), compress_ratios=list(compression_ratios), + # Ratio 1/2 layers are their own kv_source (one-layer fixtures). + kv_source_layer_ids=[ + i for i, ratio in enumerate(compression_ratios) if ratio in (1, 2) + ], ) self.hf_config.get_text_config = lambda: self.hf_config self.hf_text_config = self.hf_config @@ -411,6 +415,10 @@ class MockDSV4ModelRunner: device=device, enable_memory_saver=False, compression_ratios=list(compression_ratios), + kv_source_layers=model_config.hf_config.kv_source_layer_ids, + # Full locs are the identity-mapped SWA locs below, so the c1/c2 + # latent pools (slot = loc // ratio) only need to span swa_size. + full_size=swa_size, ) # Register identity full->swa mapping over swa_size full locs. identity = torch.arange(swa_size, dtype=torch.int64, device=device) @@ -1106,20 +1114,20 @@ def prepare_dsv4_runner_inputs( # reference needs to build that metadata itself. Stash the current batch # so `_pure_torch_dsv4_combined_reference` knows which one to use. fixture._current_batch = batch # type: ignore[attr-defined] - if case.compress_ratio in (4, 128): + if case.compress_ratio in (1, 2, 4, 128): _populate_extra_kv_cache(fixture, layer_id=0, num_entries=_DSV4_EXTRA_ENTRIES) def _seed_c4_if_needed( - fixture: DSV4AttentionFixture, *, num_entries: int = _DSV4_EXTRA_ENTRIES + fixture: DSV4AttentionFixture, *, num_entries: int | None = None ) -> None: - """For compress_ratio=4, seed the C4 metadata the exercised path consumes - (the C4Indexer would normally populate it; the compact fixture skips the - indexer): `c4_sparse_page_indices` for the dense extend path, - `c4_sparse_raw_indices` for sparse prefill. No-op for other compress_ratios. + """Seed `c4_sparse_page_indices` (dense extend) or `c4_sparse_raw_indices` + (sparse prefill); the compact fixture skips the indexer that fills them. """ - if fixture.case.compress_ratio != 4: + if fixture.case.compress_ratio not in (1, 2, 4): return + if num_entries is None: + num_entries = getattr(fixture, "extra_entries", _DSV4_EXTRA_ENTRIES) if fixture.seed_c4_for_sparse_prefill: _seed_c4_sparse_prefill_indices(fixture, num_entries=num_entries) else: @@ -1138,7 +1146,7 @@ def run_dsv4_fixture_eager(fixture: DSV4AttentionFixture) -> torch.Tensor: full_kv_locs_per_req = _populate_swa_kv_cache( fixture, max_context_len=max_context_len, device=runner.device ) - if case.compress_ratio in (4, 128): + if case.compress_ratio in (1, 2, 4, 128): _populate_extra_kv_cache(fixture, layer_id=0, num_entries=_DSV4_EXTRA_ENTRIES) q_input, _ = fixture.actual_module.project(fixture.input_hidden) with torch.no_grad(), forward_context(ForwardContext(attn_backend=fixture.backend)): @@ -1202,7 +1210,7 @@ def expected_dsv4_output_from_inputs( runner = fixture.runner max_context_len = runner.req_to_token_pool.req_to_token.shape[1] q_input, _ = fixture.actual_module.project(inputs["input_hidden"]) - if case.compress_ratio in (4, 128): + if case.compress_ratio in (1, 2, 4, 128): return _pure_torch_dsv4_combined_reference(fixture, q_input).float() full_kv_locs_per_req = _full_kv_locs_per_req( case, max_context_len=max_context_len, device=runner.device @@ -1319,7 +1327,7 @@ def _pure_torch_dsv4_combined_reference( swa_indices = md.swa_page_indices # [num_q, padded_window], full-pool locs swa_topk_lengths = md.swa_topk_lengths # [num_q] - if case.compress_ratio in (4, 128): + if case.compress_ratio in (1, 2, 4, 128): extra_indices, extra_topk_lengths = _extra_metadata_indices( md, case.compress_ratio ) @@ -1443,9 +1451,9 @@ def _seed_c4_sparse_prefill_indices( lens = (md.positions_casual + 1) // ratio max_len = int(lens.max().item()) pool = fixture.runner.token_to_kv_pool - c4_page_size = pool.get_extra_key_page_size(layer_id=0) - assert max_len <= min(num_entries, c4_page_size), ( - f"case attends {max_len} c4 entries; only {min(num_entries, c4_page_size)} populated" + c_page_size = pool.get_extra_key_page_size(layer_id=0) + assert max_len <= min(num_entries, c_page_size), ( + f"case attends {max_len} c{ratio} entries; only {min(num_entries, c_page_size)} populated" ) assert (md.page_table[:, 0] == 0).all(), ( "sparse seeding requires the raw==physical identity (first page 0)" @@ -1504,7 +1512,7 @@ def run_dsv4_target_verify_attention_case( testcase.assertEqual(fixture.backend.max_context_len, max_context_len) _populate_swa_kv_cache(fixture, max_context_len=max_context_len, device=device) - if case.compress_ratio in (4, 128): + if case.compress_ratio in (1, 2, 4, 128): _populate_extra_kv_cache(fixture, layer_id=0, num_entries=_DSV4_EXTRA_ENTRIES) _prepare_target_verify_batch(fixture.forward_batch, case, device) @@ -1614,28 +1622,21 @@ def run_dsv4_compress_attention_case( dtype: torch.dtype = torch.bfloat16, device: str = "cuda", ) -> None: - """Math-faithful test for the SWA + C4 (compress_ratio=4) / SWA + C128 - (compress_ratio=128) path through `DeepseekV4AttnBackend.forward`. - - Pre-writes random packed K into both the SWA cache and the extra - (C4/C128) cache via the production pack+set paths, lets - `init_forward_metadata` populate the compression metadata, manually seeds - the C4 metadata the exercised path consumes (see `_seed_c4_if_needed`; the - un-run indexer would otherwise leave it at `-1` / uninitialized), then - dispatches `forward(compress_ratio=case.compress_ratio)` and compares - against an independent pure-PyTorch SWA + extra reference that reads the - SAME cache bytes and metadata indices. - - `sparse_prefill` pins `SGLANG_OPT_FLASHMLA_SPARSE_PREFILL`, selecting the - dense `flash_mla_with_kvcache` extend path or `_forward_prefill_sparse`; - the C4 seeding dispatches on the same flag. + """SWA + compressed-cache path (compress ratios 1, 2, 4, 128) through + `DeepseekV4AttnBackend.forward` against a pure-PyTorch reference that reads the + same cache bytes and metadata indices. `sparse_prefill` pins + `SGLANG_OPT_FLASHMLA_SPARSE_PREFILL`; the C4 seeding dispatches on the same flag. """ - assert case.compress_ratio in ( - 4, - 128, - ), ( - f"DSV4 compact runner requires compress_ratio in (4, 128); got {case.compress_ratio}" + assert case.compress_ratio in (1, 2, 4, 128), ( + f"DSV4 compact runner requires compress_ratio in (1, 2, 4, 128); " + f"got {case.compress_ratio}" ) + # The sparse-prefill seeding attends (pos + 1) // ratio entries per query, so + # the low ratios need more populated entries than the default 32. + if case.compress_ratio in (1, 2): + extra_entries = max( + extra_entries, max(case.seq_lens) // case.compress_ratio + 1 + ) if sparse_prefill: assert case.forward_mode.is_extend_without_speculative(), ( f"sparse prefill only serves extend; got {case.forward_mode}" @@ -1648,6 +1649,7 @@ def run_dsv4_compress_attention_case( compression_ratios=[case.compress_ratio], ) fixture.seed_c4_for_sparse_prefill = sparse_prefill + fixture.extra_entries = extra_entries # type: ignore[attr-defined] runner = fixture.runner max_context_len = runner.req_to_token_pool.req_to_token.shape[1] diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index d0974039c..82d4142b9 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py @@ -126,6 +126,26 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase): extend_lens=(16,), compress_ratio=128, ), + DSV4AttentionCase( + name="dsv4_c2_extend", + backend="dsv4", + forward_mode=ForwardMode.EXTEND, + num_heads=64, + page_size=DSV4_PAGE_SIZE, + # Odd lengths: the ratio-2 causal count (pos + 1) // 2 rounds down. + prefix_lens=(33,), + extend_lens=(7,), + compress_ratio=2, + ), + DSV4AttentionCase( + name="dsv4_c2_decode", + backend="dsv4", + forward_mode=ForwardMode.DECODE, + num_heads=64, + page_size=DSV4_PAGE_SIZE, + prefix_lens=(65,), + compress_ratio=2, + ), DSV4AttentionCase( name="dsv4_c128_decode", backend="dsv4", @@ -533,6 +553,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): backend.model_runner = SimpleNamespace( spec_algorithm=SpeculativeAlgorithm.DFLASH ) + backend.token_to_kv_pool = SimpleNamespace(request_window=None) backend.forward_metadata = DSV4Metadata( self._make_core_metadata(0), indexer_metadata=None ) @@ -577,6 +598,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): backend.model_runner = SimpleNamespace( spec_algorithm=SpeculativeAlgorithm.DFLASH ) + backend.token_to_kv_pool = SimpleNamespace(request_window=None) backend.forward_metadata = DSV4Metadata( self._make_core_metadata(0), indexer_metadata=None ) @@ -794,7 +816,8 @@ class TestDSV4SwaOutCacheLocResolution(CustomTestCase): backend = object.__new__(DeepseekV4AttnBackend) backend.forward_metadata = None backend.token_to_kv_pool = SimpleNamespace( - translate_loc_from_full_to_swa=lambda loc: mapping[loc] + translate_loc_from_full_to_swa=lambda loc: mapping[loc], + request_window=None, ) return backend diff --git a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py index cebc016ec..9080914ba 100644 --- a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py +++ b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py @@ -175,7 +175,9 @@ def _make_backend( dsv4_prefill_backend: str = "auto", ) -> DeepseekV4AttnBackend: backend = DeepseekV4AttnBackend.__new__(DeepseekV4AttnBackend) - backend.forward_metadata = SimpleNamespace(sparse_prefill_cache=None) + backend.forward_metadata = SimpleNamespace( + sparse_prefill_cache=None, late_layer_tail=None + ) backend.req_to_token = req_to_token backend.sparse_prefill_workspace = SparsePrefillWorkspace(device) backend.softmax_scale = 512**-0.5 diff --git a/test/registered/spec/dspark/test_dspark_draft_path_default.py b/test/registered/spec/dspark/test_dspark_draft_path_default.py index c9cf44b58..b0da41feb 100644 --- a/test/registered/spec/dspark/test_dspark_draft_path_default.py +++ b/test/registered/spec/dspark/test_dspark_draft_path_default.py @@ -1,6 +1,8 @@ import unittest from types import SimpleNamespace +import torch + from sglang.srt.arg_groups.overrides import resolution_result from sglang.srt.arg_groups.speculative_hook import ( _handle_dspark, @@ -120,5 +122,37 @@ class TestDsparkDpAttentionMoeA2aGate(CustomTestCase): _handle_dspark(server_args) +class TestDsparkFoldedSamplingDefault(CustomTestCase): + def test_sharded_greedy_default_and_sampling_override(self): + from sglang.srt.environ import DsparkFoldedSampling, envs + from sglang.srt.speculative.dspark_components.dspark_draft_sampler import ( + _resolve_folded_sampling, + ) + + model = SimpleNamespace( + lm_head=SimpleNamespace(org_vocab_size=128, weight=torch.empty(1)), + markov_head=SimpleNamespace(supports_sharded_greedy=True), + ) + args = dict( + model=model, + gamma=5, + max_bs=64, + device="cpu", + tp_rank=0, + available_memory_gb=16, + ) + with envs.SGLANG_DSPARK_FOLDED_SAMPLING.override( + DsparkFoldedSampling.AUTO.value + ): + self.assertFalse(_resolve_folded_sampling(**args)) + model.markov_head.supports_sharded_greedy = False + self.assertTrue(_resolve_folded_sampling(**args)) + model.markov_head.supports_sharded_greedy = True + with envs.SGLANG_DSPARK_FOLDED_SAMPLING.override( + DsparkFoldedSampling.FORCE.value + ): + self.assertTrue(_resolve_folded_sampling(**args)) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py b/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py index f5a8b23b5..e4ef48bf9 100644 --- a/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py +++ b/test/registered/unit/disaggregation/test_admission_abort_not_enqueued.py @@ -86,6 +86,7 @@ def _prefill_queue(sched): def _decode_queue(sched): q = SimpleNamespace( scheduler=sched, + token_to_kv_pool_allocator=MagicMock(), retracted_queue=[], pending_reqs=[], _check_if_req_exceed_kv_capacity=MagicMock(return_value=False), diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 27e1332ae..8649b95ba 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -906,7 +906,7 @@ def _buf_infos(*ptrs): def _make_dsv4_target(*, unified, mapping=None): pool = object.__new__(DeepSeekV4TokenToKVPool) - pool.compression_ratios = [0, 4, 128] + pool.compression_ratios = [0, 2, 1, 4, 128] pool._unified_kv = unified pool.page_size = 256 pool.sliding_window = 128 diff --git a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py index 436f3934f..a36044db6 100644 --- a/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py +++ b/test/registered/unit/layers/quantization/test_fp8_blockwise_linear_backends.py @@ -7,6 +7,7 @@ SM90 / SM100 / SM120. """ import unittest +from types import SimpleNamespace from unittest import mock import torch @@ -267,24 +268,26 @@ class TestMxfp8LinearBackends(_LinearBackendCheck): is_backend_supported.assert_called_once_with("cute-dsl", 107) +def _build_block32_layer(n: int, k: int, keep_plain_weight_layout: bool = False): + quant_config = Fp8Config( + is_checkpoint_fp8_serialized=True, + activation_scheme="dynamic", + weight_block_size=[32, 32], + scale_fmt="ue8m0", + ) + layer = _make_linear(quant_config, n, k) + if keep_plain_weight_layout: + layer.keep_plain_weight_layout = True + w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10 + w_fp8, scale_e8m0, w_dequant = _quantize_fp8_block32_ue8m0(w) + load_linear_weights(layer, weight=w_fp8, weight_scale_inv=scale_e8m0) + return layer, w_dequant + + class TestBlockFp8AsMxfp8Linear(_LinearBackendCheck): """A 32-wide-K ue8m0 block-fp8 weight served through the MXFP8 GEMMs.""" - @staticmethod - def _build_layer(n: int, k: int, keep_plain_weight_layout: bool = False): - quant_config = Fp8Config( - is_checkpoint_fp8_serialized=True, - activation_scheme="dynamic", - weight_block_size=[32, 32], - scale_fmt="ue8m0", - ) - layer = _make_linear(quant_config, n, k) - if keep_plain_weight_layout: - layer.keep_plain_weight_layout = True - w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10 - w_fp8, scale_e8m0, w_dequant = _quantize_fp8_block32_ue8m0(w) - load_linear_weights(layer, weight=w_fp8, weight_scale_inv=scale_e8m0) - return layer, w_dequant + _build_layer = staticmethod(_build_block32_layer) def _run(self, backend: str): self._check_backend( @@ -339,6 +342,120 @@ class TestBlockFp8AsMxfp8Linear(_LinearBackendCheck): plain_layer.quant_method.apply(plain_layer, Mxfp8SwizzledInput(q, s)) +@unittest.skipUnless( + "flashinfer_cutedsl" in _block32_backends(), + "block-fp8-as-MXFP8 prefill tuning needs the FlashInfer CuTe-DSL kernel", +) +class TestBlockFp8AsMxfp8PrefillAutotune(_LinearBackendCheck): + """The startup hook that tunes those layers for the prefill M buckets.""" + + def setUp(self): + super().setUp() + patcher = mock.patch.object( + fp8_utils, + "FP8_GEMM_RUNNER_BACKEND", + Fp8GemmRunnerBackend.FLASHINFER_CUTEDSL, + ) + patcher.start() + self.addCleanup(patcher.stop) + torch.manual_seed(7) + + @staticmethod + def _ready_layer(n: int, k: int, keep_plain_weight_layout: bool = False): + layer, _ = _build_block32_layer(n, k, keep_plain_weight_layout) + layer.quant_method.process_weights_after_loading(layer) + return layer + + def test_model_hook_deduplicates_ready_block_fp8_weights(self): + from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM + + layers = torch.nn.ModuleList() + methods = [] + for _ in range(2): + layer = self._ready_layer(128, 128) + methods.append(layer.quant_method) + layer.quant_method.apply = mock.Mock() + layers.append(layer) + # An unprepared layer intentionally has no swizzled scale buffer. + fallback = self._ready_layer(128, 128, keep_plain_weight_layout=True) + 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): + layer = self._ready_layer(128, 128) + method = layer.quant_method + method.mxfp8_prefill_autotune_min_tokens = 4096 + call = mock.Mock(return_value=torch.empty(0)) + method.w8a8_mxfp8_linear = call + for rows, invariant, deterministic, expected in ( + (6, False, False, None), + (4096, False, False, False), + (4096, True, False, True), + (4096, False, True, True), + ): + with self.subTest( + rows=rows, invariant=invariant, deterministic=deterministic + ): + with ( + mock.patch( + "sglang.srt.batch_invariant_ops.is_batch_invariant_mode_enabled", + return_value=invariant, + ), + mock.patch( + "sglang.srt.runtime_context.get_exec", + return_value=SimpleNamespace( + deterministic=SimpleNamespace( + enable_deterministic_inference=deterministic + ) + ), + ), + ): + method.apply(layer, torch.empty(rows, 128, device="cuda")) + self.assertEqual(call.call_args.kwargs.get("pin_tactic"), expected) + + def test_prefill_tuning_leaves_decode_bit_identical(self): + """Tuning the prefill buckets must not move the decode tactic: below the + stamped min_tokens the output has to stay bit-for-bit what it was.""" + from flashinfer.autotuner import autotune + + from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM + + runtime_patch = mock.patch( + "sglang.srt.runtime_context.get_exec", + return_value=SimpleNamespace( + deterministic=SimpleNamespace(enable_deterministic_inference=False) + ), + ) + runtime_patch.start() + self.addCleanup(runtime_patch.stop) + layer = self._ready_layer(1792, 5120) + method = layer.quant_method + x = torch.randn(6, 5120, device="cuda", dtype=torch.bfloat16) + original = method.apply(layer, x) + model = SimpleNamespace( + config=SimpleNamespace(model_type="deepseek_v41"), + model=torch.nn.ModuleList([layer]), + ) + with autotune(True): + DeepseekV4ForCausalLM.autotune_prefill_kernels( + model, 4096, dtype=torch.bfloat16 + ) + self.assertEqual(method.mxfp8_prefill_autotune_min_tokens, 4096) + torch.testing.assert_close(method.apply(layer, x), original, rtol=0, atol=0) + + @unittest.skipIf(get_device_sm() < 90, "FP8 GEMM backends require SM90+") class TestModeloptFp8PerTensorLinear(_LinearBackendCheck): """Per-tensor FP8 (ModelOptFp8LinearMethod, static scales) on the auto diff --git a/test/registered/unit/layers/quantization/test_mxfp4_trtllm_padding.py b/test/registered/unit/layers/quantization/test_mxfp4_trtllm_padding.py new file mode 100644 index 000000000..1497cbadd --- /dev/null +++ b/test/registered/unit/layers/quantization/test_mxfp4_trtllm_padding.py @@ -0,0 +1,114 @@ +"""A TP-sharded MXFP4 trtllm-gen MoE whose per-rank intermediate size needs +padding must sum to the unsharded experts' output.""" + +import unittest +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput +from sglang.srt.layers.moe.topk import StandardTopKOutput +from sglang.srt.layers.quantization import mxfp4_flashinfer_trtllm_moe as mxfp4 +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +def make_layer(weights): + layer = torch.nn.Module() + names = ( + "w13_weight", + "w2_weight", + "w13_weight_scale_inv", + "w2_weight_scale_inv", + ) + for name, tensor in zip(names, weights): + layer.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False)) + layer.num_experts = weights[0].shape[0] + layer.num_local_experts = layer.num_experts + layer.moe_ep_rank = 0 + return layer + + +def make_weights(intermediate, hidden=256, device="cpu"): + experts = 8 + + def fp4_packed(*shape): + return torch.randint(-128, 128, shape, dtype=torch.int8, device=device) + + def e8m0_scales(*shape): + return torch.randint(-6, -3, shape, device=device).float().exp2() + + return ( + fp4_packed(experts, 2 * intermediate, hidden // 2), + fp4_packed(experts, hidden, intermediate // 2), + e8m0_scales(experts, 2 * intermediate, hidden // 32), + e8m0_scales(experts, hidden, intermediate // 32), + ) + + +class TestMxfp4TrtllmPadding(CustomTestCase): + @unittest.skipUnless( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10, + "Requires Blackwell", + ) + def test_tp4_matches_unsharded_experts(self): + torch.manual_seed(42) + weights = make_weights(2304, hidden=5120, device="cuda") + + def prepare(tensors): + layer = make_layer(tensors) + method = object.__new__(mxfp4.Mxfp4FlashinferTrtllmMoEMethod) + method._fp8 = Mock() + method.prefix = "test.experts" + method.flashinfer_mxfp4_moe_precision = "default" + method.process_weights_after_loading(layer) + method.create_moe_runner(layer, SimpleNamespace(swiglu_limit=10.0)) + return method, layer + + full = prepare([tensor.clone() for tensor in weights]) + shards = [] + for rank in range(4): + start = rank * 576 + end = start + 576 + w13, w2, s13, s2 = weights + shard = ( + torch.cat( + (w13[:, start:end], w13[:, 2304 + start : 2304 + end]), dim=1 + ), + w2[..., start // 2 : end // 2].contiguous(), + torch.cat( + (s13[:, start:end], s13[:, 2304 + start : 2304 + end]), dim=1 + ), + s2[..., start // 32 : end // 32].contiguous(), + ) + shards.append(prepare(shard)) + + with ( + patch.object(mxfp4, "get_tp_group", return_value=None), + patch.object(mxfp4, "is_allocation_symmetric", return_value=False), + patch.object(mxfp4, "use_symmetric_memory", return_value=nullcontext()), + ): + for tokens in (1, 64): + with self.subTest(tokens=tokens): + x = torch.randn(tokens, 5120, dtype=torch.bfloat16, device="cuda") + logits = torch.randn(tokens, 8, device="cuda") + scores, ids = logits.softmax(-1).topk(6, dim=-1) + topk = StandardTopKOutput(scores, ids.to(torch.int32), logits) + dispatch = StandardDispatchOutput(x, None, topk) + reference = full[0].apply(full[1], dispatch).hidden_states.float() + actual = sum( + method.apply(layer, dispatch).hidden_states.float() + for method, layer in shards + ) + rmse = torch.linalg.norm(actual - reference) / torch.linalg.norm( + reference + ) + self.assertLess(rmse.item(), 0.01) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py index c7807d478..c427c7b8c 100644 --- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py +++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py @@ -531,5 +531,26 @@ class TestDSV4NonPagedIndexer(CustomTestCase): self.assertEqual(call.kwargs, {"clean_logits": False, "max_seqlen_k": 128}) +class TestCandidateIndexerGating(CustomTestCase): + def test_candidate_indexer_gating(self): + from sglang.srt.layers.attention.dsv4 import candidate_indexer + + def platform(sm): + return patch.object( + candidate_indexer, "get_platform", lambda: SimpleNamespace(device_sm=sm) + ) + + flag = "sglang.srt.layers.deep_gemm_wrapper.configurer.DEEPGEMM_PAGED_SPARSE_MQA_LOGITS" + # V4 models have no candidate source; Hopper selects through masks inline. + with platform(100), patch(flag, True): + self.assertIsNone(candidate_indexer.make_candidate_indexer(0, 8)) + with platform(90), patch(flag, False): + self.assertIsNone(candidate_indexer.make_candidate_indexer(2048, 8)) + # Blackwell without DeepGEMM's sparse logits fails instead of falling back. + with platform(100), patch(flag, False): + with self.assertRaises(RuntimeError): + candidate_indexer.make_candidate_indexer(2048, 8) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 6bd4cea33..ac4453d99 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -529,6 +529,8 @@ class TestDecodePrebuilt(unittest.TestCase): scheduler.policy = MagicMock() scheduler.schedule_stream = MagicMock() scheduler.forward_stream = MagicMock() + scheduler.ngram_embedding_manager = MagicMock() + scheduler.chunked_req = None return scheduler def test_waiting_queue_is_sorted_before_prebuilt_selection(self): diff --git a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py index 65eaf1c79..0b345f3be 100644 --- a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py +++ b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py @@ -5,11 +5,17 @@ from unittest.mock import MagicMock, patch import torch +from sglang.kernels.ops.attention.dsv4.kv_layout import ( + KVLayout, + is_valid_kv_layout_pair, +) +from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( DeepSeekV4SingleKVPool, DeepSeekV4TokenToKVPool, _CompressedPoolConfig, ) +from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -25,6 +31,8 @@ class TestDSV4CompressedPools(CustomTestCase): pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool) pool._unified_kv = unified pool.uniform_fp8 = False + pool.kv_layout = KVLayout.V4 + pool.compressed_kv_layout_option = None pool.compressed_pool_configs = { 4: _CompressedPoolConfig( 256, 64, torch.bfloat16, indexer_size=1024 @@ -176,5 +184,151 @@ class TestDSV4CompressedPools(CustomTestCase): pool.get_index_k_page_size(128) +HEAD_DIM = 512 +ROPE_DIM = 64 +PAGE_SIZE = 256 +FULL_SIZE = 4 * PAGE_SIZE + + +class TestV41KVPoolLayouts(CustomTestCase): + """A V4.1-layout pool hands the attention kernel page-aligned buffers and + picks the compressed layout each ratio asks for.""" + + def setUp(self): + super().setUp() + override = get_context().override_server_args(page_size=PAGE_SIZE) + override.install() + self.addCleanup(override.restore) + + def make_pool(self, ratios, kv_source_layers, kv_layout, compressed=None, **sizes): + return DeepSeekV4TokenToKVPool( + max_num_reqs=16, + swa_size=FULL_SIZE, + c4_size=sizes.get("c4_size", 0), + c128_size=sizes.get("c128_size", 0), + c4_state_pool_size=sizes.get("c4_state_pool_size", 0), + c128_state_pool_size=sizes.get("c128_state_pool_size", 0), + page_size=PAGE_SIZE, + swa_page_size=PAGE_SIZE, + dtype=torch.float8_e4m3fn, + c4_state_dtype=torch.float32, + c128_state_dtype=torch.float32, + qk_nope_head_dim=HEAD_DIM - ROPE_DIM, + qk_rope_head_dim=ROPE_DIM, + indexer_head_dim=128, + layer_num=len(ratios), + device="cpu", + enable_memory_saver=False, + compression_ratios=ratios, + kv_source_layers=kv_source_layers, + full_size=FULL_SIZE, + kv_layout=kv_layout, + compressed_kv_layout=compressed, + ) + + def assert_kernel_requirements(self, pool, layout): + """Pages start on the kernel's alignment, and its + (num_pages, page_size, 1, bytes_per_token) view walks one token per row.""" + for buf in pool.kv_buffer: + self.assertEqual(buf.stride(0) % layout.page_align, 0) + bpt = layout.bytes_per_token + view = buf[:, : pool.page_size * bpt].view( + buf.shape[0], pool.page_size, 1, bpt + ) + self.assertEqual(view.stride(1), bpt) + self.assertEqual(view.stride(0), pool.bytes_per_page_padded) + + def test_v41_pool_buffers(self): + for option, expect in ((None, KVLayout.V41_FP4), ("fp8", KVLayout.V41)): + with self.subTest(compressed=option): + pool = self.make_pool([0, 0, 2, 1, 1], [2, 3], KVLayout.V41, option) + self.assert_kernel_requirements(pool.swa_kv_pool, KVLayout.V41) + self.assertEqual(pool.get_swa_key_bytes_per_token(), 528) + for ratio in (1, 2): + layer_id = pool.sources_by_ratio[ratio][0] + self.assertIs(pool.get_extra_key_layout(layer_id), expect) + self.assertEqual( + pool.get_extra_key_bytes_per_token(layer_id), + expect.bytes_per_token, + ) + self.assertTrue(is_valid_kv_layout_pair(pool.kv_layout, expect)) + self.assert_kernel_requirements(pool.kv_pools[ratio], expect) + # A pool of the fp4 layout cannot be the main cache. + with self.assertRaises(AssertionError): + self.make_pool([0], [], KVLayout.V41_FP4) + + def test_v41_pool_with_c4_c128(self): + pool = self.make_pool( + [0, 4, 128], + [], + KVLayout.V41, + c4_size=PAGE_SIZE, + c128_size=PAGE_SIZE, + c4_state_pool_size=16, + c128_state_pool_size=16, + ) + for ratio in (4, 128): + self.assertEqual(pool.kv_pools[ratio].page_size, PAGE_SIZE // ratio) + # The 2-token c128 page is the only production page that pads. + self.assertEqual(pool.kv_pools[128].bytes_per_page_padded, 1536) + + +class TestPagedDSparkWithEncoderReplay(CustomTestCase): + def setUp(self): + super().setUp() + override = get_context().override_server_args( + enable_encoder_swa_bounded_replay=True, + speculative_algorithm="DSPARK", + speculative_num_draft_tokens=6, + speculative_dspark_block_size=5, + page_size=256, + max_running_requests=2, + chunked_prefill_size=256, + ) + override.install() + self.addCleanup(override.restore) + + def make_pool(self, *, draft): + return DeepSeekV4TokenToKVPool( + max_num_reqs=2, + num_req_slots=3, + swa_size=1024, + c4_size=0, + c128_size=0, + c4_state_pool_size=0, + c128_state_pool_size=0, + page_size=256, + swa_page_size=256, + dtype=torch.float8_e4m3fn, + c4_state_dtype=torch.float32, + c128_state_dtype=torch.bfloat16, + qk_nope_head_dim=448, + qk_rope_head_dim=64, + indexer_head_dim=128, + layer_num=3, + device="cpu", + enable_memory_saver=False, + compression_ratios=[0, 0, 0], + online_mtp_max_draft_tokens=6, + full_size=2048, + is_draft_worker=draft, + ) + + def test_target_window_and_draft_paged_storage_share_allocator_mapping(self): + target = self.make_pool(draft=False) + draft = self.make_pool(draft=True) + allocator = SWATokenToKVPoolAllocator( + 2048, 1024, 256, torch.float8_e4m3fn, "cpu", target, False + ) + draft.register_mapping(allocator.full_to_swa_index_mapping) + allocator.full_to_swa_index_mapping[256:512] = torch.arange(768, 1024) + self.assertEqual( + draft.translate_loc_from_full_to_swa( + torch.tensor([256, 300, 511]) + ).tolist(), + [768, 812, 1023], + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py b/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py index 90a6ce417..cbab1e06e 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py +++ b/test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py @@ -48,6 +48,7 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase): ) kvcache = _mock_kvcache(DeepSeekV4TokenToKVPool) + kvcache.swa_kv_pool = MagicMock() strategy = _select_strategy(kvcache, {FULL, SWA}) self.assertIsInstance(strategy, _DeepSeekV4Strategy) @@ -141,6 +142,7 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase): for cls in (SWAKVPool, DeepSeekV4TokenToKVPool): kvcache = _mock_kvcache(cls) + kvcache.swa_kv_pool = MagicMock() with self.assertRaises(AssertionError) as cm: _select_strategy(kvcache, {FULL}) self.assertIn("No matching HiCache strategy", str(cm.exception)) diff --git a/test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py b/test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py index f090f0be5..930450a83 100644 --- a/test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py +++ b/test/registered/unit/model_executor/runner/test_flashinfer_autotune_sync.py @@ -6,9 +6,10 @@ reduction holds only if ranks also enter tuning with the same cache, so these cover that gate and the digest it decides on. """ -from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci -register_cpu_ci(est_time=52, suite="base-a-test-cpu") +register_cpu_ci(est_time=57, suite="base-a-test-cpu") +register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="1-gpu-large") import json import multiprocessing @@ -16,11 +17,15 @@ import os import tempfile import traceback import unittest +from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace +from unittest.mock import Mock, patch +import torch import torch.distributed as dist +from sglang.srt.model_executor.runner import flashinfer_autotune as autotune from sglang.srt.model_executor.runner.flashinfer_autotune import ( _autotune_cache_digest, _autotune_tactic_sync_group, @@ -160,5 +165,113 @@ class TestDropDivergedAutotuneCache(CustomTestCase): ) +class TestModelPrefillAutotune(CustomTestCase): + """Model kernel warmup must cover prefill without a speculative dummy batch.""" + + 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) + # No dummy-buffer or attention APIs: this path must not build 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_declining_model_never_enters_the_autotune_context(self): + self.mr.model.wants_prefill_autotune = lambda: False + 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_extend_pass_is_opt_in(self): + # A draft worker keeps its own warmup; a model without the hook opts out. + for draft, has_hook in ((True, True), (False, False)): + with self.subTest(draft=draft, has_hook=has_hook): + self.mr.is_draft_worker = draft + if not has_hook: + del self.mr.model.autotune_prefill_kernels + autotune.maybe_flashinfer_autotune_extend( + self.runner, decode_num_tokens=384 + ) + self.hook.assert_not_called() + self.flashinfer_autotune_context.assert_not_called() + + +@unittest.skipUnless(torch.cuda.is_available(), "FlashInfer requires CUDA") +class TestAutotuneCachePhases(CustomTestCase): + """Loaded target tactics survive draft warmup, unless cache reuse is off.""" + + def test_target_and_draft_cache_reuse(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, draft = ( + Path(directory) / name for name in ("target.json", "draft.json") + ) + for path, key, tactic in ( + (target, "target_prefill", 7), + (draft, "draft_decode", 3), + ): + path.write_text( + json.dumps( + {"_metadata": _collect_metadata(), key: ["TestRunner", tactic]} + ) + ) + with ( + patch.object( + autotune, + "flashinfer_autotune_cache_path", + side_effect=[target, draft, draft], + ), + patch.object( + autotune, "get_flashinfer_autotune_skip_ops", return_value=set() + ), + autotune.envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.override(True), + ): + with autotune.flashinfer_autotune_context(runner, run_lm_head=False): + self.assertEqual( + tuner._file_configs["target_prefill"], ("TestRunner", 7) + ) + # No profiling: this models a restart that loads tactics from disk. + self.assertFalse(tuner.profiling_cache) + with autotune.flashinfer_autotune_context(runner, run_lm_head=False): + pass + saved = json.loads(draft.read_text()) + self.assertEqual(saved["target_prefill"], ["TestRunner", 7]) + self.assertEqual(saved["draft_decode"], ["TestRunner", 3]) + with ( + autotune.envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.override(False), + autotune.flashinfer_autotune_context(runner, run_lm_head=False), + ): + self.assertNotIn("target_prefill", tuner._file_configs) + self.assertNotIn("draft_decode", tuner._file_configs) + + 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 000000000..82e27cca4 --- /dev/null +++ b/test/registered/unit/model_executor/test_compensated_mhc_update_guard.py @@ -0,0 +1,61 @@ +"""Weight-update entry points refuse a model carrying compensated-mHC derived +weight caches before any weight is written.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +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_all_update_entries_reject_before_writes(self): + for field in ("_hc_attn_tf32_parts", "_hc_ffn_tf32_parts"): + for method, args in ( + ("update_weights_from_tensor", ([], "direct")), + ("update_weights_from_distributed", ([], [], [], "unused")), + ("update_weights_from_disk", ("unused", "auto")), + ("update_weights_from_ipc", (SimpleNamespace(),)), + ): + with self.subTest(field=field, method=method, args=args): + model = torch.nn.Sequential(torch.nn.Linear(1, 1)) + original = model[0].weight.detach().clone() + setattr(model[0], field, (torch.ones(1), torch.zeros(1))) + model.load_weights = Mock() + updater = SimpleNamespace( + get_model=lambda: model, _assert_weight_cache_inactive=Mock() + ) + with patch( + "sglang.srt.model_executor.model_runner_components.weight_updater.default_weight_loader" + ) as loader: + ok, message = getattr(WeightUpdater, method)(updater, *args) + self.assertFalse(ok) + self.assertIn("compensated mHC", message) + loader.assert_not_called() + model.load_weights.assert_not_called() + torch.testing.assert_close( + model[0].weight, original, rtol=0, atol=0 + ) + + def test_models_without_derived_splits_keep_update_support(self): + model = torch.nn.Sequential(torch.nn.Linear(1, 1)) + model[0]._hc_attn_tf32_parts = model[0]._hc_ffn_tf32_parts = None + with patch( + "sglang.kernels.ops.attention.dsv4.gemm.hpc_bf16xfp32_gemm_enabled", + return_value=False, + ): + self.assertIsNone(_unsupported_derived_weight_cache_error(model)) + self.assertIsNone(_unsupported_derived_weight_cache_error()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index caefcb1f8..fe7f3bb6c 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1152,6 +1152,9 @@ class TestSWAPoolFloor(CustomTestCase): cfg.c4_ring_size = 8 cfg.c4_shrink_factor = 1 cfg._unified = unified + cfg.operator_swa_ratio = None + cfg.swa_cap_tokens = None + cfg.swa_prefix_tails = 0 return cfg._compute_dsv4_sizes(max_tokens, page_size) def test_dsv4_rejects_single_page_pool(self): @@ -1214,6 +1217,11 @@ class TestSWAPoolFloor(CustomTestCase): cfg.disaggregation_mode = None cfg.disaggregation_decode_extra_slots = 0 cfg._unified = True + cfg.operator_swa_ratio = None + cfg.swa_cap_tokens = None + cfg.swa_prefix_tails = 0 + cfg.request_window_bytes = 0 + cfg.bytes_per_swa_token = 0.0 cfg._unified_fp8 = False # object.__new__ skips __init__; bf16 unified row is 2B * latent cfg._unified_row_bytes = cfg.attn_head_dim * 2 @@ -1231,6 +1239,50 @@ class TestSWAPoolFloor(CustomTestCase): + cfg._get_c128_state_fixed_bytes(max_running_requests) ) + def test_dsv4_paged_dspark_budget_reserves_window_and_draft_layers(self): + from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator + + _publish_config( + self, + enable_encoder_swa_bounded_replay=True, + speculative_algorithm="DSPARK", + speculative_num_draft_tokens=6, + speculative_dspark_block_size=5, + page_size=256, + max_running_requests=2, + chunked_prefill_size=256, + ) + cfg = SimpleNamespace( + qk_nope_head_dim=448, + qk_rope_head_dim=64, + index_head_dim=128, + context_len=131072, + compress_ratios=[0, 0] + [2] * 18 + [1] * 20, + window_size=128, + hf_config=SimpleNamespace(kv_source_layer_ids=[2, 8, 14, 20]), + ) + spec = SimpleNamespace(is_dspark=lambda: True, is_none=lambda: False) + kvc = SimpleNamespace( + kv_cache_dtype_str="fp8_e4m3", + model_config=cfg, + layer_info=SimpleNamespace(start_layer=0, end_layer=40), + ps=SimpleNamespace(pp_size=1, attn_dp_size=1), + sliding_window_size=128, + page_size=256, + spec_algorithm=spec, + spec_aux_config=SimpleNamespace(dflash_draft_num_layers=3), + ) + planner = DSV4PoolConfigurator(kvc) + self.assertEqual(planner.bytes_per_swa_token, 3 * 584) + budget = 256 * 1024 * 1024 + sizes = planner.calculate_pool_sizes(budget, 256) + self.assertEqual(sizes.swa_max_total_num_tokens, planner.swa_cap_tokens) + self.assertLessEqual( + sizes.full_max_total_num_tokens * planner.bytes_per_full_token + + planner._get_swa_fixed_bytes(), + budget, + ) + def test_dsv4_unified_c4_state_not_token_scaled(self): # Unified-KV sizes the c4 state ring from max_running_requests in # finalize_with_max_running_requests, so it must not scale here. diff --git a/test/registered/unit/models/test_deepseek_v4_rope_policy.py b/test/registered/unit/models/test_deepseek_v4_rope_policy.py index fa8113de8..bae5e90ec 100644 --- a/test/registered/unit/models/test_deepseek_v4_rope_policy.py +++ b/test/registered/unit/models/test_deepseek_v4_rope_policy.py @@ -55,6 +55,7 @@ class TestDeepseekV4RoPEPolicy(CustomTestCase): o_lora_rank=8, rms_norm_eps=1e-6, compress_ratios=[compress_ratio], + q_head_norm=True, rope_theta=10_000, compress_rope_theta=160_000, max_position_embeddings=128, diff --git a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py index 3ee807f37..fd47b414d 100644 --- a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py +++ b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py @@ -57,6 +57,10 @@ class TestDeepseekV4SharedExpertFusionPolicy(CustomTestCase): quantization_config={}, rope_scaling={}, compress_ratios=[], + kv_source_layer_ids=[], + index_source_layer_ids=[], + engram_layer_ids=[], + engram_num_embeddings=[], n_shared_experts=1, dspark_markov_rank=1, num_nextn_predict_layers=1, diff --git a/test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py b/test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py index 013763eb0..a2fbf1329 100644 --- a/test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py +++ b/test/registered/unit/models/test_deepseek_v4_unified_fp8_q_pair.py @@ -88,7 +88,7 @@ class _Harness(deepseek_v4.MQALayer): dtype=torch.bfloat16, ) ) - self.wo_b = lambda value: (value, None) + self.wo_b = lambda value, skip_all_reduce=False: (value, None) self.prepare_kwargs = None def _forward_prepare( diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 5fc30bcaf..948a7cba0 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1718,6 +1718,14 @@ class TestGoldenModelOverrides(_IsolatedPublish): "swa_full_tokens_ratio", _deepseek_v4_overrides(_args(swa_full_tokens_ratio=0.5), hf), ) + # V4.1 leaves the ratio unset (cap-mode SWA sizing). + hf41 = SimpleNamespace( + architectures=["DeepseekV4ForCausalLM"], model_type="deepseek_v41" + ) + self.assertNotIn( + "swa_full_tokens_ratio", + _deepseek_v4_overrides(_args(fp8_gemm_runner_backend="triton"), hf41), + ) # An explicit user choice takes precedence over the model default. self.assertNotIn( "moe_runner_backend",