[Feature][JIT Kernel] Fused TP QK norm For Minimax (#20673)

Co-authored-by: Mingyang Jiang <13463932+jmydurant@users.noreply.github.com>

python/sglang/jit_kernel/all_reduce.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast
 
 import torch
 import tvm_ffi
+from tvm_ffi import Module
 
 from sglang.jit_kernel.utils import (
     cache_once,
@@ -92,7 +93,7 @@ if TYPE_CHECKING:
 
 
 @cache_once
-def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int):
+def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int) -> Module:
     args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
     return load_jit(
         "custom_all_reduce_pull",
@@ -104,7 +105,7 @@ def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int):
 
 
 @cache_once
-def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int):
+def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int) -> Module:
     args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
     return load_jit(
         "custom_all_reduce_push",
@@ -115,6 +116,24 @@ def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int):
     )
 
 
+@cache_once
+def _jit_fused_parallel_qknorm_module(
+    dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int
+) -> Module:
+    args = make_cpp_args(dtype, world_size, q_dim, k_dim, is_arch_support_pdl())
+    cls_name = f"FusedParallelQKNormAcrossHead<{args}>"
+    return load_jit(
+        "tp_qknorm",
+        *args,
+        extra_ldflags=["-lcuda"],
+        cuda_files=["distributed/tp_qknorm.cuh"],
+        cuda_wrappers=[
+            ("fused_parallel_qknorm", f"{cls_name}::run"),
+            ("get_max_occupancy", f"{cls_name}::get_max_occupancy"),
+        ],
+    )
+
+
 @cache_once
 def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
     module = load_jit(
@@ -144,17 +163,20 @@ def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
             max_pull_blocks: Optional[int] = None,
             max_push_blocks: Optional[int] = None,
         ) -> None:
+            max_pull_blocks = NUM_CTA if max_pull_blocks is None else max_pull_blocks
+            max_push_blocks = NUM_CTA if max_push_blocks is None else max_push_blocks
             self.__ffi_init__(
                 rank,
                 world_size,
-                NUM_CTA if max_pull_blocks is None else max_pull_blocks,
+                max_pull_blocks,
-                NUM_CTA if max_push_blocks is None else max_push_blocks,
+                max_push_blocks,
                 pull_buffer_bytes,
                 push_buffer_bytes,
                 graph_input_count,
             )
             self._world_size = world_size
-            self._pull_config = ConfigResult(NUM_CTA, MAX_THREADS)
+            self._pull_config = ConfigResult(min(NUM_CTA, max_pull_blocks), MAX_THREADS)
+            if max_pull_blocks > 0:  # special case: cannot configure 0 blocks
                 self.configure_pull(*self._pull_config)  # type: ignore
 
         @property
@@ -194,3 +216,25 @@ def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
             self.free_storage()  # type: ignore
 
     return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal)
+
+
+def get_fused_parallel_qknorm_max_occupancy(
+    dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int
+) -> int:
+    module = _jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim)
+    return module.get_max_occupancy()
+
+
+def fused_parallel_qknorm(
+    custom_ar: CustomAllReduceObj,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    world_size = custom_ar.world_size
+    q_dim = q.shape[-1] * world_size
+    k_dim = k.shape[-1] * world_size
+    module = _jit_fused_parallel_qknorm_module(q.dtype, world_size, q_dim, k_dim)
+    module.fused_parallel_qknorm(custom_ar, q, k, q_weight, k_weight, eps)

python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import argparse
+import os
+
+import torch
+import torch.distributed as dist
+
+import sglang.srt.distributed.parallel_state as ps
+from sglang.jit_kernel.all_reduce import (
+    fused_parallel_qknorm,
+    get_fused_parallel_qknorm_max_occupancy,
+)
+from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+    CustomAllReduceV2,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=120,
+    suite="stage-b-kernel-benchmark-1-gpu-large",
+    disabled="requires multi-GPU, self-skips in CI",
+)
+
+Q_K_DIMS = [(6144, 1024)]
+DTYPE = torch.bfloat16
+EPS = 1e-6
+BATCH_SIZES = get_ci_test_range([2**i for i in range(15)], [1, 64, 1024])
+NUM_LAYERS = 8
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--iters", type=int, default=100)
+    return parser.parse_args()
+
+
+def init_distributed():
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    rank = local_rank
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(backend="gloo")
+    ps._WORLD = coord = ps.init_world_group(
+        ranks=list(range(world_size)),
+        local_rank=local_rank,
+        backend="nccl",
+    )
+
+    cpu_group = coord.cpu_group
+    max_occupancy = get_fused_parallel_qknorm_max_occupancy(
+        DTYPE, world_size, Q_K_DIMS[0][0], Q_K_DIMS[0][1]
+    )
+    if rank == 0:
+        print(f"Max occupancy for fused_parallel_qknorm: {max_occupancy} blocks/SM")
+
+    props = torch.cuda.get_device_properties(device)
+    comm = CustomAllReduceV2(
+        cpu_group,
+        device,
+        max_pull_size=0,
+        max_push_size=8 * max(BATCH_SIZES),
+        max_push_blocks=props.multi_processor_count * max_occupancy,
+    )
+    comm_ = CustomAllReduceV2(cpu_group, device)
+    if comm.disabled or comm_.disabled:
+        raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
+    return rank, world_size, device, cpu_group, comm, comm_
+
+
+@torch.inference_mode()
+def bench_one(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn(0)
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        for i in range(NUM_LAYERS):
+            fn(i)
+
+    graph.replay()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    graph.replay()
+    start.record()
+    for i in range(iters):
+        graph.replay()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / (iters * NUM_LAYERS)
+
+
+def rmsnorm_baseline(
+    comm_,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    world_size: int,
+) -> None:
+    from sglang.srt.models.minimax_m2 import rms_apply_serial, rms_sumsq_serial
+
+    sum_sq = rms_sumsq_serial(q, k)
+    sum_sq = comm_.custom_all_reduce(sum_sq)
+    rms_apply_serial(q, k, q_weight, k_weight, sum_sq, world_size, EPS)
+
+
+def main():
+    args = parse_args()
+    rank, world_size, device, _, comm, comm_ = init_distributed()
+    torch.cuda.set_stream(torch.cuda.Stream())
+
+    if rank == 0:
+        print(
+            f"{'q_dim':>8} {'k_dim':>8} {'batch':>8} {'fused_us':>12} {'baseline_us':>12}"
+        )
+
+    for q_dim, k_dim in Q_K_DIMS:
+        local_q_dim = q_dim // world_size
+        local_k_dim = k_dim // world_size
+        for batch_size in BATCH_SIZES:
+            q = torch.randn(
+                NUM_LAYERS, batch_size, local_q_dim, device=device, dtype=DTYPE
+            )
+            k = torch.randn(
+                NUM_LAYERS, batch_size, local_k_dim, device=device, dtype=DTYPE
+            )
+            q_weight = torch.randn(NUM_LAYERS, local_q_dim, device=device, dtype=DTYPE)
+            k_weight = torch.randn(NUM_LAYERS, local_k_dim, device=device, dtype=DTYPE)
+
+            def run_fused(i: int):
+                fused_parallel_qknorm(
+                    comm.obj,
+                    q[i],
+                    k[i],
+                    q_weight[i],
+                    k_weight[i],
+                    EPS,
+                )
+
+            def run_baseline(i: int):
+                rmsnorm_baseline(
+                    comm_,
+                    q[i],
+                    k[i],
+                    q_weight[i],
+                    k_weight[i],
+                    world_size,
+                )
+
+            fused_us = bench_one(run_fused, args.warmup, args.iters)
+            baseline_us = bench_one(run_baseline, args.warmup, args.iters)
+
+            if rank == 0:
+                print(
+                    f"{q_dim:8d} {k_dim:8d} {batch_size:8d} "
+                    f"{fused_us:12.1f} {baseline_us:12.1f}"
+                )
+
+    comm.close()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh

@@ -0,0 +1,325 @@
+// Adapted from https://github.com/NVIDIA/TensorRT-LLM/pull/12163
+// We reuse the custom all reduce push buffer in SGLang
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <cstdint>
+#include <cstring>
+
+namespace {
+
+using device::distributed::PushController;
+using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
+
+struct ParallelQKNormParams {
+  void* __restrict__ buffer[device::distributed::kMaxNumGPU];
+  void* q_ptr;
+  void* k_ptr;
+  const void* __restrict__ q_weight;
+  const void* __restrict__ k_weight;
+  int64_t q_stride_bytes;
+  int64_t k_stride_bytes;
+  float eps;
+  uint32_t rank;
+  uint32_t num_tokens;
+  uint32_t epoch_bytes;
+  uint32_t num_clean_up_count = 0;
+};
+
+template <typename T>
+SGL_DEVICE void ld_global_volatile_8B(T& x, const void* addr, int64_t offset) {
+  static_assert(alignof(T) == 8 && sizeof(T) == 8);
+  addr = device::pointer::offset<T>(addr, offset);
+  uint2 val;
+  asm volatile("ld.volatile.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr));
+  x = *reinterpret_cast<const T*>(&val);
+}
+
+template <typename T>
+SGL_DEVICE void st_global_volatile_8B(const T& x, void* addr, int64_t offset) {
+  static_assert(alignof(T) == 8 && sizeof(T) == 8);
+  const uint2 val = *reinterpret_cast<const uint2*>(&x);
+  addr = device::pointer::offset<T>(addr, offset);
+  asm volatile("st.volatile.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr));
+}
+
+[[maybe_unused]]
+SGL_DEVICE float sync_float(float x) {
+  return __shfl_sync(0xffffffffu, x, 0);
+}
+
+[[maybe_unused]]
+constexpr auto next_pow_of_2(uint32_t x) {
+  uint32_t y = 1;
+  while (y < x)
+    y *= 2;
+  return y;
+}
+
+template <typename DType_, uint32_t kNumGPU_, int64_t kQDim_, int64_t kKDim_, bool kUsePDL_>
+struct KernelTrait {
+  // rename the arguments to avoid confusion with the template parameters
+  using DType = DType_;
+  static constexpr uint32_t kNumGPU = kNumGPU_;
+  static constexpr int64_t kQDim = kQDim_;
+  static constexpr int64_t kKDim = kKDim_;
+  static constexpr bool kUsePDL = kUsePDL_;
+
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static constexpr int64_t kLocalQDim = kQDim / kNumGPU;
+  static constexpr int64_t kLocalKDim = kKDim / kNumGPU;
+  static constexpr uint32_t kNumQThreads = kLocalQDim / (kVecSize * 2);
+  static constexpr uint32_t kNumKThreads = kLocalKDim / (kVecSize * 2);
+  static constexpr uint32_t kNumQWarps = kNumQThreads / device::kWarpThreads;
+  static constexpr uint32_t kNumKWarps = host::div_ceil(kNumKThreads, device::kWarpThreads);
+  static constexpr uint32_t kBlockSize = (kNumQWarps + kNumKWarps) * device::kWarpThreads;
+  static constexpr uint32_t kOccupancy = 2048 / kBlockSize;
+
+  using DType2 = packed_t<DType>;
+  using Storage = device::AlignedVector<DType2, kVecSize>;
+
+  static_assert(std::has_single_bit(kNumGPU), "must be pow of 2");
+  static_assert(kQDim % kNumGPU == 0);
+  static_assert(kKDim % kNumGPU == 0);
+  static_assert(kLocalQDim % (kVecSize * 2) == 0);
+  static_assert(kLocalKDim % (kVecSize * 2) == 0);
+  static_assert(kNumQThreads % device::kWarpThreads == 0);
+  static_assert(kBlockSize <= 1024);
+  static_assert(sizeof(Storage) == 16 && alignof(Storage) == 16);
+  static_assert(kOccupancy * kBlockSize <= 2048);
+};
+
+template <typename Trait>
+__global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel_qknorm_across_head(
+    const ParallelQKNormParams __grid_constant__ params, const PushController __grid_constant__ ctrl) {
+  using namespace device;
+
+  // each cta will handle exactly 1 token
+  using Storage = typename Trait::Storage;
+  using DType2 = typename Trait::DType2;
+  const auto &[
+      buffer, q_ptr, k_ptr, q_weight, k_weight, q_stride_bytes, k_stride_bytes, //
+      eps, rank, num_tokens, epoch_bytes, num_clean_up_count
+  ] = params;
+
+  using Package = AlignedVector<float, 2>;
+  constexpr uint32_t kNumGPU = Trait::kNumGPU;
+  constexpr uint32_t kNumQReduce = next_pow_of_2(Trait::kNumQWarps);
+  constexpr uint32_t kNumKReduce = next_pow_of_2(Trait::kNumKWarps);
+  __shared__ float smem_qk[Trait::kNumQWarps + Trait::kNumKWarps];
+  __shared__ float scale_q;
+  __shared__ float scale_k;
+  const auto tx = threadIdx.x;
+  const auto bx = blockIdx.x;
+  /// NOTE: this can hint compiler to optimize `is_valid` out when not needed
+  constexpr uint32_t kActiveThreads = Trait::kNumQThreads + Trait::kNumKThreads;
+  const auto is_valid = Trait::kBlockSize == kActiveThreads || tx < kActiveThreads;
+  const auto smem_q = smem_qk + 0;
+  const auto smem_k = smem_qk + Trait::kNumQWarps;
+  const auto load_q = tx < Trait::kNumQThreads;
+  const auto offset = load_q ? tx : tx - Trait::kNumQThreads;
+  const auto input_ptr = load_q ? q_ptr : k_ptr;
+  const auto weight_ptr = load_q ? q_weight : k_weight;
+  const auto input_stride_bytes = load_q ? q_stride_bytes : k_stride_bytes;
+  PDLWaitPrimary<Trait::kUsePDL>();
+  PDLTriggerSecondary<Trait::kUsePDL>();
+  if (bx >= num_tokens) {
+    [[unlikely]];
+    // In this case, we use the last few blocks to clean up other controllers
+    const auto start = (bx - num_tokens) * blockDim.x + threadIdx.x;
+    const auto stride = (gridDim.x - num_tokens) * blockDim.x;
+    for (uint32_t i = start; i < num_clean_up_count; i += stride)
+      ctrl.exit_unsafe(num_tokens + i);
+    return;
+  }
+  const auto epoch_offset = ctrl.epoch() * epoch_bytes;  // only for comm
+
+  __builtin_assume(bx < num_tokens);  // since we have `bx >= num_tokens`
+  Storage next_input;
+  void* input_i_ptr = pointer::offset(input_ptr, bx * input_stride_bytes);
+  if (is_valid) next_input.load(input_i_ptr, offset);
+
+  for (uint32_t i = bx; i < num_tokens; i += gridDim.x) {
+    // Stage 1. local reduce (warp-level)
+    Storage local_input;
+    {
+      float local_sum = 0.0;
+      if (is_valid) {
+        local_input = next_input;
+#pragma unroll
+        for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
+          const auto [x, y] = cast<fp32x2_t>(local_input[j]);
+          local_sum += x * x + y * y;
+        }
+      }
+      smem_qk[threadIdx.x / kWarpThreads] = warp::reduce_sum(local_sum);
+    }
+
+    // Stage 2. block reduce + push to peer ranks + poll from local rank
+    __syncthreads();
+
+    Storage local_weight;
+    const auto input_next_ptr = pointer::offset(input_i_ptr, gridDim.x * input_stride_bytes);
+    /**
+     * NOTE: Prefetch to hide the latency.
+     * This brings around 20% of performance gain in large batches
+     * The P2P communication is mainly latency bound, so during this waiting period,
+     * We can let some data loading transparently in the background.
+     */
+    if (is_valid) {
+      local_weight.load(weight_ptr, offset);
+      if (i + gridDim.x < num_tokens) next_input.load(input_next_ptr, offset);
+    }
+
+    if (tx < kWarpThreads) {
+      const auto local_sum_q = tx < Trait::kNumQWarps ? smem_q[tx] : 0.0f;
+      const auto local_sum_k = tx < Trait::kNumKWarps ? smem_k[tx] : 0.0f;
+      const auto sum_q = sync_float(warp::reduce_sum<kNumQReduce>(local_sum_q));
+      const auto sum_k = sync_float(warp::reduce_sum<kNumKReduce>(local_sum_k));
+      if (tx < kNumGPU) {  // push a float2 pack to the peer
+        Package sum_q_k;
+        /// NOTE: eps should be scaled down by kNumGPU from host side
+        /// we add here to ensure that the sum is never zero
+        sum_q_k[0] = sum_q + eps;
+        sum_q_k[1] = sum_k + eps;
+        const auto push_ptr = pointer::offset(buffer[tx], epoch_offset);
+        st_global_volatile_8B(sum_q_k, push_ptr, i * kNumGPU + rank);
+        const auto poll_ptr = pointer::offset(buffer[rank], epoch_offset);
+        while (true) {
+          ld_global_volatile_8B(sum_q_k, poll_ptr, i * kNumGPU + tx);
+          if (sum_q_k[0] != 0.0f && sum_q_k[1] != 0.0f) break;
+        }
+        constexpr uint32_t kActiveMask = (1 << kNumGPU) - 1;
+        const auto global_sum_q = warp::reduce_sum<kNumGPU>(sum_q_k[0], kActiveMask);
+        const auto global_sum_k = warp::reduce_sum<kNumGPU>(sum_q_k[1], kActiveMask);
+        scale_q = math::rsqrt(global_sum_q / static_cast<float>(Trait::kQDim));
+        scale_k = math::rsqrt(global_sum_k / static_cast<float>(Trait::kKDim));
+        Package zeros;
+        zeros.fill(0.0f);
+        zeros.store(poll_ptr, i * kNumGPU + tx);
+      }
+    }
+
+    __syncthreads();
+    const auto scale = load_q ? scale_q : scale_k;
+    if (is_valid) {
+#pragma unroll
+      for (uint32_t j = 0; j < Trait::kVecSize; ++j) {
+        const auto fp32_input = cast<fp32x2_t>(local_input[j]);
+        const auto fp32_weight = cast<fp32x2_t>(local_weight[j]);
+        const auto scaled_x = fp32_input.x * scale * fp32_weight.x;
+        const auto scaled_y = fp32_input.y * scale * fp32_weight.y;
+        local_input[j] = cast<DType2>(fp32x2_t{scaled_x, scaled_y});
+      }
+      local_input.store(input_i_ptr, offset);
+    }
+    input_i_ptr = input_next_ptr;
+  }
+  ctrl.exit();
+}
+
+template <typename DType, uint32_t kNumGPU, int64_t kQDim, int64_t kKDim, bool kUsePDL>
+struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase {
+  using Trait = KernelTrait<DType, kNumGPU, kQDim, kKDim, kUsePDL>;
+  static constexpr auto kernel = parallel_qknorm_across_head<Trait>;
+  static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
+
+  void _run(
+      const tvm::ffi::Tensor q,
+      const tvm::ffi::Tensor k,
+      const tvm::ffi::Tensor q_weight,
+      const tvm::ffi::Tensor k_weight,
+      const float eps  // passed in unscaled
+  ) {
+    using namespace host;
+    constexpr auto Q = Trait::kLocalQDim;
+    constexpr auto K = Trait::kLocalKDim;
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    TensorMatcher({N, Q})  // q
+        .with_strides({-1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q);
+    TensorMatcher({N, K})  // k
+        .with_strides({-1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(k);
+    TensorMatcher({Q})  // q_weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q_weight);
+    TensorMatcher({K})  // k_weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(k_weight);
+    const auto device = device_.unwrap();
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    // use at most `world_size` blocks to clean up,
+    // this is based on the observation that occupancy is usually linear
+    // with respect to the world size
+    const bool need_clean = num_tokens < m_max_num_cta_push;
+    const auto num_clean = need_clean ? (m_max_num_cta_push - num_tokens) : 0;
+    const auto num_blocks = need_clean ? num_tokens + div_ceil(num_clean, Trait::kBlockSize)  //
+                                       : m_max_num_cta_push;                                  //
+    const auto num_threads = Trait::kBlockSize;
+    RuntimeCheck(num_blocks <= m_max_num_cta_push, "internal error");
+    ParallelQKNormParams params;
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      params.buffer[i] = get_push_buffer(m_peer_storage[i]);
+    }
+    params.q_ptr = q.data_ptr();
+    params.k_ptr = k.data_ptr();
+    params.q_weight = q_weight.data_ptr();
+    params.k_weight = k_weight.data_ptr();
+    params.q_stride_bytes = q.stride(0) * sizeof(DType);
+    params.k_stride_bytes = k.stride(0) * sizeof(DType);
+    params.eps = eps / kNumGPU;  // scale down eps by number of GPUs
+    params.rank = m_rank;
+    params.num_tokens = num_tokens;
+    params.epoch_bytes = m_push_buffer_bytes;
+    params.num_clean_up_count = num_clean;
+
+    const auto needed_buffer_bytes = static_cast<int64_t>(num_tokens) * 2 * sizeof(float);
+    RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
+    RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.q_ptr) % 16 == 0, "q pointer is not properly aligned");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.k_ptr) % 16 == 0, "k pointer is not properly aligned");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned");
+    RuntimeCheck(std::bit_cast<intptr_t>(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned");
+    RuntimeCheck(needed_buffer_bytes <= m_push_buffer_bytes, "Push buffer is too small");
+
+    LaunchKernel(num_blocks, num_threads, device)  //
+        .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
+  }
+
+  static uint32_t get_max_occupancy() {
+    return host::runtime::get_blocks_per_sm(kernel, Trait::kBlockSize);
+  }
+
+  static void
+  run(CustomAllReduceRef obj,
+      const tvm::ffi::Tensor q,
+      const tvm::ffi::Tensor k,
+      const tvm::ffi::Tensor q_weight,
+      const tvm::ffi::Tensor k_weight,
+      const float eps) {
+    using Self = FusedParallelQKNormAcrossHead;
+    return static_cast<Self*>(obj.get())->_run(q, k, q_weight, k_weight, eps);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh

@@ -48,6 +48,8 @@ struct alignas(128) Semaphore {
 
 struct PullController {
  public:
+  using SignalType = Semaphore;
+
   PullController(void** signals, uint32_t num_gpu) {
     for (uint32_t i = 0; i < num_gpu; ++i) {
       m_signals[i] = static_cast<Semaphore*>(signals[i]);
@@ -90,25 +92,29 @@ struct PullController {
 
 struct PushController {
  public:
+  using SignalType = uint32_t;
   static constexpr int64_t kNumStages = 2;
 
-  PushController(void* ptr) : m_local_signal(static_cast<Semaphore*>(ptr)) {}
+  PushController(void* ptr) : m_local_signal(static_cast<SignalType*>(ptr)) {}
 
-  SGL_DEVICE uint32_t epoch() const {
+  SGL_DEVICE SignalType epoch() const {
-    return m_local_signal[blockIdx.x].get_counter();
+    return m_local_signal[blockIdx.x];
   }
 
   SGL_DEVICE void exit() const {
     __syncthreads();
     if (threadIdx.x == 0) {
-      auto& signal = m_local_signal[blockIdx.x];
+      this->exit_unsafe(blockIdx.x);
-      const auto epoch = signal.get_counter();
-      signal.set_counter((epoch + 1) % kNumStages);
     }
   }
 
+  SGL_DEVICE void exit_unsafe(uint32_t which) const {
+    auto& signal = m_local_signal[which];
+    signal = (signal + 1) % kNumStages;
+  }
+
  private:
-  Semaphore* m_local_signal;
+  SignalType* m_local_signal;
 };
 
 }  // namespace device::distributed

python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh

@@ -93,12 +93,14 @@ struct CustomAllReduceBase : public tvm::ffi::Object {
         // default config for pull kernel, can be updated by `configure()`
         m_num_cta(max_num_cta_pull),
         m_cta_size(256) {
-    RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
+    RuntimeCheck(pull_buffer_size % 128 == 0, "Pull buffer size should be aligned to 128 bytes");
+    RuntimeCheck(push_buffer_size % 128 == 0, "Push buffer size should be aligned to 128 bytes");
     RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank);
     const int64_t kU32Max = static_cast<int64_t>(std::numeric_limits<uint32_t>::max());
     const int64_t push_buffer_size_all = push_all_ranks_bytes();
-    RuntimeCheck(pull_buffer_size <= kU32Max, "Buffer size is too large: ", pull_buffer_size);
+    RuntimeCheck(pull_buffer_size <= kU32Max, "Pull buffer size is too large: ", pull_buffer_size);
     RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all);
+    RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
   }
 
   ExternHandle share_storage() {
@@ -252,19 +254,18 @@ struct CustomAllReduceBase : public tvm::ffi::Object {
     return static_cast<int64_t>(m_graph_capture_inputs.size());
   }
   int64_t pull_signal_bytes() const {
-    return sizeof(device::distributed::Semaphore) * m_max_num_cta_pull;
+    return _align_bytes(sizeof(PullController::SignalType) * m_max_num_cta_pull);
   }
   int64_t push_signal_bytes() const {
-    return sizeof(device::distributed::Semaphore) * m_max_num_cta_push;
+    return _align_bytes(sizeof(PushController::SignalType) * m_max_num_cta_push);
   }
-  int64_t params_bytes() const {
+  int64_t graph_param_bytes() const {
-    return sizeof(AllReduceData) * (1 + m_graph_buffer_count);  // 1 for default
+    return _align_bytes(sizeof(AllReduceData) * (1 + m_graph_buffer_count));  // 1 for default
   }
   int64_t push_all_ranks_bytes() const {
-    return PushController::kNumStages * m_num_gpu * m_push_buffer_bytes;
+    return _align_bytes(PushController::kNumStages * m_num_gpu * m_push_buffer_bytes);
   }
   int64_t storage_bytes() const {
-    // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
     return _get_offset_impl(5);
   }
   void* get_pull_signal(void* ptr) const {
@@ -283,16 +284,20 @@ struct CustomAllReduceBase : public tvm::ffi::Object {
     return pointer::offset(ptr, _get_offset_impl(4));
   }
   int64_t _get_offset_impl(int64_t which) const {
+    // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
     const int64_t offset_map[5] = {
         /*[0]=*/pull_signal_bytes(),
         /*[1]=*/push_signal_bytes(),
-        /*[2]=*/params_bytes(),
+        /*[2]=*/graph_param_bytes(),
         /*[3]=*/m_pull_buffer_bytes,
         /*[4]=*/push_all_ranks_bytes(),
     };
     RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which);
     return std::accumulate(offset_map, offset_map + which, int64_t(0));
   }
+  static int64_t _align_bytes(int64_t size) {
+    return div_ceil(size, 128) * 128;
+  }
 
   const int64_t m_pull_buffer_bytes;
   const int64_t m_push_buffer_bytes;

python/sglang/jit_kernel/include/sgl_kernel/warp.cuh

@@ -21,10 +21,12 @@ static constexpr uint32_t kFullMask = 0xffffffffu;
  * \param active_mask Bitmask of participating lanes (default: all 32).
  * \return The sum across all active lanes.
  */
-template <typename T>
+template <uint32_t kNumThreads = kWarpThreads, typename T>
 SGL_DEVICE T reduce_sum(T value, uint32_t active_mask = kFullMask) {
+  static_assert(kNumThreads >= 1 && kNumThreads <= kWarpThreads);
+  static_assert(std::has_single_bit(kNumThreads), "must be pow of 2");
 #pragma unroll
-  for (int mask = 16; mask > 0; mask >>= 1)
+  for (int mask = kNumThreads / 2; mask > 0; mask >>= 1)
     value = value + __shfl_xor_sync(active_mask, value, mask, 32);
   return value;
 }

python/sglang/jit_kernel/tests/test_custom_all_reduce.py

@@ -18,8 +18,6 @@ import itertools
 import logging
 import multiprocessing as mp
 import os
-import subprocess
-import sys
 from typing import Dict, Optional, Tuple
 
 import pytest
@@ -32,6 +30,7 @@ from sglang.jit_kernel.all_reduce import (
     _jit_custom_all_reduce_pull_module,
     _jit_custom_all_reduce_push_module,
 )
+from sglang.jit_kernel.tests.utils import multiprocess_main, multiprocess_test
 from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
     CustomAllReduceV2,
 )
@@ -79,26 +78,6 @@ TEST_LOOP = 16
 # ---------------------------------------------------------------------------
 
 
-def _run_torchrun(nproc: int, timeout: int = 300) -> None:
-    """Launch this script as a torchrun worker and assert success."""
-    cmd = [
-        "torchrun",
-        f"--nproc_per_node={nproc}",
-        __file__,
-    ]
-    result = subprocess.run(
-        cmd,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-        timeout=timeout,
-    )
-    assert result.returncode == 0, (
-        f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
-        f"{result.stdout}"
-    )
-
-
 def _compile_one(dtype: torch.dtype, world_size: int):
     _jit_custom_all_reduce_push_module(dtype, world_size)
     _jit_custom_all_reduce_pull_module(dtype, world_size)
@@ -129,7 +108,7 @@ def test_custom_allreduce(nproc: int) -> None:
         pytest.skip(
             f"Requires at least {nproc} GPUs, but only {device_count} available"
         )
-    _run_torchrun(nproc)
+    multiprocess_test(__file__, nproc)
 
 
 # ---------------------------------------------------------------------------
@@ -229,7 +208,6 @@ def worker_test(
 def worker_main() -> None:
     """Entry point for each torchrun worker process."""
     rank, device, cpu_group, nccl_group, comm = init_distributed()
-    world_size = dist.get_world_size()
 
     torch.cuda.set_stream(torch.cuda.Stream())
 
@@ -258,7 +236,4 @@ def worker_main() -> None:
 
 
 if __name__ == "__main__":
-    if "LOCAL_RANK" in os.environ:
+    multiprocess_main(__file__, worker_main)
-        worker_main()
-    else:
-        sys.exit(pytest.main([__file__, "-x", "-vv", "-s"]))

python/sglang/jit_kernel/tests/test_tp_qknorm.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import itertools
+import os
+from typing import Optional
+
+import pytest
+import torch
+import torch.distributed as dist
+import triton
+
+from sglang.jit_kernel.all_reduce import fused_parallel_qknorm
+from sglang.jit_kernel.tests.test_custom_all_reduce import multiprocess_test
+from sglang.jit_kernel.tests.utils import multiprocess_main
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=300,
+    suite="stage-b-kernel-unit-8-gpu-h200",
+)
+register_cuda_ci(
+    est_time=300,
+    suite="nightly-kernel-8-gpu-h200",
+    nightly=True,
+)
+
+
+Q_K_DIMS = [(6144, 1024)]
+EPS = 1e-6
+BATCH_SIZES = [2**n for n in range(0, 14)]
+DTYPES = [torch.float16, torch.bfloat16, torch.float32]
+TEST_CONFIG = list(itertools.product(Q_K_DIMS, BATCH_SIZES, DTYPES))
+
+
+@pytest.mark.parametrize("nproc", [2, 4, 8])
+def test_tp_qknorm(nproc: int) -> None:
+    device_count = torch.cuda.device_count()
+    if device_count < nproc:
+        pytest.skip(
+            f"Requires at least {nproc} GPUs, but only {device_count} available"
+        )
+    multiprocess_test(__file__, nproc)
+
+
+def init_distributed():
+    import sglang.srt.distributed.parallel_state as ps
+    from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+        CustomAllReduceV2,
+    )
+
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    rank = local_rank
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(backend="gloo")
+    ps._WORLD = coord = ps.init_world_group(
+        ranks=list(range(world_size)),
+        local_rank=local_rank,
+        backend="nccl",
+    )
+
+    cpu_group = coord.cpu_group
+    nccl_group = coord.device_group
+    assert nccl_group is not None
+
+    max_pull_size = 0
+    max_push_size = 8 * max(BATCH_SIZES)
+    comm = CustomAllReduceV2(cpu_group, device, max_pull_size, max_push_size)
+    if comm.disabled:
+        raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
+
+    return rank, world_size, device, cpu_group, nccl_group, comm
+
+
+def _all_gather_cat(x: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
+    gathered = [torch.empty_like(x) for _ in range(dist.get_world_size(group=group))]
+    dist.all_gather(gathered, x, group=group)
+    return torch.cat(gathered, dim=-1)
+
+
+def _rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    x_fp32 = x.float()
+    scale = (x_fp32.pow(2).mean(dim=-1, keepdim=True) + eps).rsqrt()
+    return (x_fp32 * scale * weight.float()).to(x.dtype)
+
+
+@torch.inference_mode()
+def worker_test(
+    rank: int,
+    world_size: int,
+    device: torch.device,
+    nccl_group: dist.ProcessGroup,
+    comm,
+    q_k_dim: tuple[int, int],
+    batch_size: int,
+    dtype: torch.dtype,
+) -> Optional[RuntimeError]:
+    q_dim, k_dim = q_k_dim
+    local_q_dim = q_dim // world_size
+    local_k_dim = k_dim // world_size
+
+    q = torch.randn(batch_size, local_q_dim, device=device, dtype=dtype)
+    k = torch.randn(batch_size, local_k_dim, device=device, dtype=dtype)
+    q_weight = torch.randn(local_q_dim, device=device, dtype=dtype)
+    k_weight = torch.randn(local_k_dim, device=device, dtype=dtype)
+
+    q_ref = _all_gather_cat(q, nccl_group)
+    k_ref = _all_gather_cat(k, nccl_group)
+    q_weight_ref = _all_gather_cat(q_weight.unsqueeze(0), nccl_group).squeeze(0)
+    k_weight_ref = _all_gather_cat(k_weight.unsqueeze(0), nccl_group).squeeze(0)
+
+    q_expected = _rmsnorm_ref(q_ref, q_weight_ref, EPS)
+    k_expected = _rmsnorm_ref(k_ref, k_weight_ref, EPS)
+    q_expected = q_expected[:, rank * local_q_dim : (rank + 1) * local_q_dim]
+    k_expected = k_expected[:, rank * local_k_dim : (rank + 1) * local_k_dim]
+
+    fused_parallel_qknorm(
+        comm.obj,
+        q,
+        k,
+        q_weight,
+        k_weight,
+        EPS,
+    )
+
+    try:
+        triton.testing.assert_close(q, q_expected, atol=1e-2, rtol=1e-2)
+        triton.testing.assert_close(k, k_expected, atol=1e-2, rtol=1e-2)
+    except AssertionError as err:
+        return RuntimeError(
+            f"TP QKNorm mismatch for {batch_size=}, {dtype=}, {world_size=}, {rank=}: {err}"
+        )
+    return None
+
+
+def worker_main() -> None:
+    rank, world_size, device, cpu_group, nccl_group, comm = init_distributed()
+    torch.cuda.set_stream(torch.cuda.Stream())
+
+    for q_k_dim, batch_size, dtype in TEST_CONFIG:
+        error = worker_test(
+            rank,
+            world_size,
+            device,
+            nccl_group,
+            comm,
+            q_k_dim,
+            batch_size,
+            dtype,
+        )
+        result = torch.tensor([int(error is not None)])
+        dist.all_reduce(result, group=cpu_group)
+        if error is not None:
+            print(str(error))
+        if bool(result.item()):
+            raise RuntimeError(
+                f"TP QKNorm test failed for {q_k_dim=}, {batch_size=}, {dtype=}, {world_size=}"
+            )
+
+    print(f"Rank {rank} passed all tests.")
+    comm.close()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    multiprocess_main(__file__, worker_main)

python/sglang/jit_kernel/tests/utils.py

@@ -0,0 +1,41 @@
+import os
+import subprocess
+import sys
+from typing import Callable
+
+import pytest
+
+
+def multiprocess_test(file: str, nproc: int, timeout: int = 90) -> None:
+    """Launch this script as a torchrun worker and assert success."""
+    cmd = [
+        "torchrun",
+        f"--nproc_per_node={nproc}",
+        file,
+    ]
+    try:
+        result = subprocess.run(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            timeout=timeout,
+        )
+    except subprocess.TimeoutExpired as e:
+        raise RuntimeError(
+            f"torchrun (nproc={nproc}) timed out after {timeout} seconds\n"
+            f"{e.stdout}"
+        ) from e
+
+    assert result.returncode == 0, (
+        f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
+        f"{result.stdout}"
+    )
+
+
+def multiprocess_main(file: str, main: Callable[[], None]) -> None:
+    """Helper to run a function in a multiprocess torchrun context."""
+    if "LOCAL_RANK" in os.environ:
+        main()
+    else:
+        sys.exit(pytest.main([file, "-v", "-s"]))

python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py

@@ -35,6 +35,8 @@ class CustomAllReduceV2:
         device: torch.device,
         max_pull_size: Optional[int] = None,
         max_push_size: Optional[int] = None,
+        max_pull_blocks: Optional[int] = None,
+        max_push_blocks: Optional[int] = None,
     ) -> None:
         _init_config()
         self.disabled = True
@@ -50,14 +52,15 @@ class CustomAllReduceV2:
         self.group = group
         self.rank = dist.get_rank(group=self.group)
         self.world_size = dist.get_world_size(group=self.group)
-        self.override_shot(None)
+        if max_pull_size is None:  # default to 16MB
-        if max_pull_size is None:
+            max_pull_size = 16 * 1024 * 1024
-            max_pull_size = 16 * 1024 * 1024  # default to 16MB
+        if max_push_size is None:  # default to recommended size
-        if max_push_size is None:
+            config = THRESHOLD_2_SHOT_MAP[self.world_size]
-            max_push_size = self.config.one_shot_push_threshold
+            max_push_size = config.one_shot_push_threshold
-        max_push_size = min(max_push_size, max_pull_size)
         self.max_pull_size = max_pull_size
         self.max_push_size = max_push_size
+        self.max_size = max(max_pull_size, max_push_size)
+        self.override_shot(None)  # set default config based on world size
         self.override_algo: Optional[AllReduceAlgo] = None
         self.obj = get_custom_all_reduce_cls()(
             rank=self.rank,
@@ -65,6 +68,8 @@ class CustomAllReduceV2:
             pull_buffer_bytes=self.max_pull_size,
             push_buffer_bytes=self.max_push_size,
             graph_input_count=131072,
+            max_pull_blocks=max_pull_blocks,
+            max_push_blocks=max_push_blocks,
         )
         self._post_init_obj()
         self.disabled = False
@@ -72,11 +77,19 @@ class CustomAllReduceV2:
 
     def override_shot(self, shot: int | None):
         if shot is None:
-            self.config = THRESHOLD_2_SHOT_MAP[self.world_size]
+            config = THRESHOLD_2_SHOT_MAP[self.world_size]
         else:
             assert shot in (1, 2)
             threshold = INF if shot == 1 else 0
-            self.config = replace(self.config, one_shot_pull_threshold=threshold)
+            config = replace(self.config, one_shot_pull_threshold=threshold)
+        # need to clip the config thresholds to max sizes to avoid invalid config
+        push_threshold = min(config.one_shot_push_threshold, self.max_push_size)
+        pull_threshold = min(config.one_shot_pull_threshold, self.max_pull_size)
+        self.config: ModeConfig = replace(
+            config,
+            one_shot_push_threshold=push_threshold,
+            one_shot_pull_threshold=pull_threshold,
+        )
 
     @contextmanager
     def capture(self):
@@ -109,7 +122,7 @@ class CustomAllReduceV2:
             return False
         if not is_weak_contiguous(inp):
             return False
-        return inp_size <= self.max_pull_size
+        return inp_size <= self.max_size
 
     def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
         if is_in_piecewise_cuda_graph():  # disable inplace optimization

python/sglang/srt/models/minimax_m2.py

@@ -17,7 +17,8 @@
 
 import logging
 from contextlib import nullcontext
-from typing import Iterable, Optional, Set, Tuple, Union
+from functools import lru_cache
+from typing import Any, Dict, Iterable, Optional, Set, Tuple, Union
 
 import torch
 import triton
@@ -25,9 +26,14 @@ import triton.language as tl
 from torch import nn
 from transformers import PretrainedConfig
 
+from sglang.jit_kernel.all_reduce import (
+    fused_parallel_qknorm,
+    get_fused_parallel_qknorm_max_occupancy,
+)
 from sglang.kernel_api_logging import debug_kernel_api
 from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo
 from sglang.srt.distributed import (
+    get_bool_env_var,
     get_moe_expert_parallel_world_size,
     get_pp_group,
     get_tensor_model_parallel_world_size,
@@ -42,6 +48,7 @@ from sglang.srt.layers.communicator import (
 )
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_reduce,
+    get_attention_tp_group,
     get_attention_tp_rank,
     get_attention_tp_size,
     is_dp_attention_enabled,
@@ -78,12 +85,15 @@ from sglang.srt.utils import (
     BumpAllocator,
     add_prefix,
     get_compiler_backend,
+    is_cuda,
     is_non_idle_and_non_empty,
     make_layers,
 )
+from sglang.srt.utils.custom_op import register_custom_op
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 logger = logging.getLogger(__name__)
+_is_cuda = is_cuda()
 
 
 @triton.jit
@@ -324,27 +334,114 @@ class MiniMaxM2RMSNormTP(nn.Module):
 
         return x
 
-    @staticmethod
+
-    def forward_qk(
+@register_custom_op(mutates_args=["q", "k"])
-        q_norm: "MiniMaxM2RMSNormTP",
+def fused_tp_qknorm(
-        k_norm: "MiniMaxM2RMSNormTP",
+    counter: int,
     q: torch.Tensor,
     k: torch.Tensor,
-    ) -> torch.Tensor:
+    q_weight: torch.Tensor,
-        sum_sq = rms_sumsq_serial(q, k)
+    k_weight: torch.Tensor,
-        if q_norm.attn_tp_size > 1:
+    eps: float,
-            sum_sq = attn_tp_all_reduce(sum_sq)
+) -> None:
-
+    return fused_parallel_qknorm(
-        q, k = rms_apply_serial(
+        MiniMaxM2QKRMSNorm.COMM_MAP[counter].obj,
         q,
         k,
-            q_norm.weight,
+        q_weight,
-            k_norm.weight,
+        k_weight,
-            sum_sq,
+        eps=eps,
-            q_norm.attn_tp_size,
-            q_norm.variance_epsilon,
     )
 
+
+class MiniMaxM2QKRMSNorm:
+    COUNTER = 0
+    COMM_MAP: Dict[int, Any] = {}
+
+    def __init__(
+        self,
+        q_norm: MiniMaxM2RMSNormTP,
+        k_norm: MiniMaxM2RMSNormTP,
+    ) -> None:
+        assert q_norm.variance_epsilon == k_norm.variance_epsilon
+        self._q_norm = q_norm
+        self._k_norm = k_norm
+        self._world_size = self._q_norm.attn_tp_size
+        self._eps = q_norm.variance_epsilon
+        use_fused_norm = get_bool_env_var("SGLANG_USE_FUSED_PARALLEL_QKNORM")
+
+        self._forward_impl = self._forward_naive
+        if self._world_size > 1 and _is_cuda and use_fused_norm:
+            occupancy = get_fused_parallel_qknorm_max_occupancy(
+                q_norm.weight.dtype,
+                self._world_size,
+                # NOTE: we need full dimension
+                q_dim=q_norm.weight.shape[0] * self._world_size,
+                k_dim=k_norm.weight.shape[0] * self._world_size,
+            )
+            counter = MiniMaxM2QKRMSNorm._get_comm(q_norm.weight.device, occupancy)
+            if counter is not None:
+                self._counter = counter
+                self._forward_impl = self._forward_fused
+
+    @lru_cache
+    @staticmethod
+    def _get_comm(device: torch.device, occupancy: int):
+        from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+            CustomAllReduceV2,
+        )
+
+        props = torch.cuda.get_device_properties(device)
+        # probe the maximum tokens for one prefill
+        server_args = get_global_server_args()
+        max_tokens = server_args.chunked_prefill_size
+        if max_tokens is None:
+            max_tokens = server_args.model_config.context_len
+        max_tokens = max(max_tokens, server_args.max_prefill_tokens)
+        logger.info(f"[AR] Using CustomAllReduceV2 for MiniMaxM2 with {max_tokens = }")
+        ALIGN = 512
+        # typically, this should not exceed 1M, since max_tokens is usually less than 16384
+        max_size = ((8 * max_tokens + ALIGN - 1) // ALIGN) * ALIGN
+        comm = CustomAllReduceV2(
+            group=get_attention_tp_group().cpu_group,
+            device=device,
+            max_pull_size=0,
+            max_pull_blocks=0,
+            max_push_size=max_size,
+            max_push_blocks=props.multi_processor_count * occupancy,
+        )
+        counter = MiniMaxM2QKRMSNorm.COUNTER
+        MiniMaxM2QKRMSNorm.COUNTER += 1
+        MiniMaxM2QKRMSNorm.COMM_MAP[counter] = comm
+        return counter if not comm.disabled else None
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor):
+        return self._forward_impl(q, k)
+
+    def _forward_naive(self, q: torch.Tensor, k: torch.Tensor):
+        q, k = q.contiguous(), k.contiguous()
+        sum_sq = rms_sumsq_serial(q, k)
+        if self._world_size > 1:
+            sum_sq = attn_tp_all_reduce(sum_sq)
+        return rms_apply_serial(
+            q,
+            k,
+            self._q_norm.weight,
+            self._k_norm.weight,
+            sum_sq,
+            self._world_size,
+            self._eps,
+        )
+
+    def _forward_fused(self, q: torch.Tensor, k: torch.Tensor):
+        fused_tp_qknorm(
+            self._counter,
+            q,
+            k,
+            self._q_norm.weight,
+            self._k_norm.weight,
+            self._eps,
+        )
         return q, k
 
 
@@ -681,6 +778,7 @@ class MiniMaxM2Attention(nn.Module):
                     num_heads=self.total_num_kv_heads,
                     eps=config.rms_norm_eps,
                 )
+                self.qk_norm_impl = MiniMaxM2QKRMSNorm(self.q_norm, self.k_norm)
             else:
                 raise ValueError(f"Unsupported qk_norm_type: {self.qk_norm_type}")
 
@@ -708,13 +806,7 @@ class MiniMaxM2Attention(nn.Module):
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         if self.use_qk_norm:
-            # q = self.q_norm(q.contiguous())
+            q, k = self.qk_norm_impl.forward(q, k)
-            # k = self.k_norm(k.contiguous())
-            q, k = MiniMaxM2RMSNormTP.forward_qk(
-                self.q_norm, self.k_norm, q.contiguous(), k.contiguous()
-            )
-        else:
-            q, k = q.contiguous(), k.contiguous()
         q, k = self.rotary_emb(positions, q, k)
         inner_state = q, k, v, forward_batch
         return None, forward_batch, inner_state