diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/jit_kernel/all_reduce.py index 48f763259..26167c698 100644 --- a/python/sglang/jit_kernel/all_reduce.py +++ b/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,18 +163,21 @@ 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, - NUM_CTA if max_push_blocks is None else max_push_blocks, + max_pull_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.configure_pull(*self._pull_config) # type: ignore + 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 def world_size(self) -> int: @@ -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) diff --git a/python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py b/python/sglang/jit_kernel/benchmark/bench_tp_qknorm.py new file mode 100644 index 000000000..d63545b10 --- /dev/null +++ b/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() diff --git a/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh b/python/sglang/jit_kernel/csrc/distributed/tp_qknorm.cuh new file mode 100644 index 000000000..ca80e1efc --- /dev/null +++ b/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 +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +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 +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(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(&val); +} + +template +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(&x); + addr = device::pointer::offset(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 +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; + using Storage = device::AlignedVector; + + 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 +__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; + 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(); + PDLTriggerSecondary(); + 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(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(local_sum_q)); + const auto sum_k = sync_float(warp::reduce_sum(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(sum_q_k[0], kActiveMask); + const auto global_sum_k = warp::reduce_sum(sum_q_k[1], kActiveMask); + scale_q = math::rsqrt(global_sum_q / static_cast(Trait::kQDim)); + scale_k = math::rsqrt(global_sum_k / static_cast(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(local_input[j]); + const auto fp32_weight = cast(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(fp32x2_t{scaled_x, scaled_y}); + } + local_input.store(input_i_ptr, offset); + } + input_i_ptr = input_next_ptr; + } + ctrl.exit(); +} + +template +struct FusedParallelQKNormAcrossHead : public CustomAllReduceBase { + using Trait = KernelTrait; + static constexpr auto kernel = parallel_qknorm_across_head; + 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(); + TensorMatcher({N, Q}) // q + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(q); + TensorMatcher({N, K}) // k + .with_strides({-1, 1}) + .with_dtype() + .with_device(device_) + .verify(k); + TensorMatcher({Q}) // q_weight + .with_dtype() + .with_device(device_) + .verify(q_weight); + TensorMatcher({K}) // k_weight + .with_dtype() + .with_device(device_) + .verify(k_weight); + const auto device = device_.unwrap(); + const auto num_tokens = static_cast(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(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(params.q_ptr) % 16 == 0, "q pointer is not properly aligned"); + RuntimeCheck(std::bit_cast(params.k_ptr) % 16 == 0, "k pointer is not properly aligned"); + RuntimeCheck(std::bit_cast(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned"); + RuntimeCheck(std::bit_cast(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(obj.get())->_run(q, k, q_weight, k_weight, eps); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh index c9711c5b6..e0ce2dc08 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh +++ b/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(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(ptr)) {} + PushController(void* ptr) : m_local_signal(static_cast(ptr)) {} - SGL_DEVICE uint32_t epoch() const { - return m_local_signal[blockIdx.x].get_counter(); + SGL_DEVICE SignalType epoch() const { + return m_local_signal[blockIdx.x]; } SGL_DEVICE void exit() const { __syncthreads(); if (threadIdx.x == 0) { - auto& signal = m_local_signal[blockIdx.x]; - const auto epoch = signal.get_counter(); - signal.set_counter((epoch + 1) % kNumStages); + this->exit_unsafe(blockIdx.x); } } + 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 diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh index 7b627a9e9..239fac71a 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh +++ b/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(std::numeric_limits::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(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 { - return sizeof(AllReduceData) * (1 + m_graph_buffer_count); // 1 for default + int64_t graph_param_bytes() const { + 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; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh b/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh index 52952ff81..769c16c08 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/warp.cuh +++ b/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 +template 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; } diff --git a/python/sglang/jit_kernel/tests/test_custom_all_reduce.py b/python/sglang/jit_kernel/tests/test_custom_all_reduce.py index 2d7e0253e..96a61bfbf 100644 --- a/python/sglang/jit_kernel/tests/test_custom_all_reduce.py +++ b/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: - worker_main() - else: - sys.exit(pytest.main([__file__, "-x", "-vv", "-s"])) + multiprocess_main(__file__, worker_main) diff --git a/python/sglang/jit_kernel/tests/test_tp_qknorm.py b/python/sglang/jit_kernel/tests/test_tp_qknorm.py new file mode 100644 index 000000000..98a803ca0 --- /dev/null +++ b/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) diff --git a/python/sglang/jit_kernel/tests/utils.py b/python/sglang/jit_kernel/tests/utils.py new file mode 100644 index 000000000..56f3511ab --- /dev/null +++ b/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"])) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py index 2006554f9..11599552f 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py +++ b/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: - max_pull_size = 16 * 1024 * 1024 # default to 16MB - if max_push_size is None: - max_push_size = self.config.one_shot_push_threshold - max_push_size = min(max_push_size, max_pull_size) + if max_pull_size is None: # default to 16MB + max_pull_size = 16 * 1024 * 1024 + if max_push_size is None: # default to recommended size + config = THRESHOLD_2_SHOT_MAP[self.world_size] + max_push_size = config.one_shot_push_threshold 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 diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 1904d046a..eea60e586 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/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( - q_norm: "MiniMaxM2RMSNormTP", - k_norm: "MiniMaxM2RMSNormTP", - q: torch.Tensor, - k: torch.Tensor, - ) -> torch.Tensor: - sum_sq = rms_sumsq_serial(q, k) - if q_norm.attn_tp_size > 1: - sum_sq = attn_tp_all_reduce(sum_sq) - q, k = rms_apply_serial( - q, - k, - q_norm.weight, - k_norm.weight, - sum_sq, - q_norm.attn_tp_size, - q_norm.variance_epsilon, +@register_custom_op(mutates_args=["q", "k"]) +def fused_tp_qknorm( + counter: int, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, +) -> None: + return fused_parallel_qknorm( + MiniMaxM2QKRMSNorm.COMM_MAP[counter].obj, + q, + k, + q_weight, + k_weight, + eps=eps, + ) + + +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()) - # 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.qk_norm_impl.forward(q, k) q, k = self.rotary_emb(positions, q, k) inner_state = q, k, v, forward_batch return None, forward_batch, inner_state