[DSV4.1] Reduce mHC, metadata and small-batch router overhead (#39704)

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

@@ -0,0 +1,190 @@
+// Prefill mHC post/combine/RMSNorm with the original BF16 intermediates.
+// Stage the collapsed row in shared memory to bound register use while preserving
+// the original Triton prefill normalization reduction and PTX arithmetic.
+#pragma once
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang {
+struct MhcPostCombineNormPrefillParams {
+  const bf16_t *x, *residual, *weight;
+  bf16_t *residual_out, *output;
+  const float *post, *comb, *pre;
+  float eps;
+};
+
+template <int Threads>
+__global__ __launch_bounds__(Threads) void mhc_post_combine_norm_prefill_kernel(
+    const __grid_constant__ MhcPostCombineNormPrefillParams p) {
+  using namespace device;
+  using V = AlignedVector<bf16x2_t, 4>;
+  __shared__ __align__(16) bf16_t collapsed[5120];
+  __shared__ float warp_sums[4];
+  __shared__ float inv_rms;
+  const int tid = threadIdx.x;
+  const int lane = tid % 32;
+  const int64_t row = blockIdx.x;
+  const float coeff = lane < 16   ? p.comb[row * 16 + lane]
+                      : lane < 20 ? p.post[row * 4 + lane - 16]
+                      : lane < 24 ? p.pre[row * 4 + lane - 20]
+                                  : 0.f;
+
+  // Vectorize loads but keep only one 8-element tile live at a time.
+#pragma unroll 1
+  for (int vid = tid; vid < 640; vid += Threads) {
+    V x, residual[4];
+    x.load(p.x + row * 5120, vid);
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+      residual[j].load(p.residual + row * 20480 + j * 5120, vid);
+    float combined[8] = {};
+#pragma unroll
+    for (int i = 0; i < 4; ++i) {
+      const float post = __shfl_sync(0xffffffff, coeff, 16 + i);
+      const float pre = __shfl_sync(0xffffffff, coeff, 20 + i);
+      const float c0 = __shfl_sync(0xffffffff, coeff, i);
+      const float c1 = __shfl_sync(0xffffffff, coeff, 4 + i);
+      const float c2 = __shfl_sync(0xffffffff, coeff, 8 + i);
+      const float c3 = __shfl_sync(0xffffffff, coeff, 12 + i);
+      V updated;
+#pragma unroll
+      for (int k = 0; k < 4; ++k) {
+        const auto xx = cast<fp32x2_t>(x[k]);
+        const auto r0 = cast<fp32x2_t>(residual[0][k]);
+        const auto r1 = cast<fp32x2_t>(residual[1][k]);
+        const auto r2 = cast<fp32x2_t>(residual[2][k]);
+        const auto r3 = cast<fp32x2_t>(residual[3][k]);
+        float a = __fmaf_rn(post, xx.x, __fmul_rn(c0, r0.x));
+        float b = __fmaf_rn(post, xx.y, __fmul_rn(c0, r0.y));
+        a = __fmaf_rn(c1, r1.x, a);
+        b = __fmaf_rn(c1, r1.y, b);
+        a = __fmaf_rn(c2, r2.x, a);
+        b = __fmaf_rn(c2, r2.y, b);
+        a = __fmaf_rn(c3, r3.x, a);
+        b = __fmaf_rn(c3, r3.y, b);
+        updated[k] = cast<bf16x2_t>(fp32x2_t{a, b});
+        const auto rounded = cast<fp32x2_t>(updated[k]);
+        combined[2 * k] = __fmaf_rn(pre, rounded.x, combined[2 * k]);
+        combined[2 * k + 1] = __fmaf_rn(pre, rounded.y, combined[2 * k + 1]);
+      }
+      updated.store(p.residual_out + row * 20480 + i * 5120, vid);
+    }
+    V rounded;
+#pragma unroll
+    for (int k = 0; k < 4; ++k)
+      rounded[k] = cast<bf16x2_t>(fp32x2_t{combined[2 * k], combined[2 * k + 1]});
+    rounded.store(collapsed, vid);
+  }
+  __syncthreads();
+
+  // Match the actual prefill Triton layout: 8 consecutive elements/thread,
+  // 128 threads, then stride 1024. The PTX folds each thread's elements in
+  // order, followed by XOR 16,8,4,2,1 and a four-warp XOR 2,1 reduction.
+  if (tid < 128) {
+    float sum = 0.f;
+#pragma unroll
+    for (int group = 0; group < 5; ++group) {
+      V v;
+      v.load(collapsed + group * 1024, tid);
+#pragma unroll
+      for (int k = 0; k < 4; ++k) {
+        const auto value = cast<fp32x2_t>(v[k]);
+        if (group == 0 && k == 0)
+          sum = __fadd_rn(__fmul_rn(value.y, value.y), __fmul_rn(value.x, value.x));
+        else {
+          sum = __fmaf_rn(value.x, value.x, sum);
+          sum = __fadd_rn(__fmul_rn(value.y, value.y), sum);
+        }
+      }
+    }
+#pragma unroll
+    for (int offset = 16; offset; offset >>= 1)
+      sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, offset));
+    if (lane == 0) warp_sums[tid / 32] = sum;
+  }
+  __syncthreads();
+  if (tid < 32) {
+    float sum = lane < 4 ? warp_sums[lane] : 0.f;
+    sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 2));
+    sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 1));
+    if (lane == 0) {
+      float mean, inverse;
+      asm("div.full.f32 %0, %1, %2;" : "=f"(mean) : "f"(sum), "f"(5120.f));
+      // Match Triton's unqualified PTX add: ptxas may contract the
+      // constant division's reciprocal multiply with this epsilon add.
+      asm("add.f32 %0, %1, %2;" : "=f"(mean) : "f"(mean), "f"(p.eps));
+      asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(inverse) : "f"(mean));
+      inv_rms = inverse;
+    }
+  }
+  __syncthreads();
+#pragma unroll 1
+  for (int vid = tid; vid < 640; vid += Threads) {
+    V v, w, out;
+    v.load(collapsed, vid);
+    w.load(p.weight, vid);
+#pragma unroll
+    for (int k = 0; k < 4; ++k) {
+      const auto value = cast<fp32x2_t>(v[k]);
+      const auto weight = cast<fp32x2_t>(w[k]);
+      out[k] = cast<bf16x2_t>(
+          fp32x2_t{__fmul_rn(__fmul_rn(value.x, inv_rms), weight.x), __fmul_rn(__fmul_rn(value.y, inv_rms), weight.y)});
+    }
+    out.store(p.output + row * 5120, vid);
+  }
+}
+
+template <int Threads>
+struct MhcPostCombineNormPrefill {
+  static void
+  run(tvm::ffi::TensorView x,
+      tvm::ffi::TensorView residual,
+      tvm::ffi::TensorView post,
+      tvm::ffi::TensorView comb,
+      tvm::ffi::TensorView pre,
+      tvm::ffi::TensorView weight,
+      tvm::ffi::TensorView residual_out,
+      tvm::ffi::TensorView output,
+      float eps) {
+    using namespace host;
+    auto m = SymbolicSize{"num_tokens"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLCUDA>();
+    TensorMatcher({m, 5120}).with_dtype<bf16_t>().with_device(dev).verify(x).verify(output);
+    TensorMatcher({m, 4, 5120}).with_dtype<bf16_t>().with_device(dev).verify(residual).verify(residual_out);
+    TensorMatcher({m, 4}).with_dtype<float>().with_device(dev).verify(post).verify(pre);
+    TensorMatcher({m, 4, 4}).with_dtype<float>().with_device(dev).verify(comb);
+    TensorMatcher({5120}).with_dtype<bf16_t>().with_device(dev).verify(weight);
+    CHECK_HOST(m.unwrap() >= 4096 && m.unwrap() <= 65536) << "prefill rows must be in [4096, 65536]";
+    for (const auto* ptr :
+         {x.data_ptr(), residual.data_ptr(), weight.data_ptr(), residual_out.data_ptr(), output.data_ptr()}) {
+      CHECK_HOST(reinterpret_cast<uintptr_t>(ptr) % 16 == 0) << "BF16 pointers must be 16-byte aligned";
+    }
+    CHECK_HOST(
+        residual_out.data_ptr() != residual.data_ptr() && residual_out.data_ptr() != x.data_ptr() &&
+        output.data_ptr() != residual.data_ptr() && output.data_ptr() != x.data_ptr() &&
+        output.data_ptr() != residual_out.data_ptr())
+        << "outputs must not alias inputs or each other";
+    const MhcPostCombineNormPrefillParams params{
+        static_cast<const bf16_t*>(x.data_ptr()),
+        static_cast<const bf16_t*>(residual.data_ptr()),
+        static_cast<const bf16_t*>(weight.data_ptr()),
+        static_cast<bf16_t*>(residual_out.data_ptr()),
+        static_cast<bf16_t*>(output.data_ptr()),
+        static_cast<const float*>(post.data_ptr()),
+        static_cast<const float*>(comb.data_ptr()),
+        static_cast<const float*>(pre.data_ptr()),
+        eps};
+    LaunchKernel(dim3(m.unwrap()), Threads, dev.unwrap()).launch(mhc_post_combine_norm_prefill_kernel<Threads>, params);
+  }
+};
+}  // namespace sglang

python/sglang/kernels/jit/csrc/distributed/all_reduce_fusion.cuh

@@ -286,7 +286,8 @@ template <
     bool kNorm,
     typename WeightT,
     bool kMhc = false,
-    bool kQuant = false>
+    bool kQuant = false,
+    bool kCollapse = false>
 __global__ __launch_bounds__(RowClusterTrait<kHiddenDim, kClusterSize>::kBlockSize)
     __cluster_dims__(1, kClusterSize, 1) void moe_finalize_all_reduce_kernel(
         const __grid_constant__ MoeFinalizeAllReduceParams<kWorldSize, WeightT> params) {
@@ -377,7 +378,14 @@ __global__ __launch_bounds__(RowClusterTrait<kHiddenDim, kClusterSize>::kBlockSi
   if constexpr (!kNorm) {
     const auto red = reduce_vec(vec);
     ptx::st_global_16B(red, params.out, vid);
-    if constexpr (kMhc) mhc_post_vec<kHiddenDim>(params, red, row_idx, hvec);
+    if constexpr (kMhc) {
+      if constexpr (kCollapse) {
+        const auto combined = mhc_post_vec<kHiddenDim, true>(params, red, row_idx, hvec);
+        ptx::st_global_16B(combined, params.normalized, vid);
+      } else {
+        mhc_post_vec<kHiddenDim>(params, red, row_idx, hvec);
+      }
+    }
     // ensure epoch is consumed, so flipping it won't lead to error
     barrier_cluster_wait();
   } else {
@@ -447,7 +455,8 @@ template <
     bool kUsePDL,
     typename WeightT,
     bool kMhc = false,
-    bool kQuant = false>
+    bool kQuant = false,
+    bool kCollapse = false>
 struct MoeFinalizeAllReduceKernel {
  private:
   static_assert(std::is_same_v<WeightT, bf16_t> || std::is_same_v<WeightT, fp32_t>);
@@ -466,7 +475,8 @@ struct MoeFinalizeAllReduceKernel {
       kNorm,
       WeightT,
       kMhc,
-      kQuant>;
+      kQuant,
+      kCollapse>;
 
  public:
   /// out = [allreduce over ranks of] finalize(gemm2_out, idx, weights) [+ shared] [-> RMSNorm(norm_weight, eps)].
@@ -528,6 +538,40 @@ struct MoeFinalizeAllReduceKernel {
         comb);
   }
 
+  // Keep the original RMSNorm as a separate kernel, while reusing the
+  // BF16-rounded post values for the next sublayer's pre-combine.
+  static void run_mhc_combine(
+      CommunicatorRef ref,
+      TensorView out,
+      TensorView gemm2_out,
+      TensorView permuted_idx,
+      TensorView expert_weights,
+      std::optional<TensorView> shared_output,
+      TensorView mhc_out,
+      TensorView residual,
+      TensorView post,
+      TensorView comb,
+      TensorView pre,
+      TensorView combined) {
+    static_assert(kMhc && kCollapse && !kQuant);
+    run_impl(
+        ref,
+        out,
+        gemm2_out,
+        permuted_idx,
+        expert_weights,
+        shared_output,
+        std::nullopt,
+        0.0,
+        false,
+        mhc_out,
+        residual,
+        post,
+        comb,
+        pre,
+        combined);
+  }
+
   static void run_mhc_norm(
       CommunicatorRef ref,
       TensorView out,
@@ -673,7 +717,7 @@ struct MoeFinalizeAllReduceKernel {
     }
     if constexpr (kMhc) {
       static_assert(kHiddenDim == 5120);
-      if (norm_weight.has_value()) {
+      if (norm_weight.has_value() || kCollapse) {
         TensorMatcher({T, 4}).with_dtype<fp32_t>().with_device<kDLCUDA>(device).verify(pre.value());
         TensorMatcher({T, kHiddenDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(normalized.value());
       }

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

@@ -6,6 +6,7 @@ import triton.language as tl
 from triton.language.extra import libdevice
 
 from sglang.kernels.ops.attention.dsv4.torch_quant import FP4_AMAX_FLOOR
+from sglang.srt.runtime_context import get_platform
 
 INDEX_HEAD_DIM = 128
 # One index-K slot: 64 packed e2m1 bytes and four ue8m0 block exponents.
@@ -132,6 +133,56 @@ def _quantize_fp4_indexer_kernel(
     tl.store(x_fp4 + token_id * (BLOCK_N // 2) + pair_offsets, packed)
 
 
+@triton.jit
+def _quantize_fp4_indexer_rows(
+    x,
+    x_fp4,
+    x_sf,
+    M,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    GROUP_N: tl.constexpr,
+    RNE: tl.constexpr,
+):
+    tl.static_assert(BLOCK_N == 128 and GROUP_N == 32)
+    # Each reduction covers one scale group. Keep its values for packing,
+    # avoiding four masked full-row reductions and a second input load.
+    group = tl.program_id(0) * BLOCK_M * 4 + tl.arange(0, BLOCK_M * 4)
+    offs = tl.arange(0, GROUP_N)
+    values = tl.load(
+        x + group[:, None].to(tl.int64) * GROUP_N + offs[None, :],
+        group[:, None] < M * 4,
+        0,
+    ).to(tl.float32)
+    amax = tl.max(tl.abs(values), axis=1)
+    exp = _ceil_ue8m0_exp(tl.maximum(amax / 6.0, 1.0e-4))
+    scale = (exp << 23).to(tl.float32, bitcast=True)
+    v0, v1 = tl.split(
+        tl.reshape(values / scale[:, None], (BLOCK_M * 4, GROUP_N // 2, 2))
+    )
+    if RNE:
+        code0 = _fp4_e2m1_code_rne(v0)
+        code1 = _fp4_e2m1_code_rne(v1)
+    else:
+        code0 = _fp4_e2m1_code(v0)
+        code1 = _fp4_e2m1_code(v1)
+    packed = (code0 & 0x0F) | ((code1 & 0x0F) << 4)
+    tl.store(
+        x_fp4
+        + group[:, None].to(tl.int64) * (GROUP_N // 2)
+        + tl.arange(0, GROUP_N // 2)[None, :],
+        packed,
+        group[:, None] < M * 4,
+    )
+    # The four exponents occupy disjoint bytes, so integer sum packs them.
+    shifts = tl.arange(0, 4) * 8
+    packed_sf = tl.sum(
+        tl.reshape(exp.to(tl.uint32), (BLOCK_M, 4)) << shifts[None, :], axis=1
+    )
+    token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    tl.store(x_sf + token_id, packed_sf.to(tl.int32), token_id < M)
+
+
 @triton.jit
 def _store_fp4_index_k_cache_kernel(
     k_fp4,
@@ -169,7 +220,20 @@ def quantize_fp4_indexer_tensor(
     x = x.contiguous().view(-1, x.shape[-1])
     x_fp4 = torch.empty((x.shape[0], 64), device=x.device, dtype=torch.int8)
     x_sf = torch.empty((x.shape[0],), device=x.device, dtype=torch.int32)
-    if x.shape[0] > 0:
+    if x.shape[0] >= 4096 and get_platform().is_blackwell:
+        # Independent rows share a CTA to avoid one block per 128 values.
+        _quantize_fp4_indexer_rows[(triton.cdiv(x.shape[0], 8),)](
+            x,
+            x_fp4,
+            x_sf,
+            x.shape[0],
+            8,
+            BLOCK_N=128,
+            GROUP_N=32,
+            RNE=rne,
+            num_warps=4,
+        )
+    elif x.shape[0] > 0:
         _quantize_fp4_indexer_kernel[(x.shape[0],)](
             x,
             x_fp4,

python/sglang/kernels/ops/communication/all_reduce_mhc_combine.py

@@ -0,0 +1,121 @@
+"""Finalize/all-reduce + post + combine, retaining standalone RMSNorm."""
+
+from typing import Optional
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+from sglang.kernels.ops.communication.all_reduce_fusion import (
+    default_cluster_size,
+    get_registered_comm,
+)
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+@cache_once
+def _module(world_size, top_k, cluster_size, weight_dtype):
+    args = make_cpp_args(
+        world_size,
+        5120,
+        top_k,
+        cluster_size,
+        is_arch_support_pdl(),
+        weight_dtype,
+        True,
+        False,
+        True,
+    )
+    return load_jit(
+        "moe_finalize_all_reduce_mhc_combine",
+        *args,
+        cuda_files=["distributed/all_reduce_fusion.cuh"],
+        cuda_wrappers=[("run", f"MoeFinalizeAllReduceKernel<{args}>::run_mhc_combine")],
+    )
+
+
+@register_custom_op(mutates_args=["out", "mhc_out", "combined"])
+def _moe_finalize_all_reduce_mhc_combine(
+    world_size: int,
+    top_k: int,
+    cluster_size: int,
+    out: torch.Tensor,
+    mhc_out: torch.Tensor,
+    combined: torch.Tensor,
+    gemm2: torch.Tensor,
+    idx: torch.Tensor,
+    weights: torch.Tensor,
+    shared: Optional[torch.Tensor],
+    residual: torch.Tensor,
+    post: torch.Tensor,
+    comb: torch.Tensor,
+    pre: torch.Tensor,
+) -> None:
+    comm = get_registered_comm(world_size)
+    assert comm is not None
+    _module(world_size, top_k, cluster_size, weights.dtype).run(
+        comm,
+        out,
+        gemm2,
+        idx,
+        weights,
+        shared,
+        mhc_out,
+        residual,
+        post,
+        comb,
+        pre,
+        combined,
+    )
+
+
+def moe_finalize_all_reduce_mhc_combine(
+    gemm2,
+    idx,
+    weights,
+    top_k,
+    shared,
+    residual,
+    post,
+    comb,
+    pre,
+    *,
+    world_size,
+    cluster_size=None,
+):
+    out = torch.empty(
+        (weights.shape[0], 5120), dtype=torch.bfloat16, device=gemm2.device
+    )
+    mhc_out = torch.empty_like(residual)
+    combined = torch.empty_like(out)
+    if weights.shape[0]:
+        _moe_finalize_all_reduce_mhc_combine(
+            world_size,
+            top_k,
+            cluster_size or default_cluster_size(5120),
+            out,
+            mhc_out,
+            combined,
+            gemm2,
+            idx,
+            weights,
+            shared,
+            residual,
+            post,
+            comb,
+            pre,
+        )
+    return out, mhc_out, combined
+
+
+def all_reduce_mhc_combine(x, residual, post, comb, pre, *, world_size):
+    from sglang.kernels.ops.communication.all_reduce_mhc import _identity_routing
+
+    idx, weights = _identity_routing(x.shape[0], x.device)
+    return moe_finalize_all_reduce_mhc_combine(
+        x, idx, weights, 1, None, residual, post, comb, pre, world_size=world_size
+    )

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

@@ -0,0 +1,85 @@
+"""HC=4 post-mix and pre-combine with the original BF16 intermediates."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _mhc_post_combine(X, R, P, C, A, RO, Y, H: tl.constexpr, B: tl.constexpr):
+    row = tl.program_id(0)
+    h = tl.program_id(1) * B + tl.arange(0, B)
+    mask = h < H
+    x = tl.load(X + row * H + h, mask, 0).to(tl.float32)
+    r0 = tl.load(R + (row * 4 + 0) * H + h, mask, 0).to(tl.float32)
+    r1 = tl.load(R + (row * 4 + 1) * H + h, mask, 0).to(tl.float32)
+    r2 = tl.load(R + (row * 4 + 2) * H + h, mask, 0).to(tl.float32)
+    r3 = tl.load(R + (row * 4 + 3) * H + h, mask, 0).to(tl.float32)
+    collapsed = tl.full((B,), 0, tl.float32)
+    for i in tl.static_range(4):
+        post = tl.load(P + row * 4 + i)
+        c0 = tl.load(C + row * 16 + i)
+        c1 = tl.load(C + row * 16 + 4 + i)
+        c2 = tl.load(C + row * 16 + 8 + i)
+        c3 = tl.load(C + row * 16 + 12 + i)
+        pre = tl.load(A + row * 4 + i)
+        # Match mhc_post_split_h's contraction order and BF16 store.
+        mixed = tl.fma(post, x, c0 * r0)
+        mixed = tl.fma(c1, r1, mixed)
+        mixed = tl.fma(c2, r2, mixed)
+        mixed = tl.fma(c3, r3, mixed)
+        rounded = mixed.to(tl.bfloat16)
+        tl.store(RO + (row * 4 + i) * H + h, rounded, mask)
+        # Match hc_combine's sequential FP32 accumulation, then BF16 store.
+        collapsed = tl.fma(pre, rounded.to(tl.float32), collapsed)
+    tl.store(Y + row * H + h, collapsed, mask)
+
+
+def mhc_post_combine(x, residual, post, comb, pre):
+    """Return (updated HC streams, collapsed BF16 input) out of place."""
+    assert x.dtype == residual.dtype == torch.bfloat16
+    assert post.dtype == comb.dtype == pre.dtype == torch.float32
+    assert x.ndim == 2 and x.shape[1] == 5120
+    assert residual.shape == (x.shape[0], 4, x.shape[1])
+    assert post.shape == pre.shape == (x.shape[0], 4)
+    assert comb.shape == (x.shape[0], 4, 4)
+    assert all(t.is_contiguous() for t in (x, residual, post, comb, pre))
+    updated = torch.empty_like(residual)
+    combined = torch.empty_like(x)
+    if x.shape[0]:
+        block = 1024 if x.shape[0] <= 192 or x.shape[0] >= 4096 else 512
+        _mhc_post_combine[(x.shape[0], triton.cdiv(x.shape[1], block))](
+            x,
+            residual,
+            post,
+            comb,
+            pre,
+            updated,
+            combined,
+            H=x.shape[1],
+            B=block,
+            num_warps=4,
+            enable_fp_fusion=False,
+        )
+    return updated, combined
+
+
+@triton.jit
+def _hc_norm_prefill(X, W, Y, EPS: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    h = tl.arange(0, 8192)
+    value = tl.load(X + row * 5120 + h, h < 5120, 0).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_norm_prefill(combined, weight, eps):
+    """Keep the original fused prefill norm's arithmetic on collapsed input."""
+    assert 4096 <= combined.shape[0] <= 65536 and combined.shape[1] == 5120
+    assert weight.shape == (5120,)
+    assert combined.dtype == weight.dtype == torch.bfloat16
+    assert combined.is_contiguous() and weight.is_contiguous()
+    output = torch.empty_like(combined)
+    _hc_norm_prefill[(combined.shape[0],)](combined, weight, output, eps, num_warps=4)
+    return output

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

@@ -0,0 +1,38 @@
+"""Prefill post/combine/RMSNorm with the original BF16 boundaries and norm tree."""
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit
+from sglang.srt.utils.custom_op import register_custom_op
+
+
+@cache_once
+def _module():
+    return load_jit(
+        "mhc_post_combine_norm_prefill",
+        cuda_files=["deepseek_v4/mhc_post_combine_norm_prefill.cuh"],
+        cuda_wrappers=[("run", "MhcPostCombineNormPrefill<128>::run")],
+    )
+
+
+@register_custom_op(mutates_args=["updated", "normalized"])
+def _mhc_post_combine_norm_prefill(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post: torch.Tensor,
+    comb: torch.Tensor,
+    pre: torch.Tensor,
+    weight: torch.Tensor,
+    updated: torch.Tensor,
+    normalized: torch.Tensor,
+    eps: float,
+) -> None:
+    _module().run(x, residual, post, comb, pre, weight, updated, normalized, eps)
+
+
+def mhc_post_combine_norm_prefill(x, residual, post, comb, pre, weight, eps):
+    updated, normalized = torch.empty_like(residual), torch.empty_like(x)
+    _mhc_post_combine_norm_prefill(
+        x, residual, post, comb, pre, weight, updated, normalized, eps
+    )
+    return updated, normalized

python/sglang/kernels/ops/moe/moe_fused_gate.py

@@ -8,7 +8,12 @@ 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.jit.utils import (
+    cache_once,
+    get_jit_cuda_arch,
+    is_arch_support_pdl,
+    load_jit,
+)
 from sglang.kernels.kernel_api_logging import debug_kernel_api
 from sglang.kernels.ops.moe import moe_route_radix
 
@@ -461,6 +466,12 @@ def moe_fused_gate(
     num_warps = 1 if BLOCK_N <= 512 else 4
     grid = (triton.cdiv(M, BLOCK_M),)
     use_pdl = is_arch_support_pdl()
+    if use_pdl and scoring_func_int == 1 and N == 384 and K == 6 and M <= 8:
+        # On SM103, early-launching the small DSV4.1 target router increases
+        # latency when it overlaps with mHC/shared-expert work. Use ordinary
+        # stream dependencies; keep PDL for the draft router and larger batches.
+        arch = get_jit_cuda_arch()
+        use_pdl = (arch.major, arch.minor) != (10, 3)
     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.

python/sglang/kernels/ops/speculative/dspark/dspark_accept.py

@@ -676,7 +676,10 @@ def _row_argmax(logits: torch.Tensor, fused: bool = False) -> torch.Tensor:
         and logits.dim() == 2
         and logits.dtype == torch.float32
         and logits.stride(1) == 1
-        and logits.shape[0] <= 64
+        and (
+            logits.shape[0] <= 64
+            or (logits.shape[0] <= 384 and logits.shape[1] >= 65536)
+        )
         and logits.shape[1] >= 4096
     ):
         from sglang.kernels.ops.speculative.row_argmax import row_argmax

python/sglang/kernels/ops/speculative/row_argmax.py

@@ -51,11 +51,63 @@ def _argmax_final_kernel(INV, INI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexp
 _SPLITS = 64
 
 
+@triton.jit
+def _argmax_pair(av, ai, bv, bi):
+    # Torch gives NaNs priority and picks the first index for all ties.
+    an, bn = av != av, bv != bv
+    take_a = (an & ~bn) | ((an == bn) & ((av > bv) | (((av == bv) | an) & (ai < bi))))
+    return tl.where(take_a, av, bv), tl.where(take_a, ai, bi)
+
+
+@triton.jit
+def _medium_argmax_partial_kernel(
+    X,
+    PV,
+    PI,
+    N: tl.constexpr,
+    SX: tl.constexpr,
+    SPLITS: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row, part = tl.program_id(0), tl.program_id(1)
+    ix = part * BLOCK + tl.arange(0, BLOCK)
+    valid = ix < N
+    v = tl.load(X + row * SX + ix, valid, float("-inf"))
+    i = tl.where(valid, ix, N)
+    best_v, best_i = tl.reduce((v, i), 0, _argmax_pair)
+    tl.store(PV + row * SPLITS + part, best_v)
+    tl.store(PI + row * SPLITS + part, best_i)
+
+
+@triton.jit
+def _medium_argmax_final_kernel(PV, PI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    part = tl.arange(0, BLOCK)
+    v = tl.load(PV + row * SPLITS + part, part < SPLITS, float("-inf"))
+    i = tl.load(PI + row * SPLITS + part, part < SPLITS, 0x7FFFFFFF)
+    _, index = tl.reduce((v, i), 0, _argmax_pair)
+    tl.store(OUT + row, index.to(tl.int64))
+
+
 def row_argmax(x: torch.Tensor) -> torch.Tensor:
-    """``x.argmax(dim=-1)`` for a 2D FP32 tensor with few rows and a wide vocab."""
+    """``x.argmax(dim=-1)`` for FP32 speculative logits with a wide vocab."""
     assert x.dim() == 2 and x.dtype == torch.float32 and x.stride(1) == 1
     rows, n = x.shape
     out = torch.empty((rows,), dtype=torch.int64, device=x.device)
+    if rows > 64:
+        # Whole aligned tiles avoid the loop and unaligned partition starts.
+        # Larger batches use wider tiles to limit the total number of CTAs.
+        block = 4096 if rows <= 256 else 8192
+        splits = triton.cdiv(n, block)
+        pv = torch.empty((rows, splits), dtype=torch.float32, device=x.device)
+        pi = torch.empty((rows, splits), dtype=torch.int32, device=x.device)
+        _medium_argmax_partial_kernel[(rows, splits)](
+            x, pv, pi, n, x.stride(0), splits, block, num_warps=4
+        )
+        _medium_argmax_final_kernel[(rows,)](
+            pv, pi, out, splits, triton.next_power_of_2(splits), num_warps=1
+        )
+        return out
     pv = torch.empty((rows, _SPLITS), dtype=torch.float32, device=x.device)
     pi = torch.empty((rows, _SPLITS), dtype=torch.int32, device=x.device)
     _argmax_partial_kernel[(rows, _SPLITS)](

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -4294,7 +4294,7 @@ class DeepseekV4AttnBackend(
         small_metadata = (
             not is_prefill
             and seq_lens_casual.is_cuda
-            and 0 < seq_lens_casual.numel() <= 8
+            and 0 < seq_lens_casual.numel() <= 384
             and out_loc.numel() == seq_lens_casual.numel()
             and self.low_ratios == (1, 2)
             and set(self.present_ratios) == {1, 2}

python/sglang/srt/layers/moe/mhc_post_fusion.py

@@ -20,10 +20,19 @@ class MhcPostFusion:
     norm_eps: float = 0.0
     normalized: Optional[torch.Tensor] = None
     quantized: Optional[tuple[torch.Tensor, torch.Tensor]] = None
+    combine_only: bool = False
+    combined: Optional[torch.Tensor] = None
+    overlap_only: bool = False
     record_stats: Optional[
         Callable[[], tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
     ] = None
 
+    def start_stats_before_all_reduce(self):
+        if self.overlap_only and self.record_stats is not None:
+            assert self.stats_stream is not None
+            self.stats_stream.wait_stream(torch.cuda.current_stream())
+            self.materialize_stats()
+
     def materialize_stats(self):
         # Record after the main parent; graph replay must keep the join on the
         # main stream.

python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py

@@ -416,7 +416,10 @@ class Mxfp4FlashinferTrtllmMoEMethod:
         # triple instead of the finalized [T, hidden] tensor.
         defer_finalize = is_deferred_finalize_enabled()
         symm_output = None
-        if not defer_finalize:
+        # Preserve the ordinary output shape in the medium-batch autotuner
+        # cache key. The deferred ABI ignores this allocation and returns the
+        # expanded GEMM output for the separate fused finalize epilogue.
+        if not defer_finalize or 96 < num_tokens <= 384:
             with use_symmetric_memory(
                 get_parallel().tp_group, disabled=not is_allocation_symmetric()
             ):
@@ -519,11 +522,17 @@ def maybe_fuse_routed_scale_and_shared_add(
 # Fused finalize + shared add + TP all-reduce
 _fused_finalize_all_reduce_world_size: Optional[int] = None
 _fused_finalize_all_reduce_probed = False
+_fused_finalize_all_reduce_comm = None
 
 
 def _fused_finalize_all_reduce_comm_world_size() -> Optional[int]:
+    """Reserve a separate push plane for up to 384 rows of fused MoE output."""
     global _fused_finalize_all_reduce_world_size, _fused_finalize_all_reduce_probed
+    global _fused_finalize_all_reduce_comm
     if not _fused_finalize_all_reduce_probed:
+        # The eager warmup must initialize peer workspaces before capture.
+        if torch.cuda.is_current_stream_capturing():
+            return None
         _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 (
@@ -532,8 +541,24 @@ def _fused_finalize_all_reduce_comm_world_size() -> Optional[int]:
 
         ca_comm = get_parallel().tp_group.ca_comm
         if isinstance(ca_comm, CustomAllReduceV2) and not ca_comm.disabled:
-            all_reduce_fusion.register_comm(ca_comm.obj)
+            fused_comm = ca_comm
-            _fused_finalize_all_reduce_world_size = ca_comm.world_size
+            if ca_comm.world_size == 4:
+                from sglang.kernels.ops.communication.mp import register_comm_cleanup
+
+                fused_comm = CustomAllReduceV2(
+                    ca_comm.group,
+                    ca_comm.device,
+                    max_pull_size=0,
+                    max_pull_blocks=0,
+                    max_push_size=4 * 1024 * 1024,
+                    max_push_blocks=512,
+                )
+                register_comm_cleanup(fused_comm)
+            if fused_comm.disabled:
+                return None
+            _fused_finalize_all_reduce_comm = fused_comm
+            all_reduce_fusion.register_comm(fused_comm.obj)
+            _fused_finalize_all_reduce_world_size = fused_comm.world_size
         else:
             log_info_on_rank0(
                 logger,
@@ -556,6 +581,11 @@ def should_use_fuse_finalize_all_reduce(
         return False
     if num_tokens <= 0:
         return False
+    if num_tokens > 96:
+        from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
+
+        if is_batch_invariant_mode_enabled():
+            return False
     from sglang.kernels.ops.communication import all_reduce_fusion
 
     if not all_reduce_fusion.valid_cluster_sizes(hidden_dim):
@@ -563,9 +593,8 @@ def should_use_fuse_finalize_all_reduce(
     tp_group = get_parallel().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)
+    comm = _fused_finalize_all_reduce_comm
-    if num_tokens > tp_group.ca_comm.config.num_push_blocks:
+    # Each token row owns a push phase counter on the epilogue's plane.
+    if num_tokens > comm.config.num_push_blocks:
         return False
-    return all_reduce_fusion.fits_push_slot(
+    return all_reduce_fusion.fits_push_slot(comm.max_push_size, num_tokens, hidden_dim)
-        tp_group.ca_comm.max_push_size, num_tokens, hidden_dim
-    )

python/sglang/srt/models/deepseek_v2.py

@@ -552,9 +552,9 @@ class MoEGate(nn.Module):
         return logits
 
 
-# 96 rows of 5120 bf16 fit the 1 MiB CustomAllReduceV2 push slot the whole
+# The dedicated 4 MiB push slot fits 384 rows of 5120 BF16 values.
-# [T, hidden] view is staged through.
+# The dispatch gate also checks slot capacity and available counters.
-_FUSED_FINALIZE_ALL_REDUCE_MAX_TOKENS = 96
+_FUSED_FINALIZE_ALL_REDUCE_MAX_TOKENS = 384
 
 
 class DeepseekV2MoE(nn.Module):
@@ -1129,7 +1129,17 @@ class DeepseekV2MoE(nn.Module):
                         mhc.post,
                         mhc.comb,
                     )
-                    if mhc.norm_weight is not None:
+                    if mhc.combine_only:
+                        from sglang.kernels.ops.communication.all_reduce_mhc_combine import (
+                            moe_finalize_all_reduce_mhc_combine,
+                        )
+
+                        final_hidden_states, mhc.output, mhc.combined = (
+                            moe_finalize_all_reduce_mhc_combine(
+                                *args, mhc.pre, world_size=self.tp_size
+                            )
+                        )
+                    elif mhc.norm_weight is not None:
                         from sglang.kernels.ops.communication.all_reduce_mhc import (
                             moe_finalize_all_reduce_mhc_quant,
                         )
@@ -1175,6 +1185,16 @@ class DeepseekV2MoE(nn.Module):
             )
 
         if not all_reduce_done:
+            if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
+                is_tp_path=True
+            ):
+                from sglang.srt.layers.moe.mhc_post_fusion import (
+                    current_mhc_post_fusion,
+                )
+
+                mhc = current_mhc_post_fusion()
+                if mhc is not None:
+                    mhc.start_stats_before_all_reduce()
             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.
@@ -1321,6 +1341,14 @@ class DeepseekV2MoE(nn.Module):
             self.routed_scaling_factor,
         )
 
+        if self.tp_size > 1 and not should_skip_post_experts_all_reduce(
+            is_tp_path=True,
+        ):
+            from sglang.srt.layers.moe.mhc_post_fusion import current_mhc_post_fusion
+
+            mhc = current_mhc_post_fusion()
+            if mhc is not None:
+                mhc.start_stats_before_all_reduce()
         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.

python/sglang/srt/models/deepseek_v4.py

@@ -2525,7 +2525,10 @@ class MQALayer(MqaAttentionBase):
             o if isinstance(o, Mxfp8SwizzledInput) else o.flatten(1),
             skip_all_reduce=mhc is not None,
         )
-        if mhc is not None:
+        if mhc is not None and mhc.overlap_only:
+            mhc.start_stats_before_all_reduce()
+            o = attn_tp_all_reduce(o)
+        elif mhc is not None:
             from sglang.kernels.ops.communication.all_reduce_mhc import (
                 all_reduce_mhc_norm,
             )
@@ -2533,6 +2536,20 @@ class MQALayer(MqaAttentionBase):
             mhc.materialize_stats()
             if mhc.stats_stream is not None:
                 torch.cuda.current_stream().wait_stream(mhc.stats_stream)
+            if mhc.combine_only:
+                from sglang.kernels.ops.communication.all_reduce_mhc_combine import (
+                    all_reduce_mhc_combine,
+                )
+
+                o, mhc.output, mhc.combined = all_reduce_mhc_combine(
+                    o,
+                    mhc.residual,
+                    mhc.post,
+                    mhc.comb,
+                    mhc.pre,
+                    world_size=self.attn_tp_size,
+                )
+            else:
                 o, mhc.output, mhc.normalized = all_reduce_mhc_norm(
                     o,
                     mhc.residual,
@@ -3180,6 +3197,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         quantized: Optional[list] = None,
         normalized: Optional[torch.Tensor] = None,
         precomputed: Optional[tuple] = None,
+        combined: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         from sglang.kernels.ops.layernorm.mhc import hc_combine
 
@@ -3195,8 +3213,23 @@ class DeepseekV4DecoderLayer(nn.Module):
                 quantized.append(precomputed[1])
                 return precomputed[0]
             if normalized is not None:
-                assert not quantize
+                # Prefill projections still quantize the BF16 input themselves;
+                # the optional fused-quantization list stays empty for this case.
+                assert not quantize or 4096 <= x.shape[0] <= 65536
                 return normalized
+            if combined is not None:
+                if (
+                    4096 <= combined.shape[0] <= 65536
+                    and norm.weight.dtype == torch.bfloat16
+                    and not norm.cast_x_before_out_mul
+                    and norm.variance_size_override is None
+                ):
+                    from sglang.kernels.ops.layernorm.mhc_post_combine import (
+                        hc_norm_prefill,
+                    )
+
+                    return hc_norm_prefill(combined, norm.weight, norm.variance_epsilon)
+                return norm(combined)
             if apply_pre is None:
                 return norm(x[:, 0, :].contiguous())
             from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
@@ -3378,7 +3411,24 @@ class DeepseekV4DecoderLayer(nn.Module):
         )
 
     def _get_hc_stats_stream(self, hidden_states, forward_batch):
-        # Every branch joins this stream before hc_post reads the coefficients.
+        # Prefill stats share one model-wide stream. Start them immediately
+        # before the sublayer's all-reduce, after its compute has completed.
+        if (
+            self.config.model_type == "deepseek_v41"
+            and hidden_states.is_cuda
+            and get_platform().is_blackwell
+            and forward_batch.forward_mode.is_extend_without_speculative()
+            and 4096 <= hidden_states.shape[0] <= 65536
+            and get_parallel().attn_dp_size == 1
+            and not get_forward().sp_active
+            and not self.dsa_enable_prefill_cp
+        ):
+            from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
+
+            if not is_batch_invariant_mode_enabled():
+                return self.hc_stats_stream
+        # Verify batches can also compute coefficients beside the
+        # sublayer; each branch joins before hc_post reads those coefficients.
         return (
             self.hc_stats_stream
             if (
@@ -3392,6 +3442,65 @@ class DeepseekV4DecoderLayer(nn.Module):
             else None
         )
 
+    def _hc_post_with_combine(
+        self, x, residual, post, comb, pre, forward_batch, norm=None
+    ):
+        """Return updated HC streams and optional combined/normalized inputs."""
+        from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
+
+        if (
+            self.config.model_type == "deepseek_v41"
+            and x.is_cuda
+            and get_platform().is_blackwell
+            and (
+                (
+                    128 <= x.shape[0] <= 384
+                    and (
+                        forward_batch.forward_mode.is_decode()
+                        or forward_batch.forward_mode.is_target_verify()
+                    )
+                )
+                or (
+                    4096 <= x.shape[0] <= 65536
+                    and forward_batch.forward_mode.is_extend_without_speculative()
+                    and envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get()
+                    and not envs.SGLANG_OPT_USE_FLASHINFER_MHC.get()
+                )
+            )
+            and x.shape[1] == 5120
+            and self.hc_mult == 4
+            and x.dtype == residual.dtype == torch.bfloat16
+            and post.dtype == comb.dtype == pre.dtype == torch.float32
+            and all(t.is_contiguous() for t in (x, residual, post, comb, pre))
+            and get_parallel().attn_dp_size == 1
+            and not get_forward().sp_active
+            and not self.dsa_enable_prefill_cp
+            and not is_batch_invariant_mode_enabled()
+        ):
+            if (
+                x.shape[0] >= 4096
+                and norm is not None
+                and not norm.cast_x_before_out_mul
+                and norm.variance_size_override is None
+                and norm.weight.dtype == torch.bfloat16
+                and norm.weight.shape == (5120,)
+                and norm.weight.is_contiguous()
+                and all(t.data_ptr() % 16 == 0 for t in (x, residual, norm.weight))
+            ):
+                from sglang.kernels.ops.layernorm.mhc_post_combine_norm_prefill import (
+                    mhc_post_combine_norm_prefill,
+                )
+
+                updated, normalized = mhc_post_combine_norm_prefill(
+                    x, residual, post, comb, pre, norm.weight, norm.variance_epsilon
+                )
+                return updated, None, normalized
+            from sglang.kernels.ops.layernorm.mhc_post_combine import mhc_post_combine
+
+            updated, combined = mhc_post_combine(x, residual, post, comb, pre)
+            return updated, combined, None
+        return self.hc_post(x, residual, post, comb), None, None
+
     def forward_hc_pre_from_prev(
         self,
         positions: torch.Tensor,
@@ -3403,11 +3512,20 @@ class DeepseekV4DecoderLayer(nn.Module):
         precomputed_attn: Optional[tuple] = None,
         next_norm: Optional[RMSNorm] = None,
         next_input: Optional[list] = None,
+        combined_attn: Optional[torch.Tensor] = None,
+        normalized_attn: Optional[torch.Tensor] = None,
+        next_combined: Optional[list] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Layer forward where each sublayer consumes the previous sublayer's
         pre-mix. Returns (hidden_states, ffn_pre)."""
         from functools import partial
 
+        from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
+        from sglang.srt.layers.moe.mhc_post_fusion import (
+            MhcPostFusion,
+            use_mhc_post_fusion,
+        )
+
         stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch)
         residual = hidden_states
         attn_quantized: Optional[list] = (
@@ -3428,13 +3546,23 @@ class DeepseekV4DecoderLayer(nn.Module):
             stats_stream=stats_stream,
             quantized=attn_quantized,
             precomputed=precomputed_attn,
+            combined=combined_attn,
+            normalized=normalized_attn,
+        )
+        prefill_overlap = (
+            stats_stream is not None
+            and forward_batch.forward_mode.is_extend_without_speculative()
+        )
+        medium_verify = 128 <= x.shape[0] <= 384 and (
+            forward_batch.forward_mode.is_decode()
+            or forward_batch.forward_mode.is_target_verify()
         )
         attn_mhc = None
         if (
             self.config.model_type == "deepseek_v41"
             and x.is_cuda
             and get_platform().is_blackwell
-            and 0 < x.shape[0] <= 8
+            and (0 < x.shape[0] <= 8 or medium_verify)
             and x.shape[1] == 5120
             and self.hc_mult == 4
             and x.dtype == residual.dtype == torch.bfloat16
@@ -3451,16 +3579,18 @@ class DeepseekV4DecoderLayer(nn.Module):
             from sglang.kernels.ops.communication.all_reduce_fusion import (
                 get_registered_comm,
             )
-            from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled
+
-            from sglang.srt.layers.moe.mhc_post_fusion import (
+            comm_ready = get_registered_comm(self.self_attn.attn_tp_size) is not None
-                MhcPostFusion,
+            if medium_verify and not is_batch_invariant_mode_enabled():
-                use_mhc_post_fusion,
+                from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
+                    _fused_finalize_all_reduce_comm_world_size,
                 )
 
-            if (
+                comm_ready = (
-                not is_batch_invariant_mode_enabled()
+                    _fused_finalize_all_reduce_comm_world_size()
-                and get_registered_comm(self.self_attn.attn_tp_size) is not None
+                    == self.self_attn.attn_tp_size
-            ):
+                )
+            if not is_batch_invariant_mode_enabled() and comm_ready:
                 attn_mhc = MhcPostFusion(
                     residual,
                     None,
@@ -3469,6 +3599,20 @@ class DeepseekV4DecoderLayer(nn.Module):
                     record_stats=attn_stats,
                     norm_weight=self.post_attention_layernorm.weight,
                     norm_eps=self.post_attention_layernorm.variance_epsilon,
+                    combine_only=medium_verify,
+                )
+        if (
+            prefill_overlap
+            and get_parallel().tp_size == self.self_attn.attn_tp_size == 4
+            and self.self_attn.wo_b.reduce_results
+        ):
+            attn_mhc = MhcPostFusion(
+                residual,
+                None,
+                None,
+                stats_stream,
+                overlap_only=True,
+                record_stats=attn_stats,
             )
         context = (
             use_mhc_post_fusion(attn_mhc) if attn_mhc is not None else nullcontext()
@@ -3480,15 +3624,32 @@ class DeepseekV4DecoderLayer(nn.Module):
                 forward_batch=forward_batch,
                 x_quant=attn_quantized[0] if attn_quantized else None,
             )
+        ffn_combined = None
+        ffn_normalized = None
         if attn_mhc is not None:
             attn_mhc.materialize_stats()
+        if attn_mhc is not None and attn_mhc.output is not None:
             attn_pre = attn_mhc.pre
             hidden_states = attn_mhc.output
+            ffn_combined = attn_mhc.combined
+            ffn_normalized = attn_mhc.normalized
         else:
-            attn_pre, attn_post, attn_comb = attn_stats()
+            attn_pre, attn_post, attn_comb = (
+                (attn_mhc.pre, attn_mhc.post, attn_mhc.comb)
+                if attn_mhc is not None
+                else attn_stats()
+            )
             if stats_stream is not None:
                 torch.cuda.current_stream().wait_stream(stats_stream)
-            hidden_states = self.hc_post(x, residual, attn_post, attn_comb)
+            hidden_states, ffn_combined, ffn_normalized = self._hc_post_with_combine(
+                x,
+                residual,
+                attn_post,
+                attn_comb,
+                attn_pre,
+                forward_batch,
+                norm=self.post_attention_layernorm,
+            )
 
         residual = hidden_states
         ffn_stats = partial(
@@ -3504,38 +3665,52 @@ class DeepseekV4DecoderLayer(nn.Module):
             apply_pre=attn_pre,
             norm=self.post_attention_layernorm,
             stats_stream=stats_stream,
-            normalized=attn_mhc.normalized if attn_mhc is not None else None,
+            normalized=ffn_normalized,
+            combined=ffn_combined,
         )
         mhc = None
         if (
             self.config.model_type == "deepseek_v41"
             and x.is_cuda
             and get_platform().is_blackwell
-            and 0 < x.shape[0] <= 8
+            and (0 < x.shape[0] <= 8 or (medium_verify and next_combined is not None))
             and x.shape[1] == 5120
             and self.hc_mult == 4
             and x.dtype == residual.dtype == torch.bfloat16
             and residual.is_contiguous()
             and get_parallel().attn_dp_size == 1
             and get_moe_a2a_backend().is_none()
+            and not get_forward().sp_active
             and not self.dsa_enable_prefill_cp
             and not self.mlp._shared_expert_tp1
             and self.mlp.tp_size == 4
+            and (not medium_verify or not is_batch_invariant_mode_enabled())
         ):
-            from sglang.srt.layers.moe.mhc_post_fusion import (
-                MhcPostFusion,
-                use_mhc_post_fusion,
-            )
-
             mhc = MhcPostFusion(
-                residual, None, None, stats_stream, record_stats=ffn_stats
+                residual,
+                None,
+                None,
+                stats_stream,
+                record_stats=ffn_stats,
+                combine_only=medium_verify,
             )
             if next_norm is not None:
                 mhc.norm_weight = next_norm.weight
                 mhc.norm_eps = next_norm.variance_epsilon
-            context = use_mhc_post_fusion(mhc)
+        if (
-        else:
+            prefill_overlap
-            context = nullcontext()
+            and self.mlp.tp_size == 4
+            and get_moe_a2a_backend().is_none()
+        ):
+            mhc = MhcPostFusion(
+                residual,
+                None,
+                None,
+                stats_stream,
+                overlap_only=True,
+                record_stats=ffn_stats,
+            )
+        context = use_mhc_post_fusion(mhc) if mhc is not None else nullcontext()
         with context:
             x = self._run_moe_ffn_dp_sync(
                 x, forward_batch, input_ids=input_ids, input_ids_global=input_ids_global
@@ -3549,9 +3724,24 @@ class DeepseekV4DecoderLayer(nn.Module):
             hidden_states = mhc.output
             if next_input is not None and mhc.quantized is not None:
                 next_input.append((mhc.normalized, Mxfp8SwizzledInput(*mhc.quantized)))
+            if next_combined is not None and mhc.combined is not None:
+                next_combined.append((mhc.combined, None))
         else:
             if stats_stream is not None:
                 torch.cuda.current_stream().wait_stream(stats_stream)
+            if next_combined is not None:
+                hidden_states, combined, normalized = self._hc_post_with_combine(
+                    x,
+                    residual,
+                    ffn_post,
+                    ffn_comb,
+                    ffn_pre,
+                    forward_batch,
+                    norm=next_norm,
+                )
+                if combined is not None or normalized is not None:
+                    next_combined.append((combined, normalized))
+            else:
                 hidden_states = self.hc_post(x, residual, ffn_post, ffn_comb)
         return hidden_states, ffn_pre
 
@@ -4243,8 +4433,12 @@ class DeepseekV4Model(nn.Module):
         saved_full = None
         prev_pre = None
         precomputed_attn = None
+        combined_attn = None
+        normalized_attn = None
         for i in range(self.start_layer, self.end_layer):
             if tail is not None and i == self.late_layer_start:
+                combined_attn = None
+                normalized_attn = None
                 # Decode reaches back at most SWA_WINDOW positions.
                 saved_full = attn_backend.enter_late_layer_tail(forward_batch)
                 hidden_states, prev_pre, input_ids, input_ids_global = (
@@ -4259,6 +4453,8 @@ class DeepseekV4Model(nn.Module):
             engram = self.layers[i].engram
             if engram is not None:
                 precomputed_attn = None
+                combined_attn = None
+                normalized_attn = None
                 before_engram = hidden_states
                 hidden_states = engram(
                     hidden_states,
@@ -4288,6 +4484,27 @@ class DeepseekV4Model(nn.Module):
             )
             next_norm = None
             next_input = []
+            # The next layer can consume a collapsed input only if no Engram
+            # or row selection changes the residual between the two layers.
+            next_combined = (
+                []
+                if (
+                    self.config.model_type == "deepseek_v41"
+                    and (
+                        128 <= hidden_states.shape[0] <= 384
+                        or (
+                            4096 <= hidden_states.shape[0] <= 65536
+                            and forward_batch.forward_mode.is_extend_without_speculative()
+                        )
+                    )
+                    and i + 1 < self.end_layer
+                    and tail is None
+                    and self.layers[i + 1].engram is None
+                )
+                else None
+            )
+            if next_combined is not None and hidden_states.shape[0] >= 4096:
+                next_norm = self.layers[i + 1].input_layernorm
             if (
                 self.config.model_type == "deepseek_v41"
                 and i + 1 < self.end_layer
@@ -4328,8 +4545,14 @@ class DeepseekV4Model(nn.Module):
                     precomputed_attn=precomputed_attn,
                     next_norm=next_norm,
                     next_input=next_input,
+                    combined_attn=combined_attn,
+                    normalized_attn=normalized_attn,
+                    next_combined=next_combined,
                 )
             precomputed_attn = next_input[0] if next_input else None
+            combined_attn, normalized_attn = (
+                next_combined[0] if next_combined else (None, None)
+            )
         if saved_full is not None:
             attn_backend.exit_late_layer_tail(saved_full, forward_batch)
             return hidden_states, prev_pre, tail

test/registered/kernels/ops/attention/test_fp4_indexer.py

@@ -116,6 +116,89 @@ def test_quantize_fp4_indexer_tensor(num_tokens: int) -> None:
     torch.testing.assert_close(x_sf, ref_sf)
 
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="Vectorized prefill dispatch targets Blackwell",
+)
+@pytest.mark.parametrize("rows", [4096, 4097, 16384, 524288])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("rne", [False, True])
+@pytest.mark.parametrize("strided", [False, True])
+def test_prefill_quantization_matches_single_row_and_replay(rows, dtype, rne, strided):
+    from sglang.kernels.ops.attention.dsv4.fp4_indexer import (
+        _quantize_fp4_indexer_kernel,
+    )
+
+    x = torch.randn(rows, 256 if strided else 128, device="cuda", dtype=dtype)
+    if strided:
+        x = x[:, ::2]
+
+    def reference():
+        q = torch.empty(rows, 64, device="cuda", dtype=torch.int8)
+        sf = torch.empty(rows, device="cuda", dtype=torch.int32)
+        _quantize_fp4_indexer_kernel[(rows,)](
+            x.contiguous(),
+            q,
+            sf,
+            BLOCK_N=128,
+            GROUP_N=32,
+            RNE=rne,
+        )
+        return q, sf
+
+    for _ in range(3):
+        quantize_fp4_indexer_tensor(x, rne)
+        reference()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = quantize_fp4_indexer_tensor(x, rne)
+    boundaries = torch.tensor(
+        [0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0, 6.0], device="cuda", dtype=dtype
+    )
+    for scale in (0.0, 1e-6, 1.0, 1e3):
+        x.normal_().mul_(scale)
+        x[:2] = boundaries.repeat(16)
+        x[1].neg_()
+        graph.replay()
+        for a, b in zip(actual, reference()):
+            assert torch.equal(a, b)
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10,
+    reason="Vectorized prefill dispatch targets Blackwell",
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("rne", [False, True])
+def test_prefill_quantization_nonfinite_group_replay(dtype, rne):
+    from sglang.kernels.ops.attention.dsv4.fp4_indexer import (
+        _quantize_fp4_indexer_kernel,
+    )
+
+    rows = 4097  # Also exercise the final partial CTA.
+    x = torch.randn(rows, HEAD_DIM, device="cuda", dtype=dtype)
+    expected = (
+        torch.empty(rows, FP4_DIM, device="cuda", dtype=torch.int8),
+        torch.empty(rows, device="cuda", dtype=torch.int32),
+    )
+    for _ in range(3):
+        quantize_fp4_indexer_tensor(x, rne)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        actual = quantize_fp4_indexer_tensor(x, rne)
+    for value in (float("inf"), float("-inf"), float("nan"), 0.0):
+        x.normal_()
+        x[:, 0] = value
+        x[:, 32:64] = value
+        x[-1, 96:] = value
+        graph.replay()
+        _quantize_fp4_indexer_kernel[(rows,)](
+            x, *expected, BLOCK_N=HEAD_DIM, GROUP_N=GROUP_SIZE, RNE=rne
+        )
+        for a, b in zip(actual, expected):
+            assert torch.equal(a, b)
+
+
 @pytest.mark.parametrize("num_tokens", [1, 16, 96])
 def test_fp4_index_cache_store_layout(num_tokens: int) -> None:
     torch.manual_seed(num_tokens)