diff --git a/python/sglang/kernels/jit/csrc/moe/route_radix4_hip.cuh b/python/sglang/kernels/jit/csrc/moe/route_radix4_hip.cuh new file mode 100644 index 000000000..acef54aae --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/route_radix4_hip.cuh @@ -0,0 +1,527 @@ +/// Four-bit radix-select router for K3 routing on CDNA (gfx942/gfx950). +/// +/// One block routes one token. Each thread keeps its slice of the 896 scores in +/// registers, and the block fixes four key bits of the pivot at a time: a round +/// tallies the still-live keys into a 16-bin histogram, accumulates from the top +/// bin down, and keeps the bin the k-th key falls in. The round count therefore +/// follows the key width rather than topk. +/// +/// A round's cost is dominated by the vector->scalar->vector unit trip that +/// broadcasting a cross-lane count requires, and that trip is the same price +/// whether the round resolves one bit or four, hence 16 bins. What is left is +/// the per-round work that does not shrink as the experts are spread over more +/// waves; transposing the bin totals onto lanes turns the 16-step walk over +/// them into a 4-step DPP prefix sum plus a ballot. +/// +/// Selection contract: experts rank by sigmoid(score) + bias, and a NaN ranking +/// value always ranks below every number, so it can never displace one. The +/// sigmoid itself uses aiter's approximate exp2f + rcpf combination, because +/// expf plus a divide would still disagree with aiter's result at the ULP +/// level. Experts whose key is exactly identical (a strict tie) are separated +/// by kAiterTieLaneRank below, which is exactly the order aiter's wave64 +/// traversal reaches them; +/// +/// Winners are emitted highest key first, equal keys in the same tie order. So +/// regardless of whether the input ties, this row matches what aiter would +/// produce, expert for expert and column for column -- this kernel and aiter +/// both serve K3, just split by batch size, and a routing that changed with +/// batch size would be the cost of the two disagreeing. The write race during +/// compaction does not affect the final output: it fills the staged row in +/// whatever order wins the race, and the epilogue re-ranks that row afterward. + +#pragma once + +#ifndef USE_ROCM +#error "route_radix4_hip.cuh targets CDNA; it uses amdgcn DPP and wave64 ballots" +#endif + +#include +#include + +#include +#include +#include + +#include + +#include + +namespace sglang { + +inline constexpr uint32_t kRadix4NumExperts = 896; +inline constexpr uint32_t kRadix4TopK = 16; +inline constexpr uint32_t kRadix4Block = 256; +inline constexpr uint32_t kRadix4Wave = 64; + +/// log2(e): aiter computes exp(-x) as exp2f(-kAiterSigmoidLog2E * x) rather than +/// expf(-x). Matched bit for bit with topk_softmax_kernels_group.cu's C_LOG2E. +inline constexpr float kAiterSigmoidLog2E = 1.44269504088896340736f; + +struct RouteRadix4Params { + const void* __restrict__ scores; + const void* __restrict__ bias; + fp32_t* __restrict__ out_w; + int32_t* __restrict__ out_i; + uint32_t stride_scores; + uint32_t stride_out; + fp32_t routed_scaling_factor; + bool renormalize; +}; + +namespace radix4 { + +/// Rank of a wave64 lane in aiter router's traversal order. Its tie positions +/// come from a cumulative sum over that traversal order, so two experts whose +/// ranking value is bit-for-bit equal are separated by their position in it. +/// This table was obtained by comparing item-by-item against aiter, and +/// test_moe_route_radix4 pins it back against aiter, so a change on their side +/// surfaces as a test failure rather than a silent divergence. +static __device__ __constant__ uint8_t kAiterTieLaneRank[64] = { + 56, 57, 58, 59, 63, 62, 61, 60, 52, 53, 54, 55, 51, 50, 49, 48, 40, 41, 42, 43, 47, 46, + 45, 44, 36, 37, 38, 39, 35, 34, 33, 32, 24, 25, 26, 27, 31, 30, 29, 28, 20, 21, 22, 23, + 19, 18, 17, 16, 8, 9, 10, 11, 15, 14, 13, 12, 4, 5, 6, 7, 3, 2, 1, 0, +}; + +/// Where an expert falls in that traversal order: a bijection onto [0, EXPERTS). +/// Experts come in groups of four, and a group lands on one lane; 224 groups +/// are cyclically assigned across 64 lanes, with lanes 0-31 each getting 4 +/// groups (banks 0-3) and lanes 32-63 each getting 3 groups (banks 0-2). The +/// rank < 32 branch below takes *3 (for the 3-bank lanes), otherwise *4 (for +/// the 4-bank lanes) -- this relies on the structural fact that +/// kAiterTieLaneRank happens to map lanes 0-31 to rank>=32 and lanes 32-63 to +/// rank<32, rather than branching on the bank count directly. +SGL_DEVICE uint32_t tie_priority(int expert) { + const int group = expert >> 2; + const int lane = group & 63; + const int bank = group >> 6; + const int rank = static_cast(kAiterTieLaneRank[lane]); + assert((rank < 32) == (lane >= 32)); + const int group_rank = (rank < 32) ? (rank * 3 + bank) : (96 + (rank - 32) * 4 + bank); + return static_cast((group_rank << 2) + (expert & 3)); +} + +/// How many of the bitmap's members come before p. Every thread runs the same +/// straight line over the same broadcast words, so the walk costs the block no +/// divergence, only NWORDS popcounts. +template +SGL_DEVICE int tie_rank(const uint64_t* bits, uint32_t p) { + const uint32_t w = p >> 6; + int n = 0; +#pragma unroll + for (uint32_t j = 0; j < NWORDS; ++j) { + const uint64_t below = (j < w) ? ~0ull : ((j == w) ? ((1ull << (p & 63)) - 1ull) : 0ull); + n += __popcll(bits[j] & below); + } + return n; +} + +SGL_DEVICE float load_score(const bf16_t* p, int i) { + return __uint_as_float(static_cast(reinterpret_cast(p)[i]) << 16); +} + +SGL_DEVICE float load_score(const fp32_t* p, int i) { + return p[i]; +} + +/// Monotonic float -> uint32 map, so unsigned compares order the floats. The map +/// never returns 0: a negative f gives ~u, which is 0 only for the all-ones NaN, +/// and a positive one has the sign bit set. Key 0 is therefore free to mean "NaN, +/// ranks below everything", -inf included. +SGL_DEVICE uint32_t sortable(float f) { + uint32_t u = __float_as_uint(f); + // Map -0.0 and +0.0 to the same value. + if (u == 0x80000000u) u = 0u; + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +/// The ranking key: 0 for a NaN, so the expert it belongs to can never displace +/// one whose ranking value is a number. +SGL_DEVICE uint32_t rank_key(float x) { + return (x == x) ? sortable(x) : 0u; +} + +template +SGL_DEVICE void dpp_add_stage(uint32_t (&x)[N]) { +#pragma unroll + for (int j = 0; j < N; ++j) + x[j] += static_cast(__builtin_amdgcn_update_dpp(0, static_cast(x[j]), CTRL, RM, BM, false)); +} + +/// Wave-wide inclusive prefix sum over N uint32 at a time: afterwards lane L is +/// the sum of lanes 0..L and lane 63 is the total. The narrower bank masks on the +/// row_shr:4 and row_shr:8 stages switch off exactly the lanes whose source lane +/// falls outside the row, which would have added zero, so masking them or not +/// makes no difference. +template +SGL_DEVICE void wave_sum_dpp(uint32_t (&x)[N]) { + dpp_add_stage<0x111, 0xf, 0xf>(x); // row_shr:1 + dpp_add_stage<0x112, 0xf, 0xf>(x); // row_shr:2 + dpp_add_stage<0x114, 0xf, 0xe>(x); // row_shr:4 + dpp_add_stage<0x118, 0xf, 0xc>(x); // row_shr:8 + dpp_add_stage<0x142, 0xa, 0xf>(x); // row_bcast:15 + dpp_add_stage<0x143, 0xc, 0xf>(x); // row_bcast:31 +} + +template +SGL_DEVICE float dpp_fadd_stage(float x) { + const int moved = __builtin_amdgcn_update_dpp(0, __builtin_bit_cast(int, x), CTRL, RM, BM, false); + return x + __builtin_bit_cast(float, moved); +} + +/// Sums v within the wave and leaves the total in out[wid]. The ladder fixes the +/// order the addition happens in, so the same values give the same float on two +/// runs. Not __shfl_xor: that turns into six ds_bpermute round trips through LDS, +/// measured slower. +SGL_DEVICE void stage_wave_sum(float v, int lane, int wid, float* out) { + v = dpp_fadd_stage<0x111, 0xf, 0xf>(v); // row_shr:1 + v = dpp_fadd_stage<0x112, 0xf, 0xf>(v); // row_shr:2 + v = dpp_fadd_stage<0x114, 0xf, 0xe>(v); // row_shr:4 + v = dpp_fadd_stage<0x118, 0xf, 0xc>(v); // row_shr:8 + v = dpp_fadd_stage<0x142, 0xa, 0xf>(v); // row_bcast:15 + v = dpp_fadd_stage<0x143, 0xc, 0xf>(v); // row_bcast:31 + if (lane == static_cast(kRadix4Wave) - 1) out[wid] = v; +} + +} // namespace radix4 + +template +__global__ __launch_bounds__(BLOCK) void route_radix4_kernel(__grid_constant__ const RouteRadix4Params params) { + constexpr int WAVE = static_cast(kRadix4Wave); + constexpr int NWAVE = BLOCK / WAVE; + constexpr int VPT = (EXPERTS + BLOCK - 1) / BLOCK; + // When the histogram below is tallied, each thread counts the bins of its VPT + // (values per thread, how many experts one thread handles) experts into one + // 64-bit register: 16 bins of 4 bits each, and an expert adds 1 to the 4 bits + // it belongs to. + constexpr int CHUNK = 15; + constexpr int NACC = (VPT + CHUNK - 1) / CHUNK; + static_assert(BLOCK % WAVE == 0, "block must be whole waves"); + static_assert(TOPK <= WAVE, "topk must fit in one lane-indexed row"); + static_assert(VPT <= 32, "alive mask is 32 bits"); + + const int token = blockIdx.x; + const int tid = threadIdx.x; + const int lane = tid % WAVE; + const int wid = tid / WAVE; + const auto* srow = static_cast(params.scores) + static_cast(token) * params.stride_scores; + const auto* sbias = static_cast(params.bias); + + float sig[VPT]; + uint32_t key[VPT]; + // BLOCK * VPT overshoots EXPERTS; hence the mask. + uint32_t valid = (VPT >= 32) ? 0xffffffffu : ((1u << VPT) - 1u); + + // Read the scores and the bias and compute sig[i] and key[i]: sig[i] is the + // sigmoid without the bias and is what gets emitted, key[i] is sigmoid + bias + // and is used to rank. Also accumulate the bitwise OR and AND of every key, + // needed at diff below. + uint32_t or_all = 0u, and_all = 0xffffffffu; +#pragma unroll + for (int i = 0; i < VPT; ++i) { + const int e = tid + i * BLOCK; + if (e < EXPERTS) { + const float x = radix4::load_score(srow, e); + const float g = __builtin_amdgcn_rcpf(1.0f + exp2f(-kAiterSigmoidLog2E * x)); + sig[i] = g; + key[i] = radix4::rank_key(g + radix4::load_score(sbias, e)); + or_all |= key[i]; + and_all &= key[i]; + } else { + sig[i] = 0.0f; + key[i] = 0u; + valid &= ~(1u << i); + } + } + uint32_t alive = valid; + + // One bit per expert, indexed by tie priority rather than by id. + constexpr int TIE_WORDS = (EXPERTS + 63) / 64; + + __shared__ uint32_t s_hist[2][NWAVE][8]; + __shared__ uint32_t s_pre[2][NWAVE]; + __shared__ uint64_t s_tie[TIE_WORDS]; + __shared__ float s_w[TOPK]; + __shared__ float s_wsum[NWAVE]; + __shared__ int s_id[TOPK]; + __shared__ uint32_t s_key[TOPK]; + __shared__ int s_cnt; + +#pragma unroll + for (int s = 32; s > 0; s >>= 1) { + or_all |= __shfl_xor(or_all, s, WAVE); + and_all &= __shfl_xor(and_all, s, WAVE); + } + if (lane == 0) { + s_pre[0][wid] = or_all; + s_pre[1][wid] = and_all; + } + __syncthreads(); +#pragma unroll + for (int w = 0; w < NWAVE; ++w) { + or_all |= s_pre[0][w]; + and_all &= s_pre[1][w]; + } + + // The XOR marks the bits the keys disagree on (a disagreeing bit is 1 in the + // OR and 0 in the AND); the highest one is where the search starts. Above it + // every key is the same, and that value comes out of and_all as the initial + // pivot. + const uint32_t diff = or_all ^ and_all; + const int start = diff ? (31 - __clz(diff)) : -1; + uint32_t pivot = (start >= 31) ? 0u : (and_all & ~((1u << (start + 1)) - 1u)); + + int need = TOPK; // how many still have to be picked out of the live set + // lowest bit the pivot is resolved down to + int bend = 0; + bool capped = (start < 0); + + // Main loop: fixes the TOPK-th largest key (the pivot) without sorting. A round + // takes 4 bits as the bin index, tallies the histogram and accumulates from the + // top bin down; the bin the need-th key falls in fixes those 4 pivot bits, the + // higher bins are in for good and come off need, and alive narrows to that bin. + // Survivors exactly equal to need exit early; still more than need at the + // lowest bit sets capped. +#pragma unroll 1 + for (int b = (start < 0) ? -4 : (start >> 2) << 2; b >= 0; b -= 4) { + // Two buffers used alternately, so a round's writes only have to wait for the + // reads of the round before last. + const int buf = (b >> 2) & 1; + + uint64_t h[NACC]; +#pragma unroll + for (int a = 0; a < NACC; ++a) + h[a] = 0ull; +#pragma unroll + for (int i = 0; i < VPT; ++i) { + const uint32_t bin = (key[i] >> b) & 15u; + h[i / CHUNK] += static_cast((alive >> i) & 1u) << (4 * bin); + } + + // Spread the 16 four-bit counts over 8 uint32, two bins each. + uint32_t p[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + uint32_t lo = 0u, hi = 0u; +#pragma unroll + for (int a = 0; a < NACC; ++a) { + lo += static_cast((h[a] >> (8 * j)) & 0xfull); + hi += static_cast((h[a] >> (8 * j + 4)) & 0xfull); + } + p[j] = lo | (hi << 16); + } + radix4::wave_sum_dpp(p); + if (lane == WAVE - 1) { +#pragma unroll + for (int j = 0; j < 8; ++j) + s_hist[buf][wid][j] = p[j]; + } + __syncthreads(); + + // One LDS read puts each of the 16 bins on a lane of its own, which turns the + // walk over the bins into a prefix sum across lanes. Lane L takes bin 15 - L: + // with the bins reversed, "how many keys are in this bin or a higher one" is + // exactly the direction row_shr adds in. + int c = 0; + if (lane < 16) { + const int d = 15 - lane; +#pragma unroll + for (int w = 0; w < NWAVE; ++w) + c += static_cast((s_hist[buf][w][d >> 1] >> ((d & 1) * 16)) & 0xffffu); + } + + // A prefix sum of c across lanes, so cum is the number of keys in this lane's + // bin and in every higher bin. + int cum = c; + cum += __builtin_amdgcn_update_dpp(0, cum, 0x111, 0xf, 0xf, false); // row_shr:1 + cum += __builtin_amdgcn_update_dpp(0, cum, 0x112, 0xf, 0xf, false); // row_shr:2 + cum += __builtin_amdgcn_update_dpp(0, cum, 0x114, 0xf, 0xf, false); // row_shr:4 + cum += __builtin_amdgcn_update_dpp(0, cum, 0x118, 0xf, 0xf, false); // row_shr:8 + + // cum rises monotonically with L, so the predicate turns on once and stays + // on; its lowest lane is the bin holding the k-th key. + const unsigned long long mk = __ballot(lane < 16 && cum >= need); + const int L0 = __ffsll(mk) - 1; + const int sel = 15 - L0; + const int nsel = __builtin_amdgcn_readlane(c, L0); + const int above = __builtin_amdgcn_readlane(cum, L0) - nsel; + + need -= above; + pivot |= static_cast(sel) << b; + bend = b; +#pragma unroll + for (int i = 0; i < VPT; ++i) + if (((key[i] >> b) & 15u) != static_cast(sel)) alive &= ~(1u << i); + + // Every survivor is a winner: the remaining bits cannot change the set. The + // test is uniform across the block, since every thread scanned the same LDS + // totals. + if (nsel == need) break; + if (b == 0) capped = true; + } + + const uint32_t pmask = ~((1u << bend) - 1u); + // The renorm divisor is accumulated while the winners are picked instead of + // read back off the staged row: the order that row gets filled in can differ + // from run to run, and the order of the float additions with it, whereas + // reducing across threads always adds in lane order. + float wsum = 0.0f; + if (!capped) { + // The survivors exactly fill the quota, so a key wins as soon as it reaches + // the pivot prefix. Winners span waves, so ballot cannot number them; an LDS + // bump counter can. + if (tid == 0) s_cnt = 0; + __syncthreads(); +#pragma unroll + for (int i = 0; i < VPT; ++i) { + if (((valid >> i) & 1u) && (key[i] & pmask) >= pivot) { + const int pos = atomicAdd(&s_cnt, 1); + if (pos < TOPK) { + s_w[pos] = sig[i]; + s_id[pos] = tid + i * BLOCK; + s_key[pos] = key[i]; + wsum += sig[i]; + } + } + } + radix4::stage_wave_sum(wsum, lane, wid, s_wsum); + __syncthreads(); + } else { + // Every pivot bit is fixed and the survivors still outnumber the quota, so + // what is left are keys equal bit for bit and only the tie rule separates + // them: `need` of them get in, the ones aiter's traversal reaches first. + // Marking the survivors in a bitmap indexed by that traversal turns "how + // many come before me" into a popcount, which costs the block no + // divergence and no scan. + for (int j = tid; j < TIE_WORDS; j += BLOCK) + s_tie[j] = 0ull; + if (tid == 0) s_cnt = 0; + __syncthreads(); + + uint32_t eqm = 0u; +#pragma unroll + for (int i = 0; i < VPT; ++i) { + if (((valid >> i) & 1u) && (key[i] & pmask) == pivot) { + eqm |= 1u << i; + const uint32_t p = radix4::tie_priority(tid + i * BLOCK); + atomicOr(&s_tie[p >> 6], 1ull << (p & 63)); + } + } + __syncthreads(); + + // The outright winners are TOPK - need of them, so the ties owe exactly the + // `need` the loop stopped short of and the row comes out full. +#pragma unroll + for (int i = 0; i < VPT; ++i) { + if (!((valid >> i) & 1u)) continue; + const int e = tid + i * BLOCK; + const bool tied = ((eqm >> i) & 1u) != 0u; + const bool taken = + (key[i] & pmask) > pivot || (tied && radix4::tie_rank(s_tie, radix4::tie_priority(e)) < need); + if (taken) { + const int pos = atomicAdd(&s_cnt, 1); + if (pos < TOPK) { + s_w[pos] = sig[i]; + s_id[pos] = e; + s_key[pos] = key[i]; + wsum += sig[i]; + } + } + } + radix4::stage_wave_sum(wsum, lane, wid, s_wsum); + __syncthreads(); + } + + if (tid < TOPK) { + // Where the compaction put a winner is a race; where it is emitted is not. + // Tie priorities are distinct, so (key desc, priority asc) is a strict total + // order and counting the winners that outrank this one lands each of them on + // a position of its own. The whole staged row sits in lanes 0..TOPK-1 of this + // wave, so the walk over it reads the other lanes' registers instead of LDS: + // readlane is a scalar op and the lane index is a constant of the unrolled + // loop, which leaves the vector unit with just the compares. + const uint32_t k = s_key[tid]; + const int id = s_id[tid]; + const auto p = static_cast(radix4::tie_priority(id)); + int rank = 0; +#pragma unroll + for (int q = 0; q < TOPK; ++q) { + const auto kq = static_cast(__builtin_amdgcn_readlane(static_cast(k), q)); + const int pq = __builtin_amdgcn_readlane(p, q); + rank += (kq > k || (kq == k && pq < p)) ? 1 : 0; + } + + float scale = params.routed_scaling_factor; + if (params.renormalize) { + float sum = 0.0f; +#pragma unroll + for (int w = 0; w < NWAVE; ++w) + sum += s_wsum[w]; + // Every sigmoid underflows to zero on a row of saturated scores, and a row + // of NaN sums to NaN; neither may turn a finite weight into an inf. + scale /= (sum > 0.0f) ? sum : 1.0f; + } + const size_t o = static_cast(token) * params.stride_out + rank; + params.out_w[o] = s_w[tid] * scale; + params.out_i[o] = id; + } +} + +struct RouteRadix4Kernel { + static void + run(const tvm::ffi::TensorView scores, + const tvm::ffi::TensorView bias, + const tvm::ffi::TensorView out_w, + const tvm::ffi::TensorView out_i, + int64_t topk, + double routed_scaling_factor, + bool renormalize) { + using namespace host; + + auto M_ = SymbolicSize{"num_tokens"}; + auto N_ = SymbolicSize{"num_experts"}; + auto K_ = SymbolicSize{"topk"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + auto score_dtype = SymbolicDType{}; + TensorMatcher({M_, N_}) + .with_dtype(score_dtype) + .with_device(device_) + .with_strides({-1, 1}) + .verify(scores); + // Rebinding the same symbolic dtype makes the bias track the scores. + TensorMatcher({N_}).with_dtype(score_dtype).with_device(device_).verify(bias); + TensorMatcher({M_, K_}).with_dtype().with_device(device_).verify(out_w); + TensorMatcher({M_, K_}).with_dtype().with_device(device_).verify(out_i); + + RuntimeCheck( + N_.unwrap() == kRadix4NumExperts && K_.unwrap() == kRadix4TopK && topk == kRadix4TopK, + "route_radix4 is specialized for N=896, K=16"); + + const auto M = static_cast(M_.unwrap()); + if (M == 0) return; + + const auto params = RouteRadix4Params{ + .scores = scores.data_ptr(), + .bias = bias.data_ptr(), + .out_w = static_cast(out_w.data_ptr()), + .out_i = static_cast(out_i.data_ptr()), + .stride_scores = static_cast(scores.stride(0)), + .stride_out = static_cast(out_w.stride(0)), + .routed_scaling_factor = static_cast(routed_scaling_factor), + .renormalize = renormalize, + }; + + constexpr auto kExperts = static_cast(kRadix4NumExperts); + constexpr auto kTopK = static_cast(kRadix4TopK); + constexpr auto kBlock = static_cast(kRadix4Block); + const auto device = device_.unwrap(); + if (score_dtype.is_type()) { + LaunchKernel(M, kBlock, device)(route_radix4_kernel, params); + } else { + LaunchKernel(M, kBlock, device)(route_radix4_kernel, params); + } + } +}; + +} // namespace sglang diff --git a/python/sglang/kernels/ops/moe/moe_route_radix4.py b/python/sglang/kernels/ops/moe/moe_route_radix4.py new file mode 100644 index 000000000..ff34717a1 --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_route_radix4.py @@ -0,0 +1,137 @@ +"""Four-bit radix-select router for K3 routing on CDNA (ROCm). + +The ROCm counterpart to moe_route_radix, dispatched from +biased_grouped_topk_gpu's aiter branch for covered inputs; anything else falls +back to aiter. Kernel notes live in jit/csrc/moe/route_radix4_hip.cuh. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit +from sglang.kernels.jit.utils.common import is_hip_runtime + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +_NUM_EXPERTS = 896 +_TOPK = 16 +# One block per token, so the grid outgrows the machine somewhere past a +# thousand tokens and the kernel turns throughput-bound, where spreading a token +# over four waves is a cost rather than a win. Measured break-even is ~1.5k +# tokens; below 1k the kernel still leads by 1.2x or more, and prefill-sized +# batches are far above either number. +_MAX_TOKENS = 1024 + +logger = logging.getLogger(__name__) + + +def supported_hardware() -> bool: + """Whether this device is one the kernel targets, before asking whether it + builds. The kernel is wave64 and GFX9 DPP throughout, hence gfx942/gfx950.""" + if not is_hip_runtime() or not torch.cuda.is_available(): + return False + gcn_arch = torch.cuda.get_device_properties(0).gcnArchName + return any(arch in gcn_arch for arch in ("gfx942", "gfx950")) + + +@cache_once +def build() -> Module: + """Compile and load the kernel, raising if the toolchain cannot.""" + return load_jit( + "moe_route_radix4", + cuda_files=["moe/route_radix4_hip.cuh"], + cuda_wrappers=[("run", "RouteRadix4Kernel::run")], + # No fast-math: expert-id selection must stay comparable to aiter under + # ties and NaN. + extra_cuda_cflags=["-O3"], + ) + + +@cache_once +def available() -> bool: + """Whether dispatch may use the kernel: targeted hardware, and a kernel that + builds on this toolchain. + + A build failure is swallowed on purpose, since serving would rather fall back + to aiter than refuse to start. That makes this the wrong gate for a test, + which wants a kernel that stopped compiling to be a failure and not a skip -- + the tests pair supported_hardware() with build() instead. + """ + if not supported_hardware(): + return False + try: + build() + return True + except Exception as e: # pragma: no cover - toolchain dependent + logger.warning(f"Failed to load the JIT ROCm radix router: {e}") + return False + + +def covered( + scores: torch.Tensor, + bias: torch.Tensor, + topk: int, + num_expert_group: Optional[int], + topk_group: Optional[int], +) -> bool: + """Specialized for K3 routing: [M, 896] row-contiguous scores, top-16, + ungrouped, with the bias in the score dtype (what the aiter path feeds it). + + Grouped routing is excluded rather than emulated: the kernel ranks all 896 + experts at once and has no notion of masking whole groups out first. + """ + return ( + scores.dim() == 2 + and scores.size(0) <= _MAX_TOKENS + and scores.size(1) == _NUM_EXPERTS + and int(topk) == _TOPK + and scores.dtype in (torch.bfloat16, torch.float32) + and bias.dtype == scores.dtype + and scores.stride(1) == 1 + and bias.is_contiguous() + and (num_expert_group or 1) == 1 + and (topk_group or 1) == 1 + ) + + +def route_radix4( + scores: torch.Tensor, + bias: torch.Tensor, + topk: int, + renormalize: bool, + routed_scaling_factor: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Returns (weights [M, topk] fp32, ids [M, topk] int32). Caller must have + checked covered(). + + Experts are ranked by sigmoid(score) + bias but the emitted weight is the + plain sigmoid, scaled by routed_scaling_factor and, when renormalize is set, + divided by the sum over the selected experts. A NaN ranking value keys below + every number, so it can never displace one. + + A row is what aiter would have produced, expert for expert and column for + column: winners come out highest ranking value first, and experts that tie on + the full ranking value are separated the way aiter's wave64 walk reaches them + (kAiterTieLaneRank in the kernel). Matching on tied rows too is what keeps a + token's routing from depending on the batch size, since batches past + _MAX_TOKENS fall back to aiter. None of it depends on how the kernel's + compaction raced, so a row also repeats bit for bit run to run. + """ + M = scores.shape[0] + out_w = torch.empty((M, topk), dtype=torch.float32, device=scores.device) + out_i = torch.empty((M, topk), dtype=torch.int32, device=scores.device) + build().run( + scores, + bias, + out_w, + out_i, + topk, + float(routed_scaling_factor), + bool(renormalize), + ) + return out_w, out_i diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index db01fce13..ac34b2c93 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1486,6 +1486,8 @@ class Envs: # front reads hidden_states once, and run the top-k plus the bf16 cast in one # epilogue kernel. See kernels/ops/moe/moe_front.py. Default on. SGLANG_K3_FUSED_FRONT = EnvBool(True) + # Use the ROCm radix-4 router for covered K3 top-k workloads. + SGLANG_K3_RADIX4_TOPK = EnvBool(False) SGLANG_KIMI_K3_VIT_CUDA_GRAPH_CACHE_CAPACITY = EnvInt(2) SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MIN_HITS = EnvInt(2) SGLANG_KIMI_K3_VIT_CUDA_GRAPH_MAX_SEQLEN = EnvInt(6144) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index b9393e388..5fd00a5c0 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1575,17 +1575,33 @@ def biased_grouped_topk_gpu( assert ( hidden_states.shape[0] == gating_output.shape[0] ), f"Number of tokens mismatch: hidden_states.shape[0] = {hidden_states.shape[0]}, gating_output.shape[0] = {gating_output.shape[0]}" + bias = correction_bias.to(dtype=gating_output.dtype) + scaling = routed_scaling_factor if routed_scaling_factor is not None else 1.0 + + if envs.SGLANG_K3_RADIX4_TOPK.get(): + from sglang.kernels.ops.moe import moe_route_radix4 + + # Gated on the routing shape. Kimi-K3 (896 experts, top-16, + # ungrouped) is the only config covered for now; anything else + # falls back to aiter. + if moe_route_radix4.available() and moe_route_radix4.covered( + gating_output, bias, topk, num_expert_group, topk_group + ): + return moe_route_radix4.route_radix4( + gating_output, bias, topk, renormalize, scaling + ) + topk_weights = torch.empty((token, topk), dtype=torch.float32, device=device) topk_ids = torch.empty((token, topk), dtype=torch.int32, device=device) aiter_biased_grouped_topk( gating_output, - correction_bias.to(dtype=gating_output.dtype), + bias, topk_weights, topk_ids, num_expert_group, topk_group, renormalize, - routed_scaling_factor if routed_scaling_factor is not None else 1.0, + scaling, ) return topk_weights, topk_ids elif _is_musa and ( diff --git a/test/registered/kernels/ops/moe/test_moe_route_radix4.py b/test/registered/kernels/ops/moe/test_moe_route_radix4.py new file mode 100644 index 000000000..a424b8fdc --- /dev/null +++ b/test/registered/kernels/ops/moe/test_moe_route_radix4.py @@ -0,0 +1,307 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the ROCm radix-4 K3 router. + +Two references, because the kernel owes two different things. A pure-torch fp32 +oracle spells out the contract -- what the weights are, where a NaN ranks, how a +row is ordered -- and covers the shapes and edge cases. aiter is the router this +one stands in for on decode-sized batches while larger ones still go to it, so +the test_*_aiter cases hold the kernel to matching it column for column, ties +included; anything less and a token's routing would depend on the batch size. + +Run: pytest test/registered/kernels/ops/moe/test_moe_route_radix4.py -v +""" + +import sys + +import pytest +import torch + +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_amd_ci + +if not is_hip(): + pytest.skip("The radix-4 router is the ROCm path.", allow_module_level=True) +if not torch.cuda.is_available(): + pytest.skip("Requires a GPU.", allow_module_level=True) + +from sglang.kernels.ops.moe import moe_route_radix4 + +if not moe_route_radix4.supported_hardware(): + pytest.skip("The kernel targets gfx942/gfx950.", allow_module_level=True) + +# Deliberately not gated on available(): that reports a kernel which failed to +# build as merely unavailable, so serving can fall back to aiter, and a file that +# skipped on it would report a kernel that stopped compiling as a green run. On +# hardware the kernel targets it has to build, so build it here and let the +# toolchain error reach the report. +moe_route_radix4.build() + +try: + from aiter import biased_grouped_topk as aiter_biased_grouped_topk +except ImportError: + aiter_biased_grouped_topk = None + +register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd-mi35x") + +NUM_EXPERTS = 896 +TOPK = 16 +# The kernel keys a NaN to 0, which is below every number including -inf, so -inf +# stands in for it here. The two part company only on a row that ranks an expert +# at -inf as well, which needs an infinite bias and is outside the contract. +NAN_RANK = float("-inf") + +# kAiterTieLaneRank from route_radix4_hip.cuh. Kept as a second copy on purpose: +# test_route_radix4_tie_order_is_aiters checks it against aiter directly, so a +# change on aiter's side fails here instead of drifting silently in the kernel. +_TIE_LANE_RANK = [ + 56, 57, 58, 59, 63, 62, 61, 60, 52, 53, 54, 55, 51, 50, 49, 48, + 40, 41, 42, 43, 47, 46, 45, 44, 36, 37, 38, 39, 35, 34, 33, 32, + 24, 25, 26, 27, 31, 30, 29, 28, 20, 21, 22, 23, 19, 18, 17, 16, + 8, 9, 10, 11, 15, 14, 13, 12, 4, 5, 6, 7, 3, 2, 1, 0, +] # fmt: skip + + +def _tie_priority(expert): + """Where aiter's wave64 walk reaches an expert; a bijection onto [0, 896).""" + group = expert >> 2 + lane, bank = group & 63, group >> 6 + rank = _TIE_LANE_RANK[lane] + group_rank = rank * 3 + bank if rank < 32 else 96 + (rank - 32) * 4 + bank + return (group_rank << 2) + (expert & 3) + + +TIE_PRIORITY = [_tie_priority(e) for e in range(NUM_EXPERTS)] +assert len(set(TIE_PRIORITY)) == NUM_EXPERTS, "tie priority is not a permutation" +# Experts in tie order, so a stable descending sort over these columns breaks ties +# the way the kernel does. +_BY_PRIORITY = torch.tensor( + sorted(range(NUM_EXPERTS), key=TIE_PRIORITY.__getitem__), device="cuda" +) + + +def _oracle(scores, bias, renormalize, scaling): + """Contract, from route_radix4_hip.cuh: bias ranks only and the emitted + weight stays bias-free, a NaN ranks below every number so it can never win, + ties follow aiter's walk, renormalize divides by the winners' sum (guarded to + 1 when that sum is non-positive) before scaling. Winners are emitted highest + ranking value first, equal values in that same tie order.""" + s = torch.sigmoid(scores.float()) + biased = s + bias.float() + biased = torch.where(torch.isnan(biased), torch.full_like(biased, NAN_RANK), biased) + # Reorder the columns into tie order first, so that a stable descending sort + # leaves equal values in it, then map the winners back to expert ids. + by_priority = biased[:, _BY_PRIORITY] + picked = torch.argsort(by_priority, dim=-1, descending=True, stable=True)[:, :TOPK] + ranked = _BY_PRIORITY[picked] + w = s.gather(1, ranked) + if renormalize: + total = w.sum(-1, keepdim=True) + w = w / torch.where(total > 0, total, torch.ones_like(total)) + return w * scaling, ranked.to(torch.int32) + + +def _assert_matches_oracle(scores, bias, renormalize=True, scaling=2.5): + ref_w, ref_ids = _oracle(scores, bias, renormalize, scaling) + w, ids = moe_route_radix4.route_radix4(scores, bias, TOPK, renormalize, scaling) + # Column by column: the position a winner lands in is part of the contract. + assert torch.equal(ids, ref_ids) + # The kernel's sigmoid is an approximate hardware sequence (matched to + # aiter's), whose last bits differ from torch's exact sigmoid. + torch.testing.assert_close(w, ref_w, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("m", [1, 3, 32, 255, 256, 512, 1024]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_route_radix4_matches_oracle(m, dtype, renormalize): + generator = torch.Generator(device="cuda").manual_seed(1000 + m) + # A padded row, so the kernel is exercised on a non-contiguous stride too. + backing = torch.randn( + (m, NUM_EXPERTS + 37), dtype=dtype, device="cuda", generator=generator + ) + scores = backing[:, 19 : 19 + NUM_EXPERTS] + bias = torch.randn(NUM_EXPERTS, dtype=dtype, device="cuda", generator=generator) + assert scores.stride(0) == NUM_EXPERTS + 37 + _assert_matches_oracle(scores, bias, renormalize) + + +@pytest.mark.parametrize("renormalize", [False, True]) +def test_route_radix4_ties(renormalize): + """Only the tie rule decides these, so an arrival-order compaction or a + mis-ranked tie shows up immediately.""" + bias = torch.zeros(NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + + # Every key identical: the search never runs a round and the whole field is + # a tie, which is the case an early-exit shortcut is most likely to miss. + _assert_matches_oracle( + torch.zeros(4, NUM_EXPERTS, dtype=torch.bfloat16, device="cuda"), + bias, + renormalize, + ) + + plateau = torch.full((4, NUM_EXPERTS), 0.25, dtype=torch.bfloat16, device="cuda") + plateau[:, 7] = 2.0 + plateau[:, 300] = 2.0 + plateau[:, 800] = 1.5 + _assert_matches_oracle(plateau, bias, renormalize) + + # A tied set too small to fill the quota on its own, spread across the + # block's waves and across each thread's register slots. + for seed in range(4): + generator = torch.Generator().manual_seed(seed) + picks = torch.randperm(NUM_EXPERTS, generator=generator)[:40] + scores = torch.full((1, NUM_EXPERTS), -4.0, dtype=torch.bfloat16, device="cuda") + scores[0, picks] = 4.0 + _assert_matches_oracle(scores, bias, renormalize) + + +def test_route_radix4_nan_never_wins(): + bias = torch.zeros(NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + scores = torch.randn(4, NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + scores[:, 100] = float("nan") + scores[:, 500] = float("nan") + # NaN sorts above every finite key under a raw monotone bit map, so without + # the floor these two would be picked before the real winner. + scores[:, 101] = 5.0 + ids = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5)[1] + assert not bool(((ids == 100) | (ids == 500)).any()) + _assert_matches_oracle(scores, bias) + + +def test_route_radix4_extremes(): + """Saturated sigmoids: the low key bits carry no information, so the search + has to lean on the shared-prefix skip and on the tie rule.""" + extreme = torch.linspace( + -90, 90, NUM_EXPERTS, dtype=torch.float32, device="cuda" + ).repeat(4, 1) + for dtype in (torch.bfloat16, torch.float32): + bias = torch.zeros(NUM_EXPERTS, dtype=dtype, device="cuda") + _assert_matches_oracle(extreme.to(dtype), bias) + + +def test_route_radix4_saturated_row(): + """Every sigmoid underflows to zero, so the renorm divisor is zero and only + the guard keeps the row from coming back as inf or NaN.""" + bias = torch.zeros(NUM_EXPERTS, dtype=torch.float32, device="cuda") + scores = torch.full((2, NUM_EXPERTS), -200.0, dtype=torch.float32, device="cuda") + scores[0, :16] = -120.0 + w, _ = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5) + assert torch.equal(w, torch.zeros_like(w)) + _assert_matches_oracle(scores, bias) + + +def test_route_radix4_reproducible(): + """Nothing the compaction races over reaches the output: the winner set is + settled by rank, the emitted order is settled by rank in the epilogue, and + the renorm divisor is reduced over threads. So a whole row repeats exactly, + column positions included.""" + scores = torch.randn(512, NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + bias = torch.zeros(NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + + first_w, first_ids = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5) + for _ in range(16): + w, ids = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5) + assert torch.equal(ids, first_ids) + assert torch.equal(w, first_w) + + +def _aiter_route(scores, bias, renormalize=True, scaling=2.5): + w = torch.empty((scores.shape[0], TOPK), dtype=torch.float32, device="cuda") + ids = torch.empty((scores.shape[0], TOPK), dtype=torch.int32, device="cuda") + aiter_biased_grouped_topk(scores, bias, w, ids, 1, 1, renormalize, scaling) + return w, ids + + +requires_aiter = pytest.mark.skipif( + aiter_biased_grouped_topk is None, reason="aiter is not installed" +) + + +@requires_aiter +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_route_radix4_matches_aiter(dtype, renormalize): + """The kernel stands in for the aiter router on decode-sized batches while + larger ones keep going to aiter, so a row has to agree with it column by + column, ties included -- otherwise a token's routing would depend on the batch + size. bf16 quantization makes exact ties common enough that plain random + scores reach them.""" + generator = torch.Generator(device="cuda").manual_seed(7) + scores = torch.randn( + 1024, NUM_EXPERTS, dtype=dtype, device="cuda", generator=generator + ) + bias = torch.randn(NUM_EXPERTS, dtype=dtype, device="cuda", generator=generator) + + ref_w, ref_ids = _aiter_route(scores, bias, renormalize) + w, ids = moe_route_radix4.route_radix4(scores, bias, TOPK, renormalize, 2.5) + assert torch.equal(ids, ref_ids) + torch.testing.assert_close(w, ref_w, rtol=1e-5, atol=1e-6) + + +@requires_aiter +@pytest.mark.parametrize("pool", [17, 24, 40, 128, 600, NUM_EXPERTS]) +def test_route_radix4_tie_order_is_aiters(pool): + """kAiterTieLaneRank is read off aiter rather than derived, so pin it: rows + whose whole top-k is settled by the tie rule, at pool sizes that put the cutoff + inside the tied set. A wrong entry either moves an expert across that cutoff or + swaps two columns, and both land here.""" + bias = torch.zeros(NUM_EXPERTS, dtype=torch.float32, device="cuda") + generator = torch.Generator().manual_seed(pool) + scores = torch.full((4, NUM_EXPERTS), -4.0, dtype=torch.float32, device="cuda") + for row in range(4): + scores[row, torch.randperm(NUM_EXPERTS, generator=generator)[:pool]] = 4.0 + + ref_w, ref_ids = _aiter_route(scores, bias) + w, ids = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5) + assert torch.equal(ids, ref_ids) + torch.testing.assert_close(w, ref_w, rtol=1e-5, atol=1e-6) + + +@requires_aiter +def test_route_radix4_tie_straddling_the_cutoff_matches_aiter(): + """Unambiguous winners take the first columns and a tied set fights over what + is left, so the tie rule settles both which experts get in and where they go.""" + bias = torch.zeros(NUM_EXPERTS, dtype=torch.float32, device="cuda") + scores = torch.full((8, NUM_EXPERTS), -5.0, dtype=torch.float32, device="cuda") + generator = torch.Generator().manual_seed(11) + for row in range(8): + perm = torch.randperm(NUM_EXPERTS, generator=generator) + for j, e in enumerate(perm[:10].tolist()): + scores[row, e] = 8.0 - j * 0.25 + scores[row, perm[10:30]] = 1.0 + + ref_w, ref_ids = _aiter_route(scores, bias) + w, ids = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5) + assert torch.equal(ids, ref_ids) + torch.testing.assert_close(w, ref_w, rtol=1e-5, atol=1e-6) + + +def test_route_radix4_graph_replay(): + scores = torch.randn(32, NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + bias = torch.randn(NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + ref_w, ref_ids = _oracle(scores, bias, True, 2.5) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + w, ids = moe_route_radix4.route_radix4(scores, bias, TOPK, True, 2.5) + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(ids, ref_ids) + torch.testing.assert_close(w, ref_w, rtol=1e-5, atol=1e-6) + + +def test_route_radix4_coverage(): + scores = torch.empty(32, NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + bias = torch.empty(NUM_EXPERTS, dtype=torch.bfloat16, device="cuda") + assert moe_route_radix4.covered(scores, bias, TOPK, 1, 1) + assert not moe_route_radix4.covered(scores, bias, TOPK, 8, 4) + assert not moe_route_radix4.covered(scores, bias, TOPK - 1, 1, 1) + assert not moe_route_radix4.covered(scores[:, :-1], bias[:-1], TOPK, 1, 1) + assert not moe_route_radix4.covered( + scores.new_empty((1536, NUM_EXPERTS)), bias, TOPK, 1, 1 + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))