dsv4.1: remaining model and runtime integration (#38798)

Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Xiaoyu Zhang <xiaoyu.zhang@radixark.ai>
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
Co-authored-by: Zhichen Zeng <zczeng@uw.edu>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>

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",

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

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

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

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

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

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

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

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(

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,

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:

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."
+                )

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.",

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

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:

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)
 

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.

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

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,

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

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,
+)

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(

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

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
 

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"]:

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)
 

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

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)

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

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",

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,
+    )

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

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

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

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

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)

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,

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:

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)
 
 

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

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

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

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

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.

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

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

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,

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

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]

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,

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

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,

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

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)
 

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,

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
             )
 

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,

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:

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)

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

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
 

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}"

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

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,

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
     )

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
 

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

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

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)

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}"

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(

python/sglang/srt/multimodal/dsv41/__init__.py

@@ -0,0 +1 @@
+"""Vision tower, image preprocessing and VL expert routing for DeepSeek-V4.1."""

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)

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)

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

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:

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,

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,

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

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.

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]