[DSV4.1] Optimize prefill GEMMs and compensated mHC for concurrent serving

python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh

@@ -505,10 +505,12 @@ inline PrefillPlan plan_compress_prefill(
   const auto f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr());
 
   const auto batch_size = static_cast<uint32_t>(B.unwrap());
-  constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
+  // ragged_id is a zero-based uint16 index, so a 64K-token batch is valid.
+  constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max()) + 1;
   RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
   RuntimeCheck(!use_req_ring || compress_ratio == 4);
-  RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
+  // Keep batch_id below 65535: pack_w(65535, 65535, ...) is the invalid sentinel.
+  RuntimeCheck(batch_size < kMaxTokens && batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
   // `swa_page_size` >= `ring_size` >= `compress_ratio`
   RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
   // Write pad: trailing tokens kept resident so a verify batch's committed tail survives
@@ -750,9 +752,9 @@ inline PrefillPlan plan_compress_prefill_legacy(
 
   const auto window_size = compress_ratio * (is_overlap ? 2 : 1);
   const auto batch_size = static_cast<uint32_t>(B.unwrap());
-  constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
+  constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max()) + 1;
   RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
-  RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
+  RuntimeCheck(batch_size < kMaxTokens && batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
 
   uint32_t counter = 0;
   uint32_t counter_c = 0;

python/sglang/kernels/ops/attention/dsv4/q_rope_store.py

@@ -21,6 +21,43 @@ def _q_rope_store(X, Y, F, POS, SX: tl.constexpr, SY: tl.constexpr):
     tl.store(Y + row * SY + head * 512 + r, tl.where(r >= 448, rotated, value))
 
 
+@triton.jit
+def _q_rope_store_prefill(
+    X,
+    Y,
+    F,
+    POS,
+    M,
+    H: tl.constexpr,
+    SX: tl.constexpr,
+    SY: tl.constexpr,
+    BLOCK_HEADS: tl.constexpr,
+):
+    # Keep token count dynamic to avoid compiling every prefill batch length.
+    heads = tl.program_id(0) * BLOCK_HEADS + tl.arange(0, BLOCK_HEADS)
+    row, head = heads // H, heads % H
+    r = tl.arange(0, 512)
+    # The padded output can exceed 2 GiB at the 64K prefill ceiling.
+    x_offset = row[:, None].to(tl.int64) * SX + head[:, None] * 512 + r[None, :]
+    value = tl.load(X + x_offset, row[:, None] < M, 0).to(tl.float32)
+    partner = tl.gather(value, tl.broadcast_to((r ^ 1)[None, :], (BLOCK_HEADS, 512)), 1)
+    position = tl.load(POS + row, row < M, 0)
+    freq_offset = position[:, None].to(tl.int64) * 64 + (r[None, :] - 448) // 2 * 2
+    rope_mask = (row[:, None] < M) & (r[None, :] >= 448)
+    cos = tl.load(F + freq_offset, rope_mask, 0)
+    sin = tl.load(F + freq_offset + 1, rope_mask, 0)
+    # Keep the same arithmetic and BF16 rounding as the decode kernel.
+    even = tl.fma(value, cos, -partner * sin)
+    odd = tl.fma(partner, sin, value * cos)
+    rotated = tl.where((r[None, :] & 1) == 0, even, odd)
+    y_offset = row[:, None].to(tl.int64) * SY + head[:, None] * 512 + r[None, :]
+    tl.store(
+        Y + y_offset,
+        tl.where(r[None, :] >= 448, rotated, value),
+        row[:, None] < M,
+    )
+
+
 def q_rope_store(
     q: torch.Tensor,
     output: torch.Tensor,
@@ -35,6 +72,20 @@ def q_rope_store(
     assert freqs_cis.dtype == torch.complex64 and freqs_cis.is_contiguous()
     assert freqs_cis.shape[1] == 32 and positions.shape == (q.shape[0],)
     assert positions.dtype in (torch.int32, torch.int64) and positions.is_contiguous()
+    if q.shape[0] >= 4096 and q.shape[1] == 16:
+        _q_rope_store_prefill[(triton.cdiv(q.shape[0] * q.shape[1], 4),)](
+            q,
+            output,
+            torch.view_as_real(freqs_cis),
+            positions,
+            q.shape[0],
+            q.shape[1],
+            q.stride(0),
+            output.stride(0),
+            BLOCK_HEADS=4,
+            num_warps=4,
+        )
+        return
     _q_rope_store[(q.shape[0], q.shape[1])](
         q,
         output,

python/sglang/kernels/ops/layernorm/hc_combine_norm.py

@@ -22,16 +22,40 @@ def _hc_combine_norm(X, P, W, Y, SX: tl.constexpr, SP: tl.constexpr, EPS: tl.con
     tl.store(Y + row * 5120 + h, value * inv_rms * weight, mask)
 
 
+@triton.jit
+def _hc_combine_norm_prefill(
+    X, P, W, Y, SX: tl.constexpr, SP: tl.constexpr, EPS: tl.constexpr
+):
+    # Large batches have enough rows to use one CTA per row without repeating
+    # the combine and RMS reduction for each output partition.
+    row = tl.program_id(0).to(tl.int64)
+    h = tl.arange(0, 8192)
+    value = tl.full((8192,), 0, tl.float32)
+    for c in tl.static_range(4):
+        pre = tl.load(P + row * SP + c).to(tl.float32)
+        x = tl.load(X + row * SX + c * 5120 + h, h < 5120, 0).to(tl.float32)
+        value += x * pre
+    value = value.to(tl.bfloat16).to(tl.float32)
+    inv_rms = tl.rsqrt(tl.sum(value * value, 0) / 5120 + EPS)
+    weight = tl.load(W + h, h < 5120, 0).to(tl.float32)
+    tl.store(Y + row * 5120 + h, value * inv_rms * weight, h < 5120)
+
+
 def hc_combine_norm(
     x: torch.Tensor, pre: torch.Tensor, weight: torch.Tensor, eps: float
 ) -> torch.Tensor:
-    """Fuse four-stream combine and RMSNorm for small BF16 batches of width 5120."""
+    """Fuse four-stream combine and RMSNorm for BF16 batches of width 5120."""
     m = x.shape[0]
-    assert 0 < m <= 8 and x.shape == (m, 20480)
+    assert (0 < m <= 8 or 4096 <= m <= 65536) and x.shape == (m, 20480)
     assert pre.shape == (m, 4) and pre.stride(1) == 1
     assert weight.shape == (5120,) and weight.is_contiguous()
     assert x.dtype == weight.dtype == torch.bfloat16 and x.stride(1) == 1
     y = torch.empty((m, 5120), dtype=x.dtype, device=x.device)
+    if m >= 4096:
+        _hc_combine_norm_prefill[(m,)](
+            x, pre, weight, y, x.stride(0), pre.stride(0), eps, num_warps=4
+        )
+        return y
     # Four CTAs per row trade redundant statistics for more concurrent loads
     # when only a few speculative tokens are being processed.
     _hc_combine_norm[(m, 4)](

python/sglang/kernels/ops/layernorm/hc_mix_stats_bf16x3.py

@@ -0,0 +1,100 @@
+"""Compensated mHC prefill projection with a shared activation load.
+
+Keep three BF16 components of the FP32 weights and accumulate their products
+separately. The 16 fixed K slices bound FP32 accumulation error, as in the
+compensated DeepGEMM path, while avoiding its second activation read/reduction.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+def split_bf16_hc_weight(weight: torch.Tensor):
+    assert weight.dtype == torch.float32 and weight.is_contiguous()
+    high = weight.bfloat16()
+    residual = weight - high.float()
+    middle = residual.bfloat16()
+    low = (residual - middle.float()).bfloat16()
+    return high, middle, low
+
+
+@triton.jit
+def _hc_mix_stats_bf16x3(X, W_HI, W_MID, W_LO, MIX, SQ, M, BLOCK_M: tl.constexpr):
+    # M stays runtime-valued so variable prefill lengths reuse the same binary.
+    rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    cols = tl.arange(0, 32)
+    # 20480 input features / 16 independent slices.
+    start = tl.program_id(1) * 1280
+    ks = start + tl.arange(0, 64)
+    hi = tl.zeros((BLOCK_M, 32), tl.float32)
+    mid = tl.zeros((BLOCK_M, 32), tl.float32)
+    lo = tl.zeros((BLOCK_M, 32), tl.float32)
+    sq = tl.zeros((BLOCK_M,), tl.float32)
+    for block in range(20):
+        k = ks + block * 64
+        x = tl.load(
+            X + rows[:, None].to(tl.int64) * 20480 + k[None, :],
+            rows[:, None] < M,
+            0,
+        )
+        offsets = cols[None, :] * 20480 + k[:, None]
+        w_hi = tl.load(W_HI + offsets, cols[None, :] < 24, 0)
+        w_mid = tl.load(W_MID + offsets, cols[None, :] < 24, 0)
+        w_lo = tl.load(W_LO + offsets, cols[None, :] < 24, 0)
+        hi = tl.dot(x, w_hi, hi)
+        mid = tl.dot(x, w_mid, mid)
+        lo = tl.dot(x, w_lo, lo)
+        xf = x.to(tl.float32)
+        sq += tl.sum(xf * xf, 1)
+    offsets = (tl.program_id(1) * M + rows[:, None]) * 24 + cols[None, :]
+    tl.store(MIX + offsets, (hi + mid) + lo, (rows[:, None] < M) & (cols[None, :] < 24))
+    tl.store(SQ + tl.program_id(1) * M + rows, sq, rows < M)
+
+
+def hc_mix_stats_sinkhorn_bf16x3(
+    x: torch.Tensor,
+    weight_parts,
+    scale: torch.Tensor,
+    base: torch.Tensor,
+    sinkhorn_iters: int,
+    rms_eps: float,
+    hc_eps: float,
+):
+    from sglang.kernels.ops.layernorm.mhc import _hc_mix_reduce_sinkhorn_kernel
+
+    m = x.shape[0]
+    assert x.shape == (m, 20480) and x.is_contiguous()
+    assert x.dtype == torch.bfloat16 and 4096 <= m <= 65536
+    assert len(weight_parts) == 3
+    assert all(
+        w.shape == (24, 20480) and w.dtype == torch.bfloat16 and w.is_contiguous()
+        for w in weight_parts
+    )
+    mix = torch.empty((16, m, 24), device=x.device, dtype=torch.float32)
+    sq = torch.empty((16, m), device=x.device, dtype=torch.float32)
+    pre = torch.empty((m, 4), device=x.device, dtype=torch.float32)
+    post = torch.empty_like(pre)
+    comb = torch.empty((m, 4, 4), device=x.device, dtype=torch.float32)
+    _hc_mix_stats_bf16x3[(triton.cdiv(m, 128), 16)](
+        x, *weight_parts, mix, sq, m, 128, num_warps=4, num_stages=3
+    )
+    _hc_mix_reduce_sinkhorn_kernel[(m,)](
+        mix,
+        sq,
+        scale,
+        base,
+        pre,
+        post,
+        comb,
+        m,
+        1.0 / 20480,
+        rms_eps,
+        MIX=24,
+        HC=4,
+        NUM_SLICES=16,
+        ITERS=sinkhorn_iters,
+        EPS=hc_eps,
+        num_warps=1,
+    )
+    return pre, post, comb

python/sglang/kernels/ops/layernorm/hc_mix_stats_deepgemm.py

@@ -0,0 +1,67 @@
+"""Compensated FP32 mHC projections for SM100 batches with at least 128 rows.
+
+The small-row and batch-invariant paths remain in mhc.py. Native TF32 discards
+too much of the FP32 projection weights, so evaluate their high and residual
+components separately and bound accumulation length with a fixed split count.
+"""
+
+import torch
+
+_NUM_SPLITS = 16
+
+
+def split_tf32_hc_weight(weight: torch.Tensor):
+    assert weight.dtype == torch.float32 and weight.is_contiguous()
+    high = (weight.view(torch.int32) & -8192).view(torch.float32)
+    return high, weight - high
+
+
+def hc_mix_stats_sinkhorn_deepgemm(
+    x_flat: torch.Tensor,
+    weight_parts,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    sinkhorn_iters: int,
+    rms_eps: float,
+    hc_eps: float,
+):
+    from sglang.kernels.ops.layernorm.mhc import _hc_mix_reduce_sinkhorn_kernel
+    from sglang.srt.layers.deep_gemm_wrapper.entrypoint import tf32_hc_prenorm_gemm
+
+    assert x_flat.dtype == torch.bfloat16 and x_flat.is_contiguous()
+    m, k = x_flat.shape
+    high, low = weight_parts
+    assert k == 20480 and high.shape == low.shape == (24, k)
+    dev = x_flat.device
+    pre = torch.empty((m, 4), dtype=torch.float32, device=dev)
+    post = torch.empty_like(pre)
+    comb = torch.empty((m, 4, 4), dtype=torch.float32, device=dev)
+    if m == 0:
+        return pre, post, comb
+
+    mix_hi = torch.empty((_NUM_SPLITS, m, 24), dtype=torch.float32, device=dev)
+    mix_lo = torch.empty_like(mix_hi)
+    sq = torch.empty((_NUM_SPLITS, m), dtype=torch.float32, device=dev)
+    unused_sq = torch.empty_like(sq)
+    tf32_hc_prenorm_gemm(x_flat, high, mix_hi, sq, _NUM_SPLITS)
+    tf32_hc_prenorm_gemm(x_flat, low, mix_lo, unused_sq, _NUM_SPLITS)
+    _hc_mix_reduce_sinkhorn_kernel[(m,)](
+        mix_hi,
+        sq,
+        hc_scale,
+        hc_base,
+        pre,
+        post,
+        comb,
+        m,
+        1.0 / k,
+        rms_eps,
+        MIX=24,
+        HC=4,
+        NUM_SLICES=_NUM_SPLITS,
+        ITERS=sinkhorn_iters,
+        EPS=hc_eps,
+        part_mix_residual_ptr=mix_lo,
+        num_warps=1,
+    )
+    return pre, post, comb

python/sglang/kernels/ops/layernorm/mhc.py

@@ -2246,6 +2246,7 @@ def _hc_mix_reduce_sinkhorn_kernel(
     NUM_SLICES: tl.constexpr,
     ITERS: tl.constexpr,
     EPS: tl.constexpr,
+    part_mix_residual_ptr=None,
 ):
     """One CTA per row keeps the sinkhorn reductions two-dimensional.
     Per-row arithmetic follows the slice reduction, then the Triton sinkhorn.
@@ -2263,9 +2264,16 @@ def _hc_mix_reduce_sinkhorn_kernel(
     sq = tl.zeros([], dtype=tl.float32)
     for s in tl.static_range(NUM_SLICES):
         off = (s * m + row) * MIX
-        a_pre += tl.load(part_mix_ptr + off + j)
-        a_post += tl.load(part_mix_ptr + off + HC + j)
-        a_comb += tl.load(part_mix_ptr + off + 2 * HC + jj * HC + kk)
+        v_pre = tl.load(part_mix_ptr + off + j)
+        v_post = tl.load(part_mix_ptr + off + HC + j)
+        v_comb = tl.load(part_mix_ptr + off + 2 * HC + jj * HC + kk)
+        if part_mix_residual_ptr is not None:
+            v_pre += tl.load(part_mix_residual_ptr + off + j)
+            v_post += tl.load(part_mix_residual_ptr + off + HC + j)
+            v_comb += tl.load(part_mix_residual_ptr + off + 2 * HC + jj * HC + kk)
+        a_pre += v_pre
+        a_post += v_post
+        a_comb += v_comb
         sq += tl.load(part_sq_ptr + s * m + row)
     rsqrt = 1.0 / tl.sqrt(sq * inv_k + rms_eps)
 

python/sglang/srt/environ.py

@@ -1452,6 +1452,9 @@ class Envs:
 
     # Kernels and indexer
     SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True)
+    # Compensated mHC caches FP32 weight splits; online weight updates must be
+    # disabled while this explicitly selected serving optimization is active.
+    SGLANG_DSV41_COMPENSATED_MHC = EnvBool(False)
     SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True)
     SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True)
     SGLANG_OPT_USE_FLASHINFER_MHC = EnvBool(False)