dsv4.1: communication kernels and wrappers (#39653)
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: DarkSharpness <2040703891@qq.com> Co-authored-by: DarkSharpness <76582120+DarkSharpness@users.noreply.github.com> Co-authored-by: Ke Bao <ispobaoke@gmail.com> Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: Xiaoyu Zhang <xiaoyu.zhang@radixark.ai> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Co-authored-by: Yuwei An <ayw.sirius19@gmail.com> Co-authored-by: Zhichen Zeng <zczeng@uw.edu> Co-authored-by: Ziyi Xu <ziyi.xu@radixark.ai>
This commit is contained in:
Liangsheng Yin
co-authored by
Cheng Wan
Claude Opus 5
Cursor
DarkSharpness
DarkSharpness
Ke Bao
Khoa Pham
Xiaoyu Zhang
Xiaoyu Zhang
Yuhao Yang
Yuwei An
Zhichen Zeng
Ziyi Xu
Cheng Wan
Claude Opus 5
Cursor
DarkSharpness
DarkSharpness
Ke Bao
Khoa Pham
Xiaoyu Zhang
Xiaoyu Zhang
Yuhao Yang
Yuwei An
Zhichen Zeng
Ziyi Xu
parent
5a0c1e21e9
commit
7d5696b3a1
@@ -0,0 +1,737 @@
|
||||
// Fused deferred-MoE finalize -> 1shot lamport push all-reduce [-> RMSNorm]
|
||||
// over the CustomAllReduceV2 push plane, for decode-sized batches (bf16). The
|
||||
// hidden width, top_k and cluster geometry are template parameters; the
|
||||
// shared-expert add and the RMSNorm epilogue are optional.
|
||||
//
|
||||
// `idx == -1` marks a dropped slot (EP: the token was routed to an expert that
|
||||
// is not local) and contributes nothing. Accumulation is fp32 and the bf16
|
||||
// rounding points are exactly the unfused path's (moe_runner/flashinfer_trtllm.py
|
||||
// finalize -> `shared.add_(routed)` -> fp32-accumulating bf16 all-reduce in rank
|
||||
// order), so the kNorm=false result is bit-identical to it.
|
||||
//
|
||||
// The rank-local finalize never materializes in global memory: each thread
|
||||
// computes one 16B vector of it and pushes it straight into every peer's push
|
||||
// slot with unicast `st.relaxed.sys` stores, so no multicast mapping is needed.
|
||||
//
|
||||
// Push-plane protocol (see include/sgl_kernel/distributed/communicator.cuh):
|
||||
// * every rank owns 2 phases x kWorldSize slots of `slot_bytes`; a round
|
||||
// uses phase `counter & 1`, producer r writes slot r of every peer, the
|
||||
// consumer polls its own kWorldSize slots until no +0.0 marker remains,
|
||||
// reduces, and restores the +0.0 markers before it exits;
|
||||
// * +0.0 payload words are remapped to -0.0 (numerically identical) so a
|
||||
// written word is never 0 and `word == 0` means "not arrived yet";
|
||||
// * the generic push kernel owns one phase counter per block; this kernel
|
||||
// uses one per row cluster (flipped by the cluster's leader block after a
|
||||
// cluster barrier) plus a trailing "bumper" cluster that flips every
|
||||
// remaining one, so the whole array keeps one parity and both kernel
|
||||
// families can share the plane;
|
||||
// * every rank must call with the same num_tokens / hidden / top_k / epilogue:
|
||||
// slots are addressed by 16B vector index of the [T, hidden] row view.
|
||||
#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/communicator.cuh>
|
||||
#include <sgl_kernel/distributed/ptx.cuh>
|
||||
|
||||
#include <tvm/ffi/extra/stl.h>
|
||||
|
||||
#include <cooperative_groups.h>
|
||||
#include <cstdint>
|
||||
#include <optional>
|
||||
#include <type_traits>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
using device::distributed::PushWorkSpace;
|
||||
using host::distributed::CommunicatorRef;
|
||||
|
||||
/// One 16B staging vector (8 bf16) viewed as the 4 u32 words the lamport marker
|
||||
/// protocol tests, matching the generic push kernel's `LamportTrait<T, 8, 4>`.
|
||||
using Lamport = device::distributed::LamportTrait<bf16_t, 8, /*kAtom=*/4>;
|
||||
using StageVec = device::AlignedVector<bf16x2_t, 4>;
|
||||
|
||||
SGL_DEVICE void barrier_cluster_arrive_relaxed() {
|
||||
asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory");
|
||||
}
|
||||
|
||||
SGL_DEVICE void barrier_cluster_wait() {
|
||||
asm volatile("barrier.cluster.wait.aligned;" ::: "memory");
|
||||
}
|
||||
|
||||
template <uint32_t kWorldSize, typename WeightT>
|
||||
struct MoeFinalizeAllReduceParams {
|
||||
bf16_t* out; // [num_tokens, kHiddenDim], output-only
|
||||
const bf16_t* gemm2; // [P, kHiddenDim], permuted / padded rows
|
||||
const int32_t* idx; // [num_tokens * kTopK], -1 = dropped slot
|
||||
const WeightT* weights; // [num_tokens, kTopK], scaling already folded in
|
||||
const bf16_t* shared; // [num_tokens, kHiddenDim] (kHasShared only)
|
||||
const bf16_t* norm_weight; // [kHiddenDim] (kNorm only)
|
||||
float norm_eps; // kNorm only
|
||||
// Caller's promise that everything read before the PDL wait is complete when
|
||||
// the predecessor merely *triggers*: no all-reduce on this plane right before
|
||||
// it, and the routing metadata's producers finished (PDL completion is not
|
||||
// transitive through early-triggering kernels). False (the default) waits first.
|
||||
bool prefetch_metadata;
|
||||
uint32_t rank;
|
||||
uint32_t num_tokens;
|
||||
uint32_t num_push_counters; // full counter array size (bumper range end)
|
||||
PushWorkSpace<kWorldSize> ws;
|
||||
bf16_t* mhc_out = nullptr;
|
||||
const bf16_t* residual = nullptr;
|
||||
const float* post = nullptr;
|
||||
const float* comb = nullptr;
|
||||
const float* pre = nullptr;
|
||||
bf16_t* normalized = nullptr;
|
||||
fp8_e4m3_t* quantized = nullptr;
|
||||
uint8_t* scales = nullptr;
|
||||
};
|
||||
|
||||
template <uint32_t kHiddenDim, uint32_t kWorldSize, typename WeightT>
|
||||
SGL_DEVICE void mhc_quant_vec(
|
||||
const MoeFinalizeAllReduceParams<kWorldSize, WeightT>& params,
|
||||
const StageVec& value,
|
||||
uint32_t token,
|
||||
uint32_t hvec) {
|
||||
using namespace device;
|
||||
fp32x2_t v[4];
|
||||
float amax = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
v[j] = cast<fp32x2_t>(value[j]);
|
||||
amax = fmaxf(amax, fmaxf(fabsf(v[j].x), fabsf(v[j].y)));
|
||||
}
|
||||
amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, 1, 4));
|
||||
amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, 2, 4));
|
||||
const float normalized = amax * (1.0f / 448.0f);
|
||||
const uint32_t bits = __float_as_uint(normalized);
|
||||
const uint32_t exponent = (bits >> 23) & 255;
|
||||
const uint32_t mantissa = bits & 0x7fffff;
|
||||
const bool bump = mantissa != 0 && !(exponent == 0 && mantissa <= 0x400000);
|
||||
const uint32_t sf = normalized <= 0 ? 0 : min(exponent + uint32_t(bump), 254u);
|
||||
const float inv_scale = __uint_as_float(sf == 0 ? 0 : (254 - sf) << 23);
|
||||
AlignedVector<fp8x2_e4m3_t, 4> q;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
q[j] = cast<fp8x2_e4m3_t>(
|
||||
fp32x2_t{fminf(fmaxf(v[j].x * inv_scale, -448.0f), 448.0f), fminf(fmaxf(v[j].y * inv_scale, -448.0f), 448.0f)});
|
||||
}
|
||||
q.store(params.quantized + static_cast<int64_t>(token) * kHiddenDim, hvec);
|
||||
if (hvec % 4 == 0) {
|
||||
const uint32_t g = hvec / 4;
|
||||
const uint32_t off = (g / 4) * 512 + ((token % 32) * 4 + (token / 32) % 4) * 4 + g % 4;
|
||||
params.scales[off] = sf;
|
||||
}
|
||||
}
|
||||
|
||||
/// HC=4 post mixing of an already BF16-rounded all-reduce vector, in
|
||||
/// mhc_post_split_h's order: round comb[0]*residual[0], FMA post*x, then the
|
||||
/// remaining three residual streams.
|
||||
template <uint32_t kHiddenDim, bool kCollapse = false, uint32_t kWorldSize, typename WeightT>
|
||||
SGL_DEVICE StageVec mhc_post_vec(
|
||||
const MoeFinalizeAllReduceParams<kWorldSize, WeightT>& params, const StageVec& red, uint32_t token, uint32_t hvec) {
|
||||
using namespace device;
|
||||
StageVec residual[4];
|
||||
fp32x2_t collapsed[4] = {};
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < 4; ++c) {
|
||||
residual[c].load(params.residual + (static_cast<int64_t>(token) * 4 + c) * kHiddenDim, hvec);
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < 4; ++c) {
|
||||
const float post = params.post[token * 4 + c];
|
||||
float comb[4];
|
||||
#pragma unroll
|
||||
for (uint32_t r = 0; r < 4; ++r)
|
||||
comb[r] = params.comb[token * 16 + r * 4 + c];
|
||||
StageVec out;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
const auto x = cast<fp32x2_t>(red[j]);
|
||||
const auto r0 = cast<fp32x2_t>(residual[0][j]);
|
||||
fp32x2_t acc{fmaf(post, x.x, __fmul_rn(comb[0], r0.x)), fmaf(post, x.y, __fmul_rn(comb[0], r0.y))};
|
||||
#pragma unroll
|
||||
for (uint32_t r = 1; r < 4; ++r) {
|
||||
const auto v = cast<fp32x2_t>(residual[r][j]);
|
||||
acc.x = fmaf(comb[r], v.x, acc.x);
|
||||
acc.y = fmaf(comb[r], v.y, acc.y);
|
||||
}
|
||||
out[j] = cast<bf16x2_t>(acc);
|
||||
if constexpr (kCollapse) {
|
||||
const auto rounded = cast<fp32x2_t>(out[j]);
|
||||
const float pre = params.pre[token * 4 + c];
|
||||
collapsed[j].x = fmaf(rounded.x, pre, collapsed[j].x);
|
||||
collapsed[j].y = fmaf(rounded.y, pre, collapsed[j].y);
|
||||
}
|
||||
}
|
||||
out.store(params.mhc_out + (static_cast<int64_t>(token) * 4 + c) * kHiddenDim, hvec);
|
||||
}
|
||||
StageVec result;
|
||||
if constexpr (kCollapse) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j)
|
||||
result[j] = cast<bf16x2_t>(collapsed[j]);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
/// Row geometry: one 16B vector per thread, one cluster per row, so the block
|
||||
/// size follows from the hidden width and the cluster size (the tuning knob).
|
||||
template <uint32_t kHiddenDim, uint32_t kClusterSize>
|
||||
struct RowClusterTrait {
|
||||
static constexpr uint32_t kRowVecs = kHiddenDim / 8; // 16B vectors per row
|
||||
static constexpr uint32_t kBlockSize = kRowVecs / kClusterSize; // threads per block
|
||||
static constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
|
||||
static_assert(kHiddenDim % 8 == 0, "hidden must be a whole number of 16B vectors");
|
||||
static_assert(1 <= kClusterSize && kClusterSize <= 8, "portable cluster sizes only");
|
||||
static_assert(kRowVecs % kClusterSize == 0, "cluster size must divide the row's vector count");
|
||||
static_assert(kBlockSize % device::kWarpThreads == 0, "block must be whole warps");
|
||||
static_assert(kBlockSize <= 1024, "block too large: raise the cluster size");
|
||||
};
|
||||
|
||||
// --- stage 1: the deferred finalize of one 16B vector ------------------------
|
||||
// The shared-expert vector is loaded first so that load is in flight while the
|
||||
// routing rows and the kTopK gathers are fetched.
|
||||
template <uint32_t kHiddenDim, uint32_t kTopK, bool kHasShared, bool kUsePDL, uint32_t kWorldSize, typename WeightT>
|
||||
SGL_DEVICE StageVec
|
||||
finalize_vec(const MoeFinalizeAllReduceParams<kWorldSize, WeightT>& params, uint32_t token, uint32_t hvec) {
|
||||
using namespace device;
|
||||
const auto* idx = params.idx + static_cast<int64_t>(token) * kTopK;
|
||||
const auto* weights = params.weights + static_cast<int64_t>(token) * kTopK;
|
||||
int32_t rows[kTopK];
|
||||
WeightT w[kTopK];
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < kTopK; ++k) {
|
||||
rows[k] = idx[k];
|
||||
w[k] = weights[k];
|
||||
}
|
||||
|
||||
// delay PDL wait until here
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
StageVec shared_in;
|
||||
if constexpr (kHasShared) {
|
||||
shared_in.load(params.shared + static_cast<int64_t>(token) * kHiddenDim, hvec);
|
||||
}
|
||||
|
||||
StageVec in[kTopK];
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < kTopK; ++k) {
|
||||
if (rows[k] >= 0) in[k].load(params.gemm2 + static_cast<int64_t>(rows[k]) * kHiddenDim, hvec);
|
||||
}
|
||||
|
||||
fp32x2_t acc[4];
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
acc[j] = fp32x2_t{0.0f, 0.0f};
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < kTopK; ++k) {
|
||||
if (rows[k] < 0) continue;
|
||||
#if SGL_ARCH_BLACKWELL_OR_GREATER
|
||||
if constexpr (std::is_same_v<WeightT, bf16_t>) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
acc[j].x = math::fma_f32_bf16(in[k][j].x, w[k], acc[j].x);
|
||||
acc[j].y = math::fma_f32_bf16(in[k][j].y, w[k], acc[j].y);
|
||||
}
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
const auto w_fp32 = cast<fp32_t>(w[k]);
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
const auto [x, y] = cast<fp32x2_t>(in[k][j]);
|
||||
acc[j].x = fmaf(x, w_fp32, acc[j].x);
|
||||
acc[j].y = fmaf(y, w_fp32, acc[j].y);
|
||||
}
|
||||
}
|
||||
}
|
||||
// Two deliberate roundings -- the routed combine, then the bf16 + bf16 add --
|
||||
// keep the staged vector bit-identical to the unfused rank-local result.
|
||||
StageVec out;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
if constexpr (kHasShared) {
|
||||
const auto routed = cast<fp32x2_t>(cast<bf16x2_t>(acc[j]));
|
||||
const auto sh = cast<fp32x2_t>(shared_in[j]);
|
||||
out[j] = cast<bf16x2_t>(fp32x2_t{routed.x + sh.x, routed.y + sh.y});
|
||||
} else {
|
||||
out[j] = cast<bf16x2_t>(acc[j]);
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// --- the kernel --------------------------------------------------------------
|
||||
// Grid dim3(num_tokens [+ 1], kClusterSize) with the cluster along y, so
|
||||
// blockIdx.x is the token row (and its phase counter) and blockIdx.y the rank
|
||||
// inside the cluster. The extra cluster (blockIdx.x == num_tokens) is the
|
||||
// bumper: it only flips the leftover counters [num_tokens, num_push_counters).
|
||||
template <
|
||||
uint32_t kWorldSize,
|
||||
uint32_t kHiddenDim,
|
||||
uint32_t kTopK,
|
||||
uint32_t kClusterSize,
|
||||
bool kUsePDL,
|
||||
bool kHasShared,
|
||||
bool kNorm,
|
||||
typename WeightT,
|
||||
bool kMhc = false,
|
||||
bool kQuant = false>
|
||||
__global__ __launch_bounds__(RowClusterTrait<kHiddenDim, kClusterSize>::kBlockSize)
|
||||
__cluster_dims__(1, kClusterSize, 1) void moe_finalize_all_reduce_kernel(
|
||||
const __grid_constant__ MoeFinalizeAllReduceParams<kWorldSize, WeightT> params) {
|
||||
namespace cg = cooperative_groups;
|
||||
using namespace device;
|
||||
using T = RowClusterTrait<kHiddenDim, kClusterSize>;
|
||||
constexpr uint32_t kRowVecs = T::kRowVecs;
|
||||
constexpr uint32_t kBlockSize = T::kBlockSize;
|
||||
constexpr uint32_t kNumWarps = T::kNumWarps;
|
||||
|
||||
const auto tx = threadIdx.x;
|
||||
const auto row_idx = blockIdx.x;
|
||||
const auto cluster_rank = blockIdx.y;
|
||||
// this thread's vector within a row: cluster rank picks the block's chunk
|
||||
const auto hvec = cluster_rank * kBlockSize + tx;
|
||||
|
||||
// Reading the epoch before the PDL wait can see a predecessor all-reduce
|
||||
// mid-flip on this plane; prefetch_metadata defers the wait to finalize_vec.
|
||||
if (!params.prefetch_metadata) PDLWaitPrimary<kUsePDL>();
|
||||
|
||||
if (row_idx == params.num_tokens) {
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
if constexpr (kQuant) {
|
||||
// The SF buffer is padded to 128 rows. Active rows are written by the
|
||||
// norm epilogue; the existing bumper zeros only the disjoint padding.
|
||||
for (uint32_t off = hvec; off < (kHiddenDim / 32) * 128; off += kRowVecs) {
|
||||
const uint32_t swizzled_row = (off % 512) / 4;
|
||||
const uint32_t row = swizzled_row / 4 + (swizzled_row % 4) * 32;
|
||||
if (row >= params.num_tokens) params.scales[off] = 0;
|
||||
}
|
||||
}
|
||||
if (cluster_rank == 0) {
|
||||
const auto epoch = distributed::PushEpoch<kWorldSize>{params.ws};
|
||||
__syncthreads();
|
||||
epoch.unsafe_flip_range(row_idx, params.num_push_counters);
|
||||
}
|
||||
return PDLTriggerSecondary<kUsePDL>();
|
||||
}
|
||||
|
||||
// this cluster's epoch: the counter at blockIdx.x, one per row cluster
|
||||
// (every block of the cluster reads the same one)
|
||||
const auto epoch = distributed::PushEpoch<kWorldSize>{params.ws};
|
||||
const auto r = params.rank;
|
||||
// my slot (`src = r`) inside every peer's workspace, and every peer's slot
|
||||
// inside mine (`dst = r`), for this epoch
|
||||
void* push_ptrs[kWorldSize];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
push_ptrs[i] = epoch.slot_ptr(/*dst=*/i, /*src=*/r);
|
||||
}
|
||||
|
||||
// stage 1: finalize this row's vector in registers and push it to every peer
|
||||
const auto vid = row_idx * kRowVecs + hvec;
|
||||
{
|
||||
auto vec = finalize_vec<kHiddenDim, kTopK, kHasShared, kUsePDL>(params, row_idx, hvec);
|
||||
Lamport::clear_pos_zero(vec.data());
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
ptx::st_relaxed_16B(vec, push_ptrs[i], vid);
|
||||
}
|
||||
}
|
||||
|
||||
// ensure epoch is consumed, so flipping it won't lead to error
|
||||
if constexpr (!kNorm) barrier_cluster_arrive_relaxed();
|
||||
|
||||
// stage 2: poll own slots, reduce across ranks, [norm], write, reset markers
|
||||
void* poll_ptrs[kWorldSize];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
poll_ptrs[i] = epoch.slot_ptr(/*dst=*/r, /*src=*/i);
|
||||
}
|
||||
StageVec vec[kWorldSize];
|
||||
do {
|
||||
bool has_zero = false;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
ptx::ld_relaxed_16B(vec[i], poll_ptrs[i], vid);
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
// the producer remapped +0.0 words, so a written word is never 0:
|
||||
// word == 0 <=> the slot still holds the empty marker
|
||||
has_zero |= Lamport::has_pos_zero(vec[i].data());
|
||||
}
|
||||
if (!has_zero) break;
|
||||
} while (true);
|
||||
|
||||
if constexpr (!kNorm) {
|
||||
const auto red = reduce_vec(vec);
|
||||
ptx::st_global_16B(red, params.out, vid);
|
||||
if constexpr (kMhc) mhc_post_vec<kHiddenDim>(params, red, row_idx, hvec);
|
||||
// ensure epoch is consumed, so flipping it won't lead to error
|
||||
barrier_cluster_wait();
|
||||
} else {
|
||||
// push to peer
|
||||
__shared__ float smem_sq[kClusterSize][kNumWarps];
|
||||
auto red = reduce_vec(vec);
|
||||
if constexpr (kMhc) {
|
||||
ptx::st_global_16B(red, params.out, vid);
|
||||
red = mhc_post_vec<kHiddenDim, true>(params, red, row_idx, hvec);
|
||||
}
|
||||
StageVec w;
|
||||
w.load(params.norm_weight, hvec);
|
||||
const auto cluster = cg::this_cluster();
|
||||
fp32x2_t acc[4];
|
||||
float sq = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
acc[j] = cast<fp32x2_t>(red[j]);
|
||||
sq = fmaf(acc[j].x, acc[j].x, sq);
|
||||
sq = fmaf(acc[j].y, acc[j].y, sq);
|
||||
}
|
||||
sq = warp::reduce_sum(sq);
|
||||
const auto lane = tx % kWarpThreads;
|
||||
const auto warp = tx / kWarpThreads;
|
||||
if (lane < kClusterSize) {
|
||||
*cluster.map_shared_rank(&smem_sq[cluster_rank][warp], lane) = sq;
|
||||
}
|
||||
cluster.sync();
|
||||
float total = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < kClusterSize; ++c) {
|
||||
#pragma unroll
|
||||
for (uint32_t wp = 0; wp < kNumWarps; ++wp) {
|
||||
total += smem_sq[c][wp];
|
||||
}
|
||||
}
|
||||
const auto factor = math::rsqrt(total / static_cast<float>(kHiddenDim) + params.norm_eps);
|
||||
StageVec out;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < 4; ++j) {
|
||||
const auto [wa, wb] = cast<fp32x2_t>(w[j]);
|
||||
out[j] = cast<bf16x2_t>(fp32x2_t{acc[j].x * factor * wa, acc[j].y * factor * wb});
|
||||
}
|
||||
ptx::st_global_16B(out, kMhc ? params.normalized : params.out, vid);
|
||||
if constexpr (kQuant) mhc_quant_vec<kHiddenDim>(params, out, row_idx, hvec);
|
||||
}
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
// re-establish the empty markers for the next same-phase round
|
||||
StageVec zero_vec;
|
||||
Lamport::fill_pos_zero(zero_vec.data());
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
ptx::st_global_16B(zero_vec, poll_ptrs[i], vid);
|
||||
}
|
||||
|
||||
if (cluster_rank == 0) epoch.flip();
|
||||
}
|
||||
|
||||
// --- host --------------------------------------------------------------------
|
||||
|
||||
template <
|
||||
uint32_t kWorldSize,
|
||||
uint32_t kHiddenDim,
|
||||
uint32_t kTopK,
|
||||
uint32_t kClusterSize,
|
||||
bool kUsePDL,
|
||||
typename WeightT,
|
||||
bool kMhc = false,
|
||||
bool kQuant = false>
|
||||
struct MoeFinalizeAllReduceKernel {
|
||||
private:
|
||||
static_assert(std::is_same_v<WeightT, bf16_t> || std::is_same_v<WeightT, fp32_t>);
|
||||
using TensorView = tvm::ffi::TensorView;
|
||||
using Params = MoeFinalizeAllReduceParams<kWorldSize, WeightT>;
|
||||
using Trait = RowClusterTrait<kHiddenDim, kClusterSize>;
|
||||
|
||||
template <bool kHasShared, bool kNorm>
|
||||
static constexpr auto kernel = moe_finalize_all_reduce_kernel<
|
||||
kWorldSize,
|
||||
kHiddenDim,
|
||||
kTopK,
|
||||
kClusterSize,
|
||||
kUsePDL,
|
||||
kHasShared,
|
||||
kNorm,
|
||||
WeightT,
|
||||
kMhc,
|
||||
kQuant>;
|
||||
|
||||
public:
|
||||
/// out = [allreduce over ranks of] finalize(gemm2_out, idx, weights) [+ shared] [-> RMSNorm(norm_weight, eps)].
|
||||
/// `out` ([T, kHiddenDim] bf16) is output-only. `shared_output` and `norm_weight`
|
||||
/// select the epilogue at runtime (four kernel instantiations per module).
|
||||
static void
|
||||
run(CommunicatorRef ref,
|
||||
TensorView out,
|
||||
TensorView gemm2_out,
|
||||
TensorView permuted_idx,
|
||||
TensorView expert_weights,
|
||||
std::optional<TensorView> shared_output,
|
||||
std::optional<TensorView> norm_weight,
|
||||
double eps,
|
||||
bool prefetch_metadata) {
|
||||
static_assert(!kMhc);
|
||||
run_impl(
|
||||
ref,
|
||||
out,
|
||||
gemm2_out,
|
||||
permuted_idx,
|
||||
expert_weights,
|
||||
shared_output,
|
||||
norm_weight,
|
||||
eps,
|
||||
prefetch_metadata,
|
||||
std::nullopt,
|
||||
std::nullopt,
|
||||
std::nullopt,
|
||||
std::nullopt);
|
||||
}
|
||||
|
||||
/// Finalize + all-reduce + HC=4 post; original reduced output is retained.
|
||||
static void run_mhc(
|
||||
CommunicatorRef ref,
|
||||
TensorView out,
|
||||
TensorView gemm2_out,
|
||||
TensorView permuted_idx,
|
||||
TensorView expert_weights,
|
||||
std::optional<TensorView> shared_output,
|
||||
TensorView mhc_out,
|
||||
TensorView residual,
|
||||
TensorView post,
|
||||
TensorView comb) {
|
||||
static_assert(kMhc);
|
||||
run_impl(
|
||||
ref,
|
||||
out,
|
||||
gemm2_out,
|
||||
permuted_idx,
|
||||
expert_weights,
|
||||
shared_output,
|
||||
std::nullopt,
|
||||
0.0,
|
||||
false,
|
||||
mhc_out,
|
||||
residual,
|
||||
post,
|
||||
comb);
|
||||
}
|
||||
|
||||
static void run_mhc_norm(
|
||||
CommunicatorRef ref,
|
||||
TensorView out,
|
||||
TensorView gemm2_out,
|
||||
TensorView permuted_idx,
|
||||
TensorView expert_weights,
|
||||
std::optional<TensorView> shared_output,
|
||||
TensorView mhc_out,
|
||||
TensorView residual,
|
||||
TensorView post,
|
||||
TensorView comb,
|
||||
TensorView pre,
|
||||
TensorView norm_weight,
|
||||
double eps,
|
||||
TensorView normalized) {
|
||||
static_assert(kMhc);
|
||||
run_impl(
|
||||
ref,
|
||||
out,
|
||||
gemm2_out,
|
||||
permuted_idx,
|
||||
expert_weights,
|
||||
shared_output,
|
||||
norm_weight,
|
||||
eps,
|
||||
false,
|
||||
mhc_out,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
normalized);
|
||||
}
|
||||
|
||||
static void run_mhc_quant(
|
||||
CommunicatorRef ref,
|
||||
TensorView out,
|
||||
TensorView gemm2_out,
|
||||
TensorView permuted_idx,
|
||||
TensorView expert_weights,
|
||||
std::optional<TensorView> shared_output,
|
||||
TensorView mhc_out,
|
||||
TensorView residual,
|
||||
TensorView post,
|
||||
TensorView comb,
|
||||
TensorView pre,
|
||||
TensorView norm_weight,
|
||||
double eps,
|
||||
TensorView normalized,
|
||||
TensorView quantized,
|
||||
TensorView scales) {
|
||||
static_assert(kMhc && kQuant);
|
||||
run_impl(
|
||||
ref,
|
||||
out,
|
||||
gemm2_out,
|
||||
permuted_idx,
|
||||
expert_weights,
|
||||
shared_output,
|
||||
norm_weight,
|
||||
eps,
|
||||
false,
|
||||
mhc_out,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
normalized,
|
||||
quantized,
|
||||
scales);
|
||||
}
|
||||
|
||||
private:
|
||||
static void run_impl(
|
||||
CommunicatorRef ref,
|
||||
TensorView out,
|
||||
TensorView gemm2_out,
|
||||
TensorView permuted_idx,
|
||||
TensorView expert_weights,
|
||||
std::optional<TensorView> shared_output,
|
||||
std::optional<TensorView> norm_weight,
|
||||
double eps,
|
||||
bool prefetch_metadata,
|
||||
std::optional<TensorView> mhc_out,
|
||||
std::optional<TensorView> residual,
|
||||
std::optional<TensorView> post,
|
||||
std::optional<TensorView> comb,
|
||||
std::optional<TensorView> pre = std::nullopt,
|
||||
std::optional<TensorView> normalized = std::nullopt,
|
||||
std::optional<TensorView> quantized = std::nullopt,
|
||||
std::optional<TensorView> scales = std::nullopt) {
|
||||
using namespace host;
|
||||
const auto& comm = *ref.get();
|
||||
const auto& push = comm.get_push_obj();
|
||||
CHECK_HOST(push.world_size == kWorldSize)
|
||||
<< "communicator holds " << push.world_size << " ranks, kernel built for " << kWorldSize;
|
||||
|
||||
auto T = SymbolicSize{"num_tokens"};
|
||||
auto P = SymbolicSize{"num_permuted_rows"};
|
||||
auto TK = SymbolicSize{"num_expanded"};
|
||||
SymbolicDevice device;
|
||||
device.set_options<kDLCUDA>();
|
||||
TensorMatcher({T, kHiddenDim})
|
||||
.with_strides({kHiddenDim, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(out);
|
||||
TensorMatcher({P, kHiddenDim})
|
||||
.with_strides({kHiddenDim, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(gemm2_out);
|
||||
TensorMatcher({T, kTopK})
|
||||
.with_strides({kTopK, 1})
|
||||
.with_dtype<WeightT>()
|
||||
.template with_device<kDLCUDA>(device)
|
||||
.verify(expert_weights);
|
||||
TK.set_value(T.unwrap() * kTopK);
|
||||
TensorMatcher({TK}).with_strides({1}).with_dtype<int32_t>().with_device<kDLCUDA>(device).verify(permuted_idx);
|
||||
if (shared_output.has_value()) {
|
||||
TensorMatcher({T, kHiddenDim})
|
||||
.with_strides({kHiddenDim, 1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(shared_output.value());
|
||||
}
|
||||
if (norm_weight.has_value()) {
|
||||
TensorMatcher({kHiddenDim})
|
||||
.with_strides({1})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(norm_weight.value());
|
||||
}
|
||||
const auto num_tokens = static_cast<uint32_t>(T.unwrap());
|
||||
if constexpr (kQuant) {
|
||||
CHECK_HOST(num_tokens <= 8);
|
||||
CHECK_HOST(norm_weight.has_value());
|
||||
TensorMatcher({T, kHiddenDim}).with_dtype<fp8_e4m3_t>().with_device<kDLCUDA>(device).verify(quantized.value());
|
||||
TensorMatcher({(kHiddenDim / 32) * 128})
|
||||
.with_dtype<uint8_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(scales.value());
|
||||
}
|
||||
if constexpr (kMhc) {
|
||||
static_assert(kHiddenDim == 5120);
|
||||
if (norm_weight.has_value()) {
|
||||
TensorMatcher({T, 4}).with_dtype<fp32_t>().with_device<kDLCUDA>(device).verify(pre.value());
|
||||
TensorMatcher({T, kHiddenDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(normalized.value());
|
||||
}
|
||||
TensorMatcher({T, 4, kHiddenDim})
|
||||
.with_dtype<bf16_t>()
|
||||
.with_device<kDLCUDA>(device)
|
||||
.verify(mhc_out.value())
|
||||
.verify(residual.value());
|
||||
TensorMatcher({T, 4}).with_dtype<fp32_t>().with_device<kDLCUDA>(device).verify(post.value());
|
||||
TensorMatcher({T, 4, 4}).with_dtype<fp32_t>().with_device<kDLCUDA>(device).verify(comb.value());
|
||||
}
|
||||
CHECK_HOST(num_tokens > 0) << "num_tokens must be positive";
|
||||
CHECK_HOST(reinterpret_cast<uintptr_t>(gemm2_out.data_ptr()) % 16 == 0) << "gemm2_out must be 16B aligned";
|
||||
CHECK_HOST(reinterpret_cast<uintptr_t>(out.data_ptr()) % 16 == 0) << "out must be 16B aligned";
|
||||
|
||||
// the whole [T, hidden] row view is staged by vector index, so it must fit
|
||||
// one push slot; the generic push kernel's callers pick the slot size
|
||||
const int64_t nbytes = num_tokens * int64_t(kHiddenDim) * sizeof(bf16_t);
|
||||
CHECK_HOST(nbytes <= push.slot_bytes) << "num_tokens * hidden * 2 = " << nbytes << " bytes exceeds the "
|
||||
<< push.slot_bytes << "-byte push slot (reduce the batch or enlarge "
|
||||
<< "max_push_size)";
|
||||
// one cluster (and phase counter) per row; the bumper cluster is launched
|
||||
// only when counters are left over for it to flip
|
||||
CHECK_HOST(num_tokens <= push.num_blocks)
|
||||
<< "num_tokens = " << num_tokens << " exceeds the " << push.num_blocks << " push phase counters of the plane";
|
||||
const uint32_t num_clusters = num_tokens + (num_tokens < push.num_blocks ? 1 : 0);
|
||||
|
||||
const auto params = Params{
|
||||
.out = static_cast<bf16_t*>(out.data_ptr()),
|
||||
.gemm2 = static_cast<const bf16_t*>(gemm2_out.data_ptr()),
|
||||
.idx = static_cast<const int32_t*>(permuted_idx.data_ptr()),
|
||||
.weights = static_cast<const WeightT*>(expert_weights.data_ptr()),
|
||||
.shared = shared_output.has_value() ? static_cast<const bf16_t*>(shared_output.value().data_ptr()) : nullptr,
|
||||
.norm_weight = norm_weight.has_value() ? static_cast<const bf16_t*>(norm_weight.value().data_ptr()) : nullptr,
|
||||
.norm_eps = static_cast<float>(eps),
|
||||
.prefetch_metadata = prefetch_metadata,
|
||||
.rank = push.rank,
|
||||
.num_tokens = num_tokens,
|
||||
.num_push_counters = push.num_blocks,
|
||||
.ws = push.get_workspace<kWorldSize>(nbytes),
|
||||
.mhc_out = mhc_out.has_value() ? static_cast<bf16_t*>(mhc_out.value().data_ptr()) : nullptr,
|
||||
.residual = residual.has_value() ? static_cast<const bf16_t*>(residual.value().data_ptr()) : nullptr,
|
||||
.post = post.has_value() ? static_cast<const float*>(post.value().data_ptr()) : nullptr,
|
||||
.comb = comb.has_value() ? static_cast<const float*>(comb.value().data_ptr()) : nullptr,
|
||||
.pre = pre.has_value() ? static_cast<const float*>(pre.value().data_ptr()) : nullptr,
|
||||
.normalized = normalized.has_value() ? static_cast<bf16_t*>(normalized.value().data_ptr()) : nullptr,
|
||||
.quantized = quantized.has_value() ? static_cast<fp8_e4m3_t*>(quantized.value().data_ptr()) : nullptr,
|
||||
.scales = scales.has_value() ? static_cast<uint8_t*>(scales.value().data_ptr()) : nullptr,
|
||||
};
|
||||
|
||||
const auto has_shared = shared_output.has_value();
|
||||
const auto has_norm = norm_weight.has_value();
|
||||
const auto kern = has_shared ? (has_norm ? kernel<true, true> : kernel<true, false>)
|
||||
: (has_norm ? kernel<false, true> : kernel<false, false>);
|
||||
// __cluster_dims__(1, kClusterSize, 1) is compiled in, so a plain launch
|
||||
// already forms the clusters along y
|
||||
LaunchKernel(dim3(num_clusters, kClusterSize), Trait::kBlockSize, out.device()).enable_pdl(kUsePDL)(kern, params);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
@@ -142,7 +142,16 @@ ALL_REDUCE_KERNEL void all_reduce_1shot_push_kernel(const __grid_constant__ AllR
|
||||
const auto r = params.rank;
|
||||
const auto num_vecs = params.num_vecs;
|
||||
const auto num_threads = blockDim.x * gridDim.x;
|
||||
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
// Round-robin warps to blocks rather than giving each block a contiguous run.
|
||||
// The grid is pinned to the counter array, so once `num_vecs` stops filling
|
||||
// `gridDim * blockDim` a block-major index leaves the tail CTAs with nothing
|
||||
// to do; that happens over a whole 2x band of sizes, between the point where
|
||||
// `choose_block_size` gives up on 512 and the point where 1024 threads fill
|
||||
// the grid again.
|
||||
const auto warp_in_block = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto global_warp_id = blockIdx.x + gridDim.x * warp_in_block;
|
||||
const auto global_tid = global_warp_id * kWarpThreads + lane_id;
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
const auto epoch = distributed::PushEpoch<kWorldSize>{params.ws};
|
||||
|
||||
|
||||
@@ -0,0 +1,660 @@
|
||||
#include <sgl_kernel/ffi.h>
|
||||
#include <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/runtime.cuh>
|
||||
#include <sgl_kernel/type.cuh>
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
#include <sgl_kernel/vec.cuh>
|
||||
|
||||
#include <sgl_kernel/distributed/communicator.cuh>
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/extra/stl.h>
|
||||
#include <tvm/ffi/object.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <bit>
|
||||
#include <cstdint>
|
||||
|
||||
namespace sglang {
|
||||
|
||||
using device::distributed::PushWorkSpace;
|
||||
using device::distributed::Semaphore;
|
||||
|
||||
// Runtime uint32 division as a multiply-high and a shift (round-up magic,
|
||||
// exact below 2^31); cuda::fast_mod_div needs a newer CCCL than CUDA 13 bundles.
|
||||
struct fast_mod_div_u32_t {
|
||||
uint32_t divisor;
|
||||
uint32_t magic;
|
||||
uint32_t shift;
|
||||
|
||||
__host__ explicit fast_mod_div_u32_t(uint32_t d) : divisor(d), magic(0), shift(0) {
|
||||
if (d > 1) {
|
||||
const uint32_t log2_ceil = 32 - std::countl_zero(d - 1);
|
||||
const uint32_t p = 31 + log2_ceil;
|
||||
magic = static_cast<uint32_t>(((uint64_t{1} << p) + d - 1) / d);
|
||||
shift = p - 32;
|
||||
}
|
||||
}
|
||||
|
||||
__device__ friend uint32_t operator/(uint32_t n, const fast_mod_div_u32_t& fd) {
|
||||
return fd.divisor == 1 ? n : __umulhi(n, fd.magic) >> fd.shift;
|
||||
}
|
||||
|
||||
__device__ friend uint32_t operator%(uint32_t n, const fast_mod_div_u32_t& fd) {
|
||||
return n - (n / fd) * fd.divisor;
|
||||
}
|
||||
};
|
||||
|
||||
template <uint32_t kWorldSize>
|
||||
struct NVLinkCommPushParams {
|
||||
const void* __restrict__ input;
|
||||
const void* __restrict__ residual;
|
||||
void* __restrict__ output;
|
||||
uint32_t dst_offset; // AR = rank slot stride; AG = packed token prefix
|
||||
uint32_t rank;
|
||||
uint32_t num_push_vecs;
|
||||
uint32_t num_poll_vecs;
|
||||
uint32_t num_vecs_per_token;
|
||||
// Ragged split, reduce-scatter only: rank r owns `avg + (r < rem)` tokens of
|
||||
// the input starting at `r * avg + min(r, rem)`.
|
||||
uint32_t tokens_avg;
|
||||
uint32_t tokens_rem;
|
||||
fast_mod_div_u32_t vecs_per_token_div;
|
||||
PushWorkSpace<kWorldSize> ws;
|
||||
};
|
||||
|
||||
struct NVLinkCommPullParams {
|
||||
const void* __restrict__ input;
|
||||
const void* __restrict__ residual;
|
||||
void* __restrict__ output;
|
||||
uint32_t num_vecs;
|
||||
// multicast buffer
|
||||
uint8_t* input_mc;
|
||||
uint8_t* output_mc;
|
||||
Semaphore* sem_local;
|
||||
Semaphore* sem_mc;
|
||||
uint32_t rank;
|
||||
uint32_t world_size;
|
||||
};
|
||||
|
||||
template <bool kHasResidual>
|
||||
inline constexpr uint32_t get_poll_group(uint32_t world_size) {
|
||||
if (world_size <= 8) return world_size;
|
||||
return kHasResidual ? 6 : 8;
|
||||
}
|
||||
|
||||
inline constexpr uint32_t kPushCTASize = 1024; // max value
|
||||
inline constexpr uint32_t kPullCTASize = 512; // max value
|
||||
|
||||
#define PUSH_KERNEL __global__ __launch_bounds__(kPushCTASize, 1)
|
||||
#define PULL_KERNEL __global__ __launch_bounds__(kPullCTASize, 1)
|
||||
|
||||
enum Primitive {
|
||||
RS = 0b01, // Reduce-Scatter
|
||||
AG = 0b10, // All-Gather
|
||||
AR = RS | AG, // All-Reduce = RS + AG
|
||||
};
|
||||
|
||||
template <typename vec_t>
|
||||
SGL_DEVICE vec_t reduce_vec(vec_t x, vec_t y) {
|
||||
vec_t arr[2] = {x, y};
|
||||
return device::reduce_vec(arr);
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Layout:
|
||||
* 1. `AG`/`RS`: each rank push to its own slot
|
||||
* [rank0] | [rank 1] | [rank 2] | ...
|
||||
* 2. `AG`: each rank push to a contiguous region
|
||||
* [rank0, rank1, rank2, ...]
|
||||
*
|
||||
* `RS` use swizzle layout for push kernel \n
|
||||
* `AG` use normal linear layout for push kernel
|
||||
*/
|
||||
template <typename T, bool kHasResidual, Primitive kPrim, uint32_t kWorldSize, bool kUsePDL>
|
||||
PUSH_KERNEL void nvlink_push_kernel(const __grid_constant__ NVLinkCommPushParams<kWorldSize> params) {
|
||||
using namespace device;
|
||||
enable_smem_spilling();
|
||||
constexpr uint32_t kVecSize = 16 / sizeof(T); // 16 bytes per vector
|
||||
using vec_t = device::AlignedVector<packed_t<T>, kVecSize / 2>;
|
||||
using Lamport = distributed::LamportTrait<T, kVecSize, /*kAtom=*/4>;
|
||||
constexpr uint32_t kGroup = get_poll_group<kHasResidual>(kWorldSize);
|
||||
|
||||
// Round-robin warps to blocks: the poll domain is this rank's shard, so a
|
||||
// block-major index would park all of it on the first few CTAs and idle the
|
||||
// rest of the SMs.
|
||||
const auto warp_in_block = threadIdx.x / kWarpThreads;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto global_warp_id = blockIdx.x + gridDim.x * warp_in_block;
|
||||
const auto global_tid = global_warp_id * kWarpThreads + lane_id;
|
||||
const auto num_threads = blockDim.x * gridDim.x;
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
const auto epoch = distributed::PushEpoch<kWorldSize>{params.ws};
|
||||
|
||||
void* push_ptrs[kWorldSize];
|
||||
/// NOTE: broadcast write is only fast when world size is large
|
||||
if constexpr (kWorldSize < 8 && (kPrim & Primitive::AG)) {
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
// `dst_offset` is a slot stride for the all-reduce but a packed token
|
||||
// prefix for the gather, whose consumer reads the plane linearly; this
|
||||
// must stay the same address arithmetic as the multicast branch below.
|
||||
push_ptrs[i] = static_cast<uint8_t*>(epoch.slot_ptr(/*dst=*/i)) + params.dst_offset;
|
||||
}
|
||||
}
|
||||
|
||||
const auto dst_ptr_mc = params.ws.mc_workspace + params.dst_offset + epoch.slot_offset();
|
||||
const auto vpt = params.num_vecs_per_token;
|
||||
|
||||
for (auto vid = global_tid; vid < params.num_push_vecs; vid += num_threads) {
|
||||
if constexpr (kPrim & Primitive::AG) {
|
||||
vec_t vec;
|
||||
vec.load(params.input, vid);
|
||||
if constexpr (kHasResidual && kPrim == Primitive::AG) {
|
||||
vec_t res;
|
||||
res.load(params.residual, vid);
|
||||
vec = reduce_vec(vec, res);
|
||||
}
|
||||
Lamport::clear_pos_zero(vec.data());
|
||||
if constexpr (kWorldSize < 8) {
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
ptx::st_relaxed_16B(vec, push_ptrs[i], vid);
|
||||
}
|
||||
} else {
|
||||
ptx::st_multimem_16B(vec, dst_ptr_mc, vid);
|
||||
}
|
||||
} else /* reduce-scatter only */ {
|
||||
const auto token_id = vid / params.vecs_per_token_div;
|
||||
const auto offset = vid % params.vecs_per_token_div;
|
||||
// Both by a compile-time constant, so this is a mask and a shift.
|
||||
const auto dst_rank = token_id % kWorldSize;
|
||||
const auto dst_token_id = token_id / kWorldSize;
|
||||
// Round-robin over peers so every link stays busy instead of one congesting
|
||||
const auto avg_tokens = params.tokens_avg;
|
||||
const auto rem_tokens = params.tokens_rem;
|
||||
const auto rank_prefix = dst_rank * avg_tokens + std::min(dst_rank, rem_tokens);
|
||||
const auto src_token = rank_prefix + dst_token_id;
|
||||
vec_t vec;
|
||||
vec.load(params.input, src_token * vpt + offset);
|
||||
const auto dst_ptr = epoch.slot_ptr(dst_rank, params.rank);
|
||||
Lamport::clear_pos_zero(vec.data());
|
||||
ptx::st_relaxed_16B(vec, dst_ptr, dst_token_id * vpt + offset);
|
||||
}
|
||||
}
|
||||
|
||||
// Poll addresses are linear in the source rank -- one base, `slot_bytes`
|
||||
// apart -- so a base plus a vector-index bias replaces a per-peer pointer table.
|
||||
const auto poll_base = epoch.slot_ptr(params.rank);
|
||||
const auto slot_vecs = params.ws.slot_bytes / sizeof(vec_t);
|
||||
vec_t pos_zero_vec;
|
||||
Lamport::fill_pos_zero(pos_zero_vec.data());
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
|
||||
for (auto vid = global_tid; vid < params.num_poll_vecs; vid += num_threads) {
|
||||
if constexpr (kPrim & Primitive::RS) {
|
||||
constexpr uint32_t kNumPairs = kVecSize / 2;
|
||||
vec_t out_vec;
|
||||
|
||||
if constexpr (kGroup >= kWorldSize) {
|
||||
vec_t vec[kWorldSize + kHasResidual];
|
||||
if constexpr (kHasResidual) vec[kWorldSize].load(params.residual, vid);
|
||||
do {
|
||||
bool has_zero = false;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
ptx::ld_relaxed_16B(vec[i], poll_base, i * slot_vecs + vid);
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
has_zero |= Lamport::has_pos_zero(vec[i].data());
|
||||
}
|
||||
if (!has_zero) break;
|
||||
} while (true);
|
||||
out_vec = reduce_vec(vec);
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kWorldSize; ++i) {
|
||||
ptx::st_global_16B(pos_zero_vec, poll_base, i * slot_vecs + vid);
|
||||
}
|
||||
} else /* > 1 group: divide into chunks */ {
|
||||
fp32x2_t acc[kNumPairs];
|
||||
constexpr uint32_t kNumGroups = div_ceil(kWorldSize, kGroup);
|
||||
vec_t vec[kGroup];
|
||||
vec_t res;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t g = 0; g < kNumGroups; ++g) {
|
||||
const auto for_each = [&](auto&& fn) {
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < kGroup; ++j) {
|
||||
const auto i = g * kGroup + j;
|
||||
if (i >= kWorldSize) continue;
|
||||
fn(i, j);
|
||||
}
|
||||
};
|
||||
|
||||
// Loaded a group early so the fetch overlaps the last poll; it is
|
||||
// folded into the accumulator once the groups are done.
|
||||
if constexpr (kHasResidual) {
|
||||
if (g + 1 == kNumGroups) res.load(params.residual, vid);
|
||||
}
|
||||
|
||||
do {
|
||||
bool has_zero = false;
|
||||
for_each([&](uint32_t i, uint32_t j) {
|
||||
// load all the vectors
|
||||
ptx::ld_relaxed_16B(vec[j], poll_base, i * slot_vecs + vid);
|
||||
});
|
||||
for_each([&](uint32_t, uint32_t j) {
|
||||
// check for zeros
|
||||
has_zero |= Lamport::has_pos_zero(vec[j].data());
|
||||
});
|
||||
if (!has_zero) break;
|
||||
} while (true);
|
||||
|
||||
for_each([&](uint32_t i, uint32_t j) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < kNumPairs; ++k) {
|
||||
const auto [x, y] = cast<fp32x2_t>(vec[j][k]);
|
||||
acc[k].x = i == 0 ? x : acc[k].x + x;
|
||||
acc[k].y = i == 0 ? y : acc[k].y + y;
|
||||
}
|
||||
ptx::st_global_16B(pos_zero_vec, poll_base, i * slot_vecs + vid);
|
||||
});
|
||||
}
|
||||
if constexpr (kHasResidual) {
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < kNumPairs; ++k) {
|
||||
const auto [x, y] = cast<fp32x2_t>(res[k]);
|
||||
acc[k].x += x;
|
||||
acc[k].y += y;
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < kNumPairs; ++k) {
|
||||
out_vec[k] = cast<packed_t<T>>(acc[k]);
|
||||
}
|
||||
}
|
||||
|
||||
out_vec.store(params.output, vid);
|
||||
} else /* all-gather only */ {
|
||||
vec_t vec;
|
||||
do {
|
||||
ptx::ld_relaxed_16B(vec, poll_base, vid);
|
||||
} while (Lamport::has_pos_zero(vec.data()));
|
||||
vec.store(params.output, vid);
|
||||
ptx::st_global_16B(pos_zero_vec, poll_base, vid);
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
epoch.flip();
|
||||
}
|
||||
|
||||
template <typename T, bool kHasResidual, Primitive kPrim, bool kUsePDL, uint32_t kPullUnroll>
|
||||
PULL_KERNEL void nvlink_pull_kernel(const __grid_constant__ NVLinkCommPullParams params) {
|
||||
using namespace device;
|
||||
constexpr uint32_t kVecSize = 16 / sizeof(T); // 16 bytes per vector
|
||||
using vec_t = device::AlignedVector<packed_t<T>, kVecSize / 2>;
|
||||
constexpr uint32_t kNumWarpVecs = kPullUnroll * kWarpThreads;
|
||||
|
||||
// Round-robin chunks to blocks: the global warp index runs block-fastest, so
|
||||
// neighbouring chunks are driven by different CTAs.
|
||||
const auto warp_in_block = threadIdx.x / kWarpThreads;
|
||||
const auto global_warp_id = blockIdx.x + gridDim.x * warp_in_block;
|
||||
const auto lane_id = threadIdx.x % kWarpThreads;
|
||||
const auto num_warps = gridDim.x * (kPullCTASize / kWarpThreads);
|
||||
|
||||
PDLWaitPrimary<kUsePDL>();
|
||||
const auto barrier = distributed::McBarrier{params.sem_local, params.sem_mc, params.world_size, 2};
|
||||
barrier.arrive_relaxed(/*n=*/0);
|
||||
__syncthreads();
|
||||
|
||||
const auto num_whole_chunks = params.num_vecs / kNumWarpVecs;
|
||||
// warp uniform unrolled path, 0 predicate
|
||||
for (auto chunk = global_warp_id; chunk < num_whole_chunks; chunk += num_warps) {
|
||||
vec_t vecs[kPullUnroll];
|
||||
const auto base = chunk * kNumWarpVecs + lane_id;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kPullUnroll; ++i) {
|
||||
const auto vid = base + i * kWarpThreads;
|
||||
if constexpr (kPrim & Primitive::RS) {
|
||||
ptx::ld_multimem_16B(vecs[i], params.input_mc, vid);
|
||||
} else {
|
||||
ptx::ld_global_16B(vecs[i], params.input, vid);
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (kHasResidual) {
|
||||
vec_t residuals[kPullUnroll];
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kPullUnroll; ++i) {
|
||||
residuals[i].load(params.residual, base + i * kWarpThreads);
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kPullUnroll; ++i) {
|
||||
vecs[i] = reduce_vec(vecs[i], residuals[i]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kPullUnroll; ++i) {
|
||||
if constexpr (kPrim & Primitive::AG) {
|
||||
ptx::st_multimem_16B(vecs[i], params.output_mc, base + i * kWarpThreads);
|
||||
} else {
|
||||
ptx::st_global_16B(vecs[i], params.output, base + i * kWarpThreads);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const auto chunk_offset = num_whole_chunks * kNumWarpVecs;
|
||||
const auto global_tid = global_warp_id * kWarpThreads + lane_id;
|
||||
const auto global_threads = num_warps * kWarpThreads;
|
||||
for (auto vid = chunk_offset + global_tid; vid < params.num_vecs; vid += global_threads) {
|
||||
vec_t vec;
|
||||
if constexpr (kPrim & Primitive::RS) {
|
||||
ptx::ld_multimem_16B(vec, params.input_mc, vid);
|
||||
} else {
|
||||
ptx::ld_global_16B(vec, params.input, vid);
|
||||
}
|
||||
if constexpr (kHasResidual) {
|
||||
vec_t res;
|
||||
res.load(params.residual, vid);
|
||||
vec = reduce_vec(vec, res);
|
||||
}
|
||||
if constexpr (kPrim & Primitive::AG) {
|
||||
ptx::st_multimem_16B(vec, params.output_mc, vid);
|
||||
} else {
|
||||
ptx::st_global_16B(vec, params.output, vid);
|
||||
}
|
||||
}
|
||||
|
||||
PDLTriggerSecondary<kUsePDL>();
|
||||
__syncthreads();
|
||||
if constexpr (kPrim & Primitive::AG) {
|
||||
barrier.arrive_rel_acq(/*n=*/1);
|
||||
} else { // no store multimem, only local store
|
||||
barrier.arrive_relaxed(/*n=*/1);
|
||||
}
|
||||
}
|
||||
|
||||
/// Block size for the push kernel: the smallest that still spreads the work
|
||||
/// over every SM, capped at the launch bound.
|
||||
inline auto choose_push_block_size(uint32_t num_vecs) -> uint32_t {
|
||||
static const uint32_t kNumSM = [] {
|
||||
int device = 0;
|
||||
CHECK_CUDA(cudaGetDevice(&device));
|
||||
return host::runtime::get_sm_count(device);
|
||||
}();
|
||||
for (const uint32_t block_size : {128u, 256u, 384u, 512u}) {
|
||||
if (host::div_ceil(num_vecs, block_size) <= kNumSM) return block_size;
|
||||
}
|
||||
return 1024u;
|
||||
}
|
||||
|
||||
template <typename T, bool kUsePDL>
|
||||
struct NVLinkComm {
|
||||
private:
|
||||
using TensorView = tvm::ffi::TensorView;
|
||||
using PushPlaneObj = host::distributed::PushPlaneObj;
|
||||
using PullPlaneObj = host::distributed::PullPlaneObj;
|
||||
using CommunicatorObj = host::distributed::CommunicatorObj;
|
||||
using CommunicatorRef = host::distributed::CommunicatorRef;
|
||||
static constexpr uint32_t kVecBytes = 16;
|
||||
static constexpr uint32_t kVecSize = kVecBytes / sizeof(T);
|
||||
|
||||
public:
|
||||
struct RouteInfo {
|
||||
uint32_t prefix_tokens; // exclusive prefix sum
|
||||
uint32_t num_rank_tokens; // current rank
|
||||
};
|
||||
|
||||
static RouteInfo get_routing(uint32_t num_tokens, uint32_t rank, uint32_t world_size) {
|
||||
const auto avg = num_tokens / world_size;
|
||||
const auto rem = num_tokens % world_size;
|
||||
return {rank * avg + std::min(rank, rem), avg + (rank < rem ? 1 : 0)};
|
||||
}
|
||||
|
||||
struct HostParams {
|
||||
int64_t hidden_size;
|
||||
DLDevice device;
|
||||
};
|
||||
|
||||
/// \brief Base pointer of the residual, shifted onto this rank's slice when
|
||||
/// the caller hands over the whole tensor.
|
||||
///
|
||||
/// A shard-shaped residual passes straight through; a full tensor is sliced
|
||||
/// via `get_routing`, not a uniform stride, so ragged splits keep working.
|
||||
static const void* get_residual_ptr(
|
||||
const tvm::ffi::Optional<TensorView>& residual,
|
||||
uint32_t domain_tokens,
|
||||
uint32_t total_tokens,
|
||||
uint32_t prefix_bytes) {
|
||||
if (!residual.has_value()) return nullptr;
|
||||
const auto tokens = static_cast<uint32_t>(residual.value().size(0));
|
||||
const auto* base = static_cast<const uint8_t*>(residual.value().data_ptr());
|
||||
if (tokens == domain_tokens) return base;
|
||||
CHECK_HOST(tokens == total_tokens) << "residual has " << tokens << " tokens, expected " << domain_tokens
|
||||
<< " (this rank's shard) or " << total_tokens << " (the whole tensor)";
|
||||
return base + prefix_bytes;
|
||||
}
|
||||
|
||||
static HostParams check_params(
|
||||
const TensorView in,
|
||||
const TensorView out,
|
||||
const tvm::ffi::Optional<TensorView>& residual = {},
|
||||
host::DebugInfo info = {}) {
|
||||
using namespace host;
|
||||
auto D = SymbolicSize{"hidden_size"};
|
||||
auto device_ = SymbolicDevice{};
|
||||
auto dtype_ = SymbolicDType{};
|
||||
if constexpr (!std::is_same_v<T, void>) dtype_.set_options<T>();
|
||||
device_.set_options<kDLCUDA>();
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_dtype(dtype_)
|
||||
.with_device(device_)
|
||||
.verify(in, info);
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_dtype(dtype_)
|
||||
.with_device(device_)
|
||||
.verify(out, info);
|
||||
if (residual.has_value()) {
|
||||
TensorMatcher({-1, D}) //
|
||||
.with_dtype(dtype_)
|
||||
.with_device(device_)
|
||||
.verify(residual.value(), info);
|
||||
}
|
||||
return {D.unwrap(), device_.unwrap()};
|
||||
}
|
||||
|
||||
private:
|
||||
template <Primitive kPrim, uint32_t kWorldSize>
|
||||
static void run_push(
|
||||
const PushPlaneObj& push,
|
||||
const TensorView in,
|
||||
const TensorView out,
|
||||
const tvm::ffi::Optional<TensorView> residual) {
|
||||
CHECK_HOST(push.world_size == kWorldSize) << push.world_size << " != " << kWorldSize;
|
||||
const auto [hidden_size, device] = check_params(in, out, residual);
|
||||
const auto rank = push.rank;
|
||||
const auto num_vecs_per_token = static_cast<uint32_t>(hidden_size / kVecSize);
|
||||
const auto num_push_vecs = static_cast<uint32_t>(in.numel() / kVecSize);
|
||||
const auto num_poll_vecs = static_cast<uint32_t>(out.numel() / kVecSize);
|
||||
const auto slot_bytes = static_cast<int64_t>(push.slot_bytes);
|
||||
const auto out_nbytes = static_cast<int64_t>(out.numel() * sizeof(T));
|
||||
const auto num_tokens = static_cast<uint32_t>(in.size(0));
|
||||
const auto out_tokens = static_cast<uint32_t>(out.size(0));
|
||||
const auto total_tokens = kPrim == Primitive::AG ? out_tokens : num_tokens;
|
||||
const auto routing = get_routing(total_tokens, rank, kWorldSize);
|
||||
|
||||
uint32_t dst_offset = 0;
|
||||
if constexpr (kPrim == Primitive::AR) {
|
||||
CHECK_HOST(num_tokens == out_tokens);
|
||||
CHECK_HOST(out_nbytes <= push.slot_bytes);
|
||||
dst_offset = static_cast<uint32_t>(rank * slot_bytes);
|
||||
} else if constexpr (kPrim == Primitive::RS) {
|
||||
CHECK_HOST(out_tokens == routing.num_rank_tokens);
|
||||
CHECK_HOST(out_nbytes <= push.slot_bytes);
|
||||
// dst_offset is not used for this case
|
||||
} else {
|
||||
static_assert(kPrim == Primitive::AG);
|
||||
CHECK_HOST(num_tokens == routing.num_rank_tokens);
|
||||
CHECK_HOST(out_nbytes <= slot_bytes * kWorldSize);
|
||||
dst_offset = routing.prefix_tokens * static_cast<uint32_t>(num_vecs_per_token * kVecBytes);
|
||||
}
|
||||
|
||||
// Slice the whole plane: the kernel reaches every slot, not just this rank's.
|
||||
const auto block_size = choose_push_block_size(std::max(num_push_vecs, num_poll_vecs));
|
||||
CHECK_HOST(num_vecs_per_token > 0) << "fast div-mod rejects a zero divisor";
|
||||
const auto in_tokens_total = static_cast<uint32_t>(in.size(0));
|
||||
// The all-reduce reduces the whole tensor on every rank; the other two work
|
||||
// on this rank's shard, so a full-length residual is sliced.
|
||||
const auto residual_domain = kPrim == Primitive::AR ? total_tokens : routing.num_rank_tokens;
|
||||
const auto residual_ptr = get_residual_ptr(
|
||||
residual,
|
||||
residual_domain,
|
||||
total_tokens,
|
||||
routing.prefix_tokens * static_cast<uint32_t>(num_vecs_per_token * kVecBytes));
|
||||
const auto params = NVLinkCommPushParams<kWorldSize>{
|
||||
.input = in.data_ptr(),
|
||||
.residual = residual_ptr,
|
||||
.output = out.data_ptr(),
|
||||
.dst_offset = dst_offset,
|
||||
.rank = rank,
|
||||
.num_push_vecs = num_push_vecs,
|
||||
.num_poll_vecs = num_poll_vecs,
|
||||
.num_vecs_per_token = num_vecs_per_token,
|
||||
.tokens_avg = in_tokens_total / kWorldSize,
|
||||
.tokens_rem = in_tokens_total % kWorldSize,
|
||||
.vecs_per_token_div = fast_mod_div_u32_t{num_vecs_per_token},
|
||||
.ws = push.get_workspace<kWorldSize>(/*size=*/0),
|
||||
};
|
||||
const auto kernel = residual.has_value() ? nvlink_push_kernel<T, true, kPrim, kWorldSize, kUsePDL>
|
||||
: nvlink_push_kernel<T, false, kPrim, kWorldSize, kUsePDL>;
|
||||
host::LaunchKernel(push.num_blocks, block_size, device).enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
|
||||
template <Primitive kPrim, uint32_t kPullUnroll>
|
||||
static void run_pull(
|
||||
const PullPlaneObj& pull,
|
||||
const TensorView in,
|
||||
const TensorView out,
|
||||
const tvm::ffi::Optional<TensorView> residual,
|
||||
uintptr_t in_mc_ptr,
|
||||
uintptr_t out_mc_ptr,
|
||||
uint32_t num_blocks_hint) {
|
||||
CHECK_HOST(pull.mc_semaphore != nullptr);
|
||||
const auto [hidden_size, device] = check_params(in, out, residual);
|
||||
const auto rank = pull.rank;
|
||||
const auto world_size = pull.world_size;
|
||||
const auto num_tokens = static_cast<uint32_t>(in.size(0));
|
||||
const auto out_tokens = static_cast<uint32_t>(out.size(0));
|
||||
const auto total_tokens = kPrim == Primitive::AG ? out_tokens : num_tokens;
|
||||
const auto routing = get_routing(total_tokens, rank, world_size);
|
||||
const auto num_vecs_per_token = static_cast<uint32_t>(hidden_size / kVecSize);
|
||||
const auto bytes_per_token = static_cast<uint32_t>(num_vecs_per_token * kVecBytes);
|
||||
const auto prefix_bytes = static_cast<uint32_t>(routing.prefix_tokens * bytes_per_token);
|
||||
|
||||
// 0 = no hint, autotune; > 0 always use hint but clip to upper bound
|
||||
if constexpr (kPrim == Primitive::AR) {
|
||||
CHECK_HOST(num_tokens == out_tokens && in_mc_ptr != 0 && out_mc_ptr != 0);
|
||||
in_mc_ptr += prefix_bytes;
|
||||
out_mc_ptr += prefix_bytes;
|
||||
if (num_blocks_hint == 0) num_blocks_hint = host::div_ceil(256u, kPullUnroll * world_size);
|
||||
} else if constexpr (kPrim == Primitive::RS) {
|
||||
CHECK_HOST(out_tokens == routing.num_rank_tokens && in_mc_ptr != 0);
|
||||
in_mc_ptr += prefix_bytes;
|
||||
if (num_blocks_hint == 0) num_blocks_hint = pull.num_blocks; // use all the blocks for RS
|
||||
} else {
|
||||
static_assert(kPrim == Primitive::AG);
|
||||
CHECK_HOST(num_tokens == routing.num_rank_tokens && out_mc_ptr != 0);
|
||||
out_mc_ptr += prefix_bytes;
|
||||
if (num_blocks_hint == 0) num_blocks_hint = host::div_ceil(128u, kPullUnroll * world_size);
|
||||
}
|
||||
/// NOTE: hard limit upper bound is `pull.num_blocks`
|
||||
num_blocks_hint = std::min(num_blocks_hint, pull.num_blocks);
|
||||
|
||||
// Every pull primitive works on this rank's shard, so a full-length
|
||||
// residual is sliced onto it.
|
||||
const auto residual_ptr = get_residual_ptr(residual, routing.num_rank_tokens, total_tokens, prefix_bytes);
|
||||
const auto params = NVLinkCommPullParams{
|
||||
.input = in.data_ptr(),
|
||||
.residual = residual_ptr,
|
||||
.output = out.data_ptr(),
|
||||
.num_vecs = static_cast<uint32_t>(routing.num_rank_tokens * num_vecs_per_token),
|
||||
.input_mc = std::bit_cast<uint8_t*>(in_mc_ptr),
|
||||
.output_mc = std::bit_cast<uint8_t*>(out_mc_ptr),
|
||||
.sem_local = pull.semaphores[rank],
|
||||
.sem_mc = pull.mc_semaphore,
|
||||
.rank = rank,
|
||||
.world_size = pull.world_size,
|
||||
};
|
||||
|
||||
/// NOTE: the final num_blocks resolution must be world unified, otherwise may deadlock
|
||||
const auto max_vecs_in_world = host::div_ceil(total_tokens, world_size) * num_vecs_per_token;
|
||||
const auto max_num_blocks = host::div_ceil(max_vecs_in_world, kPullUnroll * kPullCTASize);
|
||||
const auto num_blocks = std::max(1u, std::min(max_num_blocks, num_blocks_hint));
|
||||
const auto kernel = residual.has_value() ? nvlink_pull_kernel<T, true, kPrim, kUsePDL, kPullUnroll>
|
||||
: nvlink_pull_kernel<T, false, kPrim, kUsePDL, kPullUnroll>;
|
||||
host::LaunchKernel(num_blocks, kPullCTASize, device).enable_pdl(kUsePDL)(kernel, params);
|
||||
}
|
||||
|
||||
public:
|
||||
// specialized for each world size
|
||||
template <uint32_t kWorldSize>
|
||||
static void
|
||||
all_reduce_push(CommunicatorRef comm, TensorView in, TensorView out, tvm::ffi::Optional<TensorView> residual) {
|
||||
return run_push<Primitive::AR, kWorldSize>(comm->get_push_obj(), in, out, residual);
|
||||
}
|
||||
template <uint32_t kWorldSize>
|
||||
static void
|
||||
all_gather_push(CommunicatorRef comm, TensorView in, TensorView out, tvm::ffi::Optional<TensorView> residual) {
|
||||
return run_push<Primitive::AG, kWorldSize>(comm->get_push_obj(), in, out, residual);
|
||||
}
|
||||
template <uint32_t kWorldSize>
|
||||
static void
|
||||
reduce_scatter_push(CommunicatorRef comm, TensorView in, TensorView out, tvm::ffi::Optional<TensorView> residual) {
|
||||
return run_push<Primitive::RS, kWorldSize>(comm->get_push_obj(), in, out, residual);
|
||||
}
|
||||
|
||||
// only compile once for each world size
|
||||
template <uint32_t kPullUnroll>
|
||||
static void all_reduce_pull(
|
||||
CommunicatorRef comm, // only pull is needed
|
||||
TensorView in,
|
||||
TensorView out,
|
||||
tvm::ffi::Optional<TensorView> residual,
|
||||
int64_t in_mc_ptr,
|
||||
int64_t out_mc_ptr,
|
||||
uint32_t num_blocks_hint) {
|
||||
return run_pull<Primitive::AR, kPullUnroll>(
|
||||
comm->get_pull_obj(), in, out, residual, in_mc_ptr, out_mc_ptr, num_blocks_hint);
|
||||
}
|
||||
template <uint32_t kPullUnroll>
|
||||
static void all_gather_pull(
|
||||
CommunicatorRef comm, // only pull is needed
|
||||
TensorView in,
|
||||
TensorView out,
|
||||
tvm::ffi::Optional<TensorView> residual,
|
||||
int64_t out_mc_ptr,
|
||||
uint32_t num_blocks_hint) {
|
||||
return run_pull<Primitive::AG, kPullUnroll>(
|
||||
comm->get_pull_obj(), in, out, residual, 0, out_mc_ptr, num_blocks_hint);
|
||||
}
|
||||
template <uint32_t kPullUnroll>
|
||||
static void reduce_scatter_pull(
|
||||
CommunicatorRef comm, // only pull is needed
|
||||
TensorView in,
|
||||
TensorView out,
|
||||
tvm::ffi::Optional<TensorView> residual,
|
||||
int64_t in_mc_ptr,
|
||||
uint32_t num_blocks_hint) {
|
||||
return run_pull<Primitive::RS, kPullUnroll>(comm->get_pull_obj(), in, out, residual, in_mc_ptr, 0, num_blocks_hint);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace sglang
|
||||
@@ -97,7 +97,7 @@ SGL_DEVICE void ld_multimem_16B(V& x, const void* mc_addr, int64_t vec_offset) {
|
||||
mc_addr = static_cast<const uint8_t*>(mc_addr) + vec_offset * 16;
|
||||
if constexpr (std::is_same_v<V, device::AlignedVector<fp32x2_t, 2>>) {
|
||||
float4 val;
|
||||
asm volatile("multimem.ld_reduce.weak.add.v4.f32 {%0, %1, %2, %3}, [%4];"
|
||||
asm volatile("multimem.ld_reduce.weak.global.add.v4.f32 {%0, %1, %2, %3}, [%4];"
|
||||
: "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w)
|
||||
: "l"(mc_addr));
|
||||
x = *reinterpret_cast<const V*>(&val);
|
||||
@@ -107,12 +107,12 @@ SGL_DEVICE void ld_multimem_16B(V& x, const void* mc_addr, int64_t vec_offset) {
|
||||
// rejects .f32 ("=f") destinations with "Arguments mismatch".
|
||||
uint4 val;
|
||||
if constexpr (std::is_same_v<V, device::AlignedVector<fp16x2_t, 4>>) {
|
||||
asm volatile("multimem.ld_reduce.weak.add.acc::f32.v4.f16x2 {%0, %1, %2, %3}, [%4];"
|
||||
asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.v4.f16x2 {%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
|
||||
: "l"(mc_addr));
|
||||
} else {
|
||||
static_assert(std::is_same_v<V, device::AlignedVector<bf16x2_t, 4>>); // 4x bf16x2
|
||||
asm volatile("multimem.ld_reduce.weak.add.acc::f32.v4.bf16x2 {%0, %1, %2, %3}, [%4];"
|
||||
asm volatile("multimem.ld_reduce.weak.global.add.acc::f32.v4.bf16x2 {%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
|
||||
: "l"(mc_addr));
|
||||
}
|
||||
@@ -150,7 +150,7 @@ SGL_DEVICE void st_multimem_16B(const V& x, void* mc_addr, int64_t vec_offset) {
|
||||
static_assert(alignof(V) == 16 && sizeof(V) == 16);
|
||||
const auto val = *reinterpret_cast<const float4*>(&x);
|
||||
mc_addr = static_cast<uint8_t*>(mc_addr) + vec_offset * 16;
|
||||
asm volatile("multimem.st.weak.v4.f32 [%4], {%0, %1, %2, %3};"
|
||||
asm volatile("multimem.st.weak.global.v4.f32 [%4], {%0, %1, %2, %3};"
|
||||
:
|
||||
: "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w), "l"(mc_addr));
|
||||
#else
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Fused deferred-MoE finalize + 1shot push all-reduce [+ RMSNorm] (bf16).
|
||||
|
||||
One entry point, :func:`moe_finalize_all_reduce`, over
|
||||
``csrc/distributed/all_reduce_fusion.cuh``::
|
||||
|
||||
out[t] = allreduce( sum_k expert_weights[t, k] * gemm2_out[idx[t*top_k + k]]
|
||||
(+ shared_output[t]) ) # then, optionally,
|
||||
out[t] = out[t] * rsqrt(mean(out[t]^2) + eps) * norm_weight
|
||||
|
||||
The rank-local finalize (the trtllm-gen ``do_finalize=False`` triple, see
|
||||
``moe_runner/flashinfer_trtllm.py``) is computed in registers and pushed
|
||||
straight into every peer's CustomAllReduceV2 push slot, so it never
|
||||
materializes; ``idx == -1`` slots (EP: non-local expert) contribute nothing.
|
||||
Small-batch only: the whole ``[T, hidden]`` bf16 row view must fit one push
|
||||
slot (checked C++-side; :func:`fits_push_slot` lets callers pre-check).
|
||||
|
||||
Needs :func:`register_comm` once per process (the CustomAllReduceV2
|
||||
``Communicator``); the ops key on ``world_size`` alone.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
get_jit_cuda_arch,
|
||||
is_arch_support_pdl,
|
||||
is_hip_runtime,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
from sglang.kernels.ops.communication.all_reduce import Communicator
|
||||
|
||||
|
||||
# Storage plane: the CustomAllReduceV2 Communicator (push plane only)
|
||||
|
||||
_COMM_MAP: dict[int, Communicator] = {}
|
||||
|
||||
|
||||
def register_comm(comm: Communicator) -> None:
|
||||
"""Register the CustomAllReduceV2 communicator whose push plane the fused
|
||||
kernel stages through. ``world_size`` is the whole key, so at most one
|
||||
communicator per size may be registered in a process.
|
||||
"""
|
||||
prev = _COMM_MAP.get(comm.world_size)
|
||||
assert prev is None or prev is comm, (
|
||||
f"a different communicator is already registered for world_size="
|
||||
f"{comm.world_size}; these ops key only on world_size, so two groups of "
|
||||
f"the same size cannot coexist in one process"
|
||||
)
|
||||
_COMM_MAP[comm.world_size] = comm
|
||||
|
||||
|
||||
def get_registered_comm(world_size: int) -> Optional[Communicator]:
|
||||
return _COMM_MAP.get(world_size)
|
||||
|
||||
|
||||
# Geometry
|
||||
|
||||
_VEC_ELEMS = 8 # bf16 per 16B vector = per thread
|
||||
_MAX_CLUSTER_SIZE = 8 # portable cluster size limit
|
||||
|
||||
|
||||
def valid_cluster_sizes(hidden_dim: int) -> list[int]:
|
||||
"""Cluster sizes the kernel can be built for at this hidden width: whole
|
||||
16B vectors per row, whole warps per block, <= 1024 threads, <= 8 blocks."""
|
||||
if hidden_dim % _VEC_ELEMS != 0:
|
||||
return []
|
||||
row_vecs = hidden_dim // _VEC_ELEMS
|
||||
return [
|
||||
c
|
||||
for c in range(1, _MAX_CLUSTER_SIZE + 1)
|
||||
if row_vecs % c == 0 and (row_vecs // c) % 32 == 0 and row_vecs // c <= 1024
|
||||
]
|
||||
|
||||
|
||||
@cache_once
|
||||
def default_cluster_size(hidden_dim: int) -> int:
|
||||
if hidden_dim % 1024 == 0 and hidden_dim <= 8192:
|
||||
return hidden_dim // 1024
|
||||
if hidden_dim % 512 == 0 and hidden_dim <= 3584:
|
||||
return hidden_dim // 512
|
||||
candidates = valid_cluster_sizes(hidden_dim)
|
||||
if not candidates:
|
||||
raise ValueError(
|
||||
f"hidden_dim={hidden_dim} has no valid cluster geometry (needs a "
|
||||
f"multiple of {_VEC_ELEMS * 32} bf16)"
|
||||
)
|
||||
# closest to 128 threads per block, larger block on ties
|
||||
return min(candidates, key=lambda c: (abs(hidden_dim // _VEC_ELEMS // c - 128), c))
|
||||
|
||||
|
||||
def fits_push_slot(max_push_size: int, num_tokens: int, hidden_dim: int) -> bool:
|
||||
"""Whether a ``[num_tokens, hidden_dim]`` bf16 row view fits one push slot
|
||||
(``CustomAllReduceV2.max_push_size``)."""
|
||||
return 0 < num_tokens * hidden_dim * 2 <= max_push_size
|
||||
|
||||
|
||||
# JIT module: one per (world_size, hidden_dim, top_k, cluster_size, weight_dtype); the
|
||||
# shared-add and norm variants are compiled into it and picked at call time.
|
||||
|
||||
|
||||
def require_cluster_launch_arch() -> None:
|
||||
if is_hip_runtime() or get_jit_cuda_arch().major < 9:
|
||||
raise RuntimeError(
|
||||
"fused all-reduce cluster kernels require CUDA SM90 or newer"
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_module(
|
||||
world_size: int,
|
||||
hidden_dim: int,
|
||||
top_k: int,
|
||||
cluster_size: int,
|
||||
weight_dtype: torch.dtype,
|
||||
) -> Module:
|
||||
require_cluster_launch_arch()
|
||||
assert cluster_size in valid_cluster_sizes(hidden_dim), (
|
||||
f"cluster_size={cluster_size} is not valid for hidden_dim={hidden_dim}; "
|
||||
f"choose from {valid_cluster_sizes(hidden_dim)}"
|
||||
)
|
||||
args = make_cpp_args(
|
||||
world_size, hidden_dim, top_k, cluster_size, is_arch_support_pdl(), weight_dtype
|
||||
)
|
||||
return load_jit(
|
||||
"moe_finalize_all_reduce",
|
||||
*args,
|
||||
cuda_files=["distributed/all_reduce_fusion.cuh"],
|
||||
cuda_wrappers=[("run", f"MoeFinalizeAllReduceKernel<{args}>::run")],
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["out"])
|
||||
def _moe_finalize_all_reduce_op(
|
||||
world_size: int,
|
||||
hidden_dim: int,
|
||||
top_k: int,
|
||||
cluster_size: int,
|
||||
out: torch.Tensor,
|
||||
gemm2_out: torch.Tensor,
|
||||
expanded_idx_to_permuted_idx: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
shared_output: Optional[torch.Tensor],
|
||||
norm_weight: Optional[torch.Tensor],
|
||||
norm_eps: float,
|
||||
prefetch_metadata: bool,
|
||||
) -> None:
|
||||
comm = _COMM_MAP.get(world_size)
|
||||
assert comm is not None, (
|
||||
f"no communicator registered for world_size={world_size}; call "
|
||||
"all_reduce_fusion.register_comm(comm.obj) first"
|
||||
)
|
||||
_jit_module(world_size, hidden_dim, top_k, cluster_size, expert_weights.dtype).run(
|
||||
comm,
|
||||
out,
|
||||
gemm2_out,
|
||||
expanded_idx_to_permuted_idx,
|
||||
expert_weights,
|
||||
shared_output,
|
||||
norm_weight,
|
||||
norm_eps,
|
||||
prefetch_metadata,
|
||||
)
|
||||
|
||||
|
||||
def moe_finalize_all_reduce(
|
||||
gemm2_out: torch.Tensor,
|
||||
expanded_idx_to_permuted_idx: torch.Tensor,
|
||||
expert_weights: torch.Tensor,
|
||||
top_k: int,
|
||||
shared_output: Optional[torch.Tensor] = None,
|
||||
norm_weight: Optional[torch.Tensor] = None,
|
||||
norm_eps: Optional[float] = None,
|
||||
*,
|
||||
world_size: int,
|
||||
hidden_dim: int,
|
||||
cluster_size: Optional[int] = None,
|
||||
prefetch_metadata: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Deferred MoE finalize [+ shared add] -> 1shot push all-reduce [-> RMSNorm].
|
||||
|
||||
:param gemm2_out: ``[P, hidden_dim]`` bf16, trtllm-gen permuted / padded rows.
|
||||
:param expanded_idx_to_permuted_idx: ``[T * top_k]`` int32, ``-1`` = dropped slot.
|
||||
:param expert_weights: ``[T, top_k]`` bf16 or fp32; any routed scaling factor is
|
||||
already folded in (nothing is rescaled here).
|
||||
:param shared_output: optional ``[T, hidden_dim]`` bf16 added before the reduce.
|
||||
:param norm_weight: optional ``[hidden_dim]`` bf16 RMSNorm weight; with
|
||||
``norm_eps`` it turns on the fused norm epilogue.
|
||||
:param prefetch_metadata: read the plane's phase counter and the routing
|
||||
metadata before the PDL wait; valid only when the
|
||||
preceding kernel is not an all-reduce on the same
|
||||
plane and the producers of
|
||||
``expanded_idx_to_permuted_idx`` /
|
||||
``expert_weights`` are complete. Defaults to False.
|
||||
:returns: a new ``[T, hidden_dim]`` bf16 tensor (not in place).
|
||||
"""
|
||||
num_tokens = expert_weights.shape[0]
|
||||
assert expert_weights.dtype in (torch.bfloat16, torch.float32)
|
||||
assert expert_weights.shape[1] == top_k, (expert_weights.shape, top_k)
|
||||
assert (norm_weight is None) == (norm_eps is None), (
|
||||
"norm_weight and norm_eps must be given together"
|
||||
)
|
||||
out = torch.empty(
|
||||
num_tokens, hidden_dim, dtype=torch.bfloat16, device=gemm2_out.device
|
||||
)
|
||||
if num_tokens == 0: # nothing staged: no phase flip on any rank, stays in step
|
||||
return out
|
||||
_moe_finalize_all_reduce_op(
|
||||
world_size,
|
||||
hidden_dim,
|
||||
top_k,
|
||||
cluster_size or default_cluster_size(hidden_dim),
|
||||
out,
|
||||
gemm2_out,
|
||||
expanded_idx_to_permuted_idx,
|
||||
expert_weights,
|
||||
shared_output,
|
||||
norm_weight,
|
||||
float(norm_eps) if norm_eps is not None else 0.0,
|
||||
prefetch_metadata,
|
||||
)
|
||||
return out
|
||||
@@ -0,0 +1,361 @@
|
||||
"""MoE finalize and TP all-reduce with an HC=4 post-mixing epilogue."""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.ops.communication.all_reduce_fusion import (
|
||||
default_cluster_size,
|
||||
get_registered_comm,
|
||||
require_cluster_launch_arch,
|
||||
)
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
# The kernel is built for this width only (static_assert in all_reduce_fusion.cuh).
|
||||
_MHC_HIDDEN_DIM = 5120
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_mhc_module(world_size, top_k, cluster_size, weight_dtype):
|
||||
require_cluster_launch_arch()
|
||||
args = make_cpp_args(
|
||||
world_size,
|
||||
_MHC_HIDDEN_DIM,
|
||||
top_k,
|
||||
cluster_size,
|
||||
is_arch_support_pdl(),
|
||||
weight_dtype,
|
||||
True,
|
||||
)
|
||||
return load_jit(
|
||||
"moe_finalize_all_reduce_mhc",
|
||||
*args,
|
||||
cuda_files=["distributed/all_reduce_fusion.cuh"],
|
||||
cuda_wrappers=[
|
||||
("run", f"MoeFinalizeAllReduceKernel<{args}>::run_mhc"),
|
||||
("run_norm", f"MoeFinalizeAllReduceKernel<{args}>::run_mhc_norm"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["out", "mhc_out"])
|
||||
def _moe_finalize_all_reduce_mhc_op(
|
||||
world_size: int,
|
||||
top_k: int,
|
||||
cluster_size: int,
|
||||
out: torch.Tensor,
|
||||
mhc_out: torch.Tensor,
|
||||
gemm2: torch.Tensor,
|
||||
idx: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
shared: Optional[torch.Tensor],
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
) -> None:
|
||||
comm = get_registered_comm(world_size)
|
||||
assert comm is not None
|
||||
_jit_mhc_module(world_size, top_k, cluster_size, weights.dtype).run(
|
||||
comm,
|
||||
out,
|
||||
gemm2,
|
||||
idx,
|
||||
weights,
|
||||
shared,
|
||||
mhc_out,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
)
|
||||
|
||||
|
||||
def moe_finalize_all_reduce_mhc(
|
||||
gemm2: torch.Tensor,
|
||||
idx: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
top_k: int,
|
||||
shared: Optional[torch.Tensor],
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
*,
|
||||
world_size: int,
|
||||
cluster_size: Optional[int] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Deferred MoE finalize -> push all-reduce -> HC=4 post mixing.
|
||||
|
||||
Same finalize inputs as :func:`all_reduce_fusion.moe_finalize_all_reduce`;
|
||||
``residual`` is ``[T, 4, hidden]`` bf16, ``post`` ``[T, 4]`` and ``comb``
|
||||
``[T, 4, 4]`` fp32. Returns ``(reduced [T, hidden], mhc_out [T, 4, hidden])``.
|
||||
"""
|
||||
out = torch.empty(
|
||||
(weights.shape[0], _MHC_HIDDEN_DIM), dtype=torch.bfloat16, device=gemm2.device
|
||||
)
|
||||
mhc_out = torch.empty_like(residual)
|
||||
if weights.shape[0]:
|
||||
_moe_finalize_all_reduce_mhc_op(
|
||||
world_size,
|
||||
top_k,
|
||||
cluster_size or default_cluster_size(_MHC_HIDDEN_DIM),
|
||||
out,
|
||||
mhc_out,
|
||||
gemm2,
|
||||
idx,
|
||||
weights,
|
||||
shared,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
)
|
||||
return out, mhc_out
|
||||
|
||||
|
||||
@register_custom_op(mutates_args=["out", "mhc_out", "normalized"])
|
||||
def _moe_finalize_all_reduce_mhc_norm_op(
|
||||
world_size: int,
|
||||
top_k: int,
|
||||
cluster_size: int,
|
||||
out: torch.Tensor,
|
||||
mhc_out: torch.Tensor,
|
||||
normalized: torch.Tensor,
|
||||
gemm2: torch.Tensor,
|
||||
idx: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
shared: Optional[torch.Tensor],
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
pre: torch.Tensor,
|
||||
norm_weight: torch.Tensor,
|
||||
eps: float,
|
||||
) -> None:
|
||||
comm = get_registered_comm(world_size)
|
||||
assert comm is not None
|
||||
_jit_mhc_module(world_size, top_k, cluster_size, weights.dtype).run_norm(
|
||||
comm,
|
||||
out,
|
||||
gemm2,
|
||||
idx,
|
||||
weights,
|
||||
shared,
|
||||
mhc_out,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
norm_weight,
|
||||
eps,
|
||||
normalized,
|
||||
)
|
||||
|
||||
|
||||
def moe_finalize_all_reduce_mhc_norm(
|
||||
gemm2: torch.Tensor,
|
||||
idx: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
top_k: int,
|
||||
shared: Optional[torch.Tensor],
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
pre: torch.Tensor,
|
||||
norm_weight: torch.Tensor,
|
||||
eps: float,
|
||||
*,
|
||||
world_size: int,
|
||||
cluster_size: Optional[int] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
""":func:`moe_finalize_all_reduce_mhc` plus the HC=4 pre-collapse
|
||||
(``pre`` ``[T, 4]`` fp32) and RMSNorm of the collapsed row. Returns
|
||||
``(reduced, mhc_out, normalized [T, hidden])``.
|
||||
"""
|
||||
out = torch.empty(
|
||||
(weights.shape[0], _MHC_HIDDEN_DIM), dtype=torch.bfloat16, device=gemm2.device
|
||||
)
|
||||
mhc_out = torch.empty_like(residual)
|
||||
normalized = torch.empty_like(out)
|
||||
if weights.shape[0]:
|
||||
_moe_finalize_all_reduce_mhc_norm_op(
|
||||
world_size,
|
||||
top_k,
|
||||
cluster_size or default_cluster_size(_MHC_HIDDEN_DIM),
|
||||
out,
|
||||
mhc_out,
|
||||
normalized,
|
||||
gemm2,
|
||||
idx,
|
||||
weights,
|
||||
shared,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
norm_weight,
|
||||
eps,
|
||||
)
|
||||
return out, mhc_out, normalized
|
||||
|
||||
|
||||
@cache_once
|
||||
def _identity_routing(rows, device):
|
||||
return (
|
||||
torch.arange(rows, device=device, dtype=torch.int32),
|
||||
torch.ones(rows, 1, device=device, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
def all_reduce_mhc_norm(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
pre: torch.Tensor,
|
||||
norm_weight: torch.Tensor,
|
||||
eps: float,
|
||||
*,
|
||||
world_size: int,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Plain all-reduce of ``x`` with the mHC + norm epilogue: the finalize
|
||||
kernel driven with identity routing (top_k = 1, unit weights)."""
|
||||
idx, weights = _identity_routing(x.shape[0], x.device)
|
||||
return moe_finalize_all_reduce_mhc_norm(
|
||||
x,
|
||||
idx,
|
||||
weights,
|
||||
1,
|
||||
None,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
norm_weight,
|
||||
eps,
|
||||
world_size=world_size,
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_mhc_quant_module(world_size, top_k, cluster_size, weight_dtype):
|
||||
require_cluster_launch_arch()
|
||||
args = make_cpp_args(
|
||||
world_size,
|
||||
_MHC_HIDDEN_DIM,
|
||||
top_k,
|
||||
cluster_size,
|
||||
is_arch_support_pdl(),
|
||||
weight_dtype,
|
||||
True,
|
||||
True,
|
||||
)
|
||||
return load_jit(
|
||||
"moe_finalize_all_reduce_mhc_quant",
|
||||
*args,
|
||||
cuda_files=["distributed/all_reduce_fusion.cuh"],
|
||||
cuda_wrappers=[("run", f"MoeFinalizeAllReduceKernel<{args}>::run_mhc_quant")],
|
||||
)
|
||||
|
||||
|
||||
@register_custom_op(
|
||||
mutates_args=["out", "mhc_out", "normalized", "quantized", "scales"]
|
||||
)
|
||||
def _moe_finalize_all_reduce_mhc_quant_op(
|
||||
world_size: int,
|
||||
top_k: int,
|
||||
cluster_size: int,
|
||||
out: torch.Tensor,
|
||||
mhc_out: torch.Tensor,
|
||||
normalized: torch.Tensor,
|
||||
quantized: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
gemm2: torch.Tensor,
|
||||
idx: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
shared: Optional[torch.Tensor],
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
pre: torch.Tensor,
|
||||
norm_weight: torch.Tensor,
|
||||
eps: float,
|
||||
) -> None:
|
||||
comm = get_registered_comm(world_size)
|
||||
assert comm is not None
|
||||
_jit_mhc_quant_module(world_size, top_k, cluster_size, weights.dtype).run(
|
||||
comm,
|
||||
out,
|
||||
gemm2,
|
||||
idx,
|
||||
weights,
|
||||
shared,
|
||||
mhc_out,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
norm_weight,
|
||||
eps,
|
||||
normalized,
|
||||
quantized,
|
||||
scales,
|
||||
)
|
||||
|
||||
|
||||
def moe_finalize_all_reduce_mhc_quant(
|
||||
gemm2: torch.Tensor,
|
||||
idx: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
top_k: int,
|
||||
shared: Optional[torch.Tensor],
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
comb: torch.Tensor,
|
||||
pre: torch.Tensor,
|
||||
norm_weight: torch.Tensor,
|
||||
eps: float,
|
||||
*,
|
||||
world_size: int,
|
||||
cluster_size: Optional[int] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
""":func:`moe_finalize_all_reduce_mhc_norm` plus fp8 e4m3 quantization of
|
||||
the normalized row with ue8m0 group scales (rows <= 8). Returns
|
||||
``(reduced, mhc_out, normalized, quantized, scales)``.
|
||||
"""
|
||||
rows = weights.shape[0]
|
||||
assert 0 < rows <= 8
|
||||
out = torch.empty(
|
||||
(rows, _MHC_HIDDEN_DIM), dtype=torch.bfloat16, device=gemm2.device
|
||||
)
|
||||
mhc_out = torch.empty_like(residual)
|
||||
normalized = torch.empty_like(out)
|
||||
quantized = torch.empty_like(out, dtype=torch.float8_e4m3fn)
|
||||
# ue8m0 scale layout: one byte per 32-wide group, rows padded to 128
|
||||
scales = torch.empty(
|
||||
(_MHC_HIDDEN_DIM // 32) * 128, device=gemm2.device, dtype=torch.uint8
|
||||
)
|
||||
_moe_finalize_all_reduce_mhc_quant_op(
|
||||
world_size,
|
||||
top_k,
|
||||
cluster_size or default_cluster_size(_MHC_HIDDEN_DIM),
|
||||
out,
|
||||
mhc_out,
|
||||
normalized,
|
||||
quantized,
|
||||
scales,
|
||||
gemm2,
|
||||
idx,
|
||||
weights,
|
||||
shared,
|
||||
residual,
|
||||
post,
|
||||
comb,
|
||||
pre,
|
||||
norm_weight,
|
||||
eps,
|
||||
)
|
||||
return out, mhc_out, normalized, quantized, scales
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi import Module
|
||||
|
||||
from sglang.kernels.ops.communication.all_reduce import Communicator
|
||||
|
||||
|
||||
_PRIMITIVES: Final = ["all_reduce", "all_gather", "reduce_scatter"]
|
||||
|
||||
|
||||
def get_multicast_ptr(tensor: torch.Tensor) -> int:
|
||||
"""Multicast alias of a symmetric-memory tensor. Collective on first call;
|
||||
torch caches the handle per allocation, so repeats stay cheap.
|
||||
"""
|
||||
from torch._C._distributed_c10d import _SymmetricMemory
|
||||
|
||||
ptr = _SymmetricMemory.rendezvous(tensor).multicast_ptr
|
||||
assert ptr != 0, "tensor has no multicast alias; was it allocated p2p?"
|
||||
return ptr
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_pull_module(dtype: torch.dtype, num_unroll: int) -> Module:
|
||||
args = make_cpp_args(dtype, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"nvlink_comm_pull",
|
||||
*args,
|
||||
f"unroll{num_unroll}",
|
||||
cuda_files=["distributed/nvlink_comm.cuh"],
|
||||
cuda_wrappers=[
|
||||
(n, f"NVLinkComm<{args}>::{n}_pull<{num_unroll}>") for n in _PRIMITIVES
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@cache_once
|
||||
def _jit_push_module(dtype: torch.dtype, world_size: int) -> Module:
|
||||
args = make_cpp_args(dtype, is_arch_support_pdl())
|
||||
return load_jit(
|
||||
"nvlink_comm_push",
|
||||
*args,
|
||||
f"world{world_size}",
|
||||
cuda_files=["distributed/nvlink_comm.cuh"],
|
||||
cuda_wrappers=[
|
||||
(n, f"NVLinkComm<{args}>::{n}_push<{world_size}>") for n in _PRIMITIVES
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# `residual` on any of these is folded into the reduction; it may be shaped like
|
||||
# this rank's shard or like the whole tensor, of which this rank's slice is taken.
|
||||
def all_reduce_push(
|
||||
comm: Communicator,
|
||||
input: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
_jit_push_module(input.dtype, comm.world_size).all_reduce(
|
||||
comm, input, output, residual
|
||||
)
|
||||
|
||||
|
||||
def all_gather_push(
|
||||
comm: Communicator,
|
||||
input: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
_jit_push_module(input.dtype, comm.world_size).all_gather(
|
||||
comm, input, output, residual
|
||||
)
|
||||
|
||||
|
||||
def reduce_scatter_push(
|
||||
comm: Communicator,
|
||||
input: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
) -> None:
|
||||
_jit_push_module(input.dtype, comm.world_size).reduce_scatter(
|
||||
comm, input, output, residual
|
||||
)
|
||||
|
||||
|
||||
def all_reduce_pull(
|
||||
comm: Communicator,
|
||||
input: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
*,
|
||||
in_mc_ptr: int = 0,
|
||||
out_mc_ptr: int = 0,
|
||||
num_unroll=4,
|
||||
num_blocks_hint: int = 0,
|
||||
) -> None:
|
||||
_jit_pull_module(input.dtype, num_unroll).all_reduce(
|
||||
comm,
|
||||
input,
|
||||
output,
|
||||
residual,
|
||||
in_mc_ptr or get_multicast_ptr(input),
|
||||
out_mc_ptr or get_multicast_ptr(output),
|
||||
num_blocks_hint,
|
||||
)
|
||||
|
||||
|
||||
def all_gather_pull(
|
||||
comm: Communicator,
|
||||
input: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
*,
|
||||
out_mc_ptr: int = 0,
|
||||
num_unroll=4,
|
||||
num_blocks_hint: int = 0,
|
||||
) -> None:
|
||||
_jit_pull_module(input.dtype, num_unroll).all_gather(
|
||||
comm,
|
||||
input,
|
||||
output,
|
||||
residual,
|
||||
out_mc_ptr or get_multicast_ptr(output),
|
||||
num_blocks_hint,
|
||||
)
|
||||
|
||||
|
||||
def reduce_scatter_pull(
|
||||
comm: Communicator,
|
||||
input: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
residual: torch.Tensor | None = None,
|
||||
*,
|
||||
in_mc_ptr: int = 0,
|
||||
num_unroll=4,
|
||||
num_blocks_hint: int = 0,
|
||||
) -> None:
|
||||
_jit_pull_module(input.dtype, num_unroll).reduce_scatter(
|
||||
comm,
|
||||
input,
|
||||
output,
|
||||
residual,
|
||||
in_mc_ptr or get_multicast_ptr(input),
|
||||
num_blocks_hint,
|
||||
)
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Exact greedy selection from TP-local base logits and rounded Markov bias."""
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _sharded_greedy_partial_kernel(
|
||||
B,
|
||||
X,
|
||||
P,
|
||||
BS: tl.constexpr,
|
||||
XS: tl.constexpr,
|
||||
WIDTH: tl.constexpr,
|
||||
OFFSET: tl.constexpr,
|
||||
PARTS: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
row, part = tl.program_id(0), tl.program_id(1)
|
||||
i = part * BLOCK + tl.arange(0, BLOCK)
|
||||
valid = i < WIDTH
|
||||
bias = tl.load(B + row * BS + i, valid, 0).to(tl.float32)
|
||||
base = tl.load(X + row * XS + i, valid, 0).to(tl.float32)
|
||||
value = base + bias
|
||||
nan = valid & (value != value)
|
||||
has_nan = tl.sum(nan.to(tl.int32), 0) > 0
|
||||
maximum = tl.max(tl.where(valid & ~nan, value, -float("inf")), 0)
|
||||
wins = tl.where(has_nan, nan, valid & (value == maximum))
|
||||
idx = tl.min(tl.where(wins, i + OFFSET, 2147483647), 0)
|
||||
maximum = tl.where(has_nan, float("nan"), maximum)
|
||||
tl.store(P + (row * PARTS + part) * 2, maximum)
|
||||
tl.store(P + (row * PARTS + part) * 2 + 1, idx.to(tl.float32, bitcast=True))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _sharded_greedy_finish_kernel(
|
||||
P,
|
||||
OUT,
|
||||
BS: tl.constexpr,
|
||||
PARTS: tl.constexpr,
|
||||
WORLD: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0)
|
||||
i = tl.arange(0, BLOCK)
|
||||
valid = i < PARTS * WORLD
|
||||
rank, part = i // PARTS, i % PARTS
|
||||
offset = ((rank * BS + row) * PARTS + part) * 2
|
||||
value = tl.load(P + offset, valid, -float("inf"))
|
||||
idx = tl.load(P + offset + 1, valid, 0).to(tl.int32, bitcast=True)
|
||||
# The NVLink push transport changes +0 to -0 as its Lamport sentinel.
|
||||
# Indices are nonnegative int32 bits, so strip that sign bit to recover ID 0.
|
||||
idx &= 2147483647
|
||||
valid &= idx != 2147483647
|
||||
nan = valid & (value != value)
|
||||
has_nan = tl.sum(nan.to(tl.int32), 0) > 0
|
||||
maximum = tl.max(tl.where(valid & ~nan, value, -float("inf")), 0)
|
||||
wins = tl.where(has_nan, nan, valid & (value == maximum))
|
||||
result = tl.min(tl.where(wins, idx, 2147483647), 0)
|
||||
tl.store(OUT + row, result.to(tl.int64))
|
||||
|
||||
|
||||
def sharded_greedy_step(bias, base_local, *, group, vocab_start, gather=None):
|
||||
"""Fused BuildStepLocal + vocab gather + argmax, without materializing logits.
|
||||
|
||||
Equivalent to argmax of rank-ordered ``build_step_local``/all_gather over the
|
||||
sharded vocab, excluding padding. The transport carries one (value,
|
||||
global-index-bits) pair per 4096-wide block of the shard per row; indices
|
||||
move as bits and are never converted numerically to float.
|
||||
"""
|
||||
assert bias.ndim == base_local.ndim == 2
|
||||
assert bias.shape[0] == base_local.shape[0]
|
||||
assert bias.shape[1] <= base_local.shape[1]
|
||||
assert bias.stride(1) == base_local.stride(1) == 1
|
||||
rows, width = bias.shape
|
||||
block = 4096
|
||||
parts = triton.cdiv(base_local.shape[1], block)
|
||||
assert parts > 0
|
||||
partial = torch.empty((rows, parts, 2), device=bias.device, dtype=torch.float32)
|
||||
_sharded_greedy_partial_kernel[(rows, parts)](
|
||||
bias,
|
||||
base_local,
|
||||
partial,
|
||||
bias.stride(0),
|
||||
base_local.stride(0),
|
||||
width,
|
||||
vocab_start,
|
||||
parts,
|
||||
block,
|
||||
num_warps=4,
|
||||
)
|
||||
# The padded partition width fixes the transport shape on all ranks;
|
||||
# WIDTH masks real entries, including a completely empty final shard.
|
||||
if gather is not None:
|
||||
gathered = gather(partial.view(rows, parts * 2))
|
||||
else:
|
||||
gathered = group.all_gather(partial, dim=0) if group.world_size > 1 else partial
|
||||
result = torch.empty(rows, device=bias.device, dtype=torch.int64)
|
||||
_sharded_greedy_finish_kernel[(rows,)](
|
||||
gathered,
|
||||
result,
|
||||
rows,
|
||||
parts,
|
||||
group.world_size,
|
||||
triton.next_power_of_2(parts * group.world_size),
|
||||
num_warps=4,
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,226 @@
|
||||
"""All-gather of a vocab-parallel row block across a TP group.
|
||||
|
||||
Every implementation takes this rank's ``[rows, local_width]`` slice and returns
|
||||
``[rows, world_size * local_width]`` with the ranks' slices side by side, the
|
||||
layout ``GroupCoordinator.all_gather(dim=-1)`` produces.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VocabGather(ABC):
|
||||
"""``[rows, local] -> [rows, world_size * local]``, ranks side by side."""
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, local: torch.Tensor) -> torch.Tensor: ...
|
||||
|
||||
@abstractmethod
|
||||
def gather_stacked(self, local: torch.Tensor) -> torch.Tensor:
|
||||
"""Gather compact row blocks as [world_size * rows, local_width]."""
|
||||
...
|
||||
|
||||
|
||||
class LocalVocabGather(VocabGather):
|
||||
"""A group of one: the slice is the whole row."""
|
||||
|
||||
def __call__(self, local: torch.Tensor) -> torch.Tensor:
|
||||
return local
|
||||
|
||||
def gather_stacked(self, local: torch.Tensor) -> torch.Tensor:
|
||||
return local
|
||||
|
||||
|
||||
class NcclVocabGather(VocabGather):
|
||||
"""The group coordinator's all_gather along the last dim (NCCL ring)."""
|
||||
|
||||
def __init__(self, group) -> None:
|
||||
self.group = group
|
||||
|
||||
def __call__(self, local: torch.Tensor) -> torch.Tensor:
|
||||
return self.group.all_gather(local, dim=-1)
|
||||
|
||||
def gather_stacked(self, local: torch.Tensor) -> torch.Tensor:
|
||||
return self.group.all_gather(local, dim=0)
|
||||
|
||||
|
||||
# Collective: every rank of the group must call this, in the same order, outside
|
||||
# CUDA-graph capture. The returned multicast alias is 0 when the group has none.
|
||||
def _alloc_symm(
|
||||
group, shape: Tuple[int, int], dtype: torch.dtype
|
||||
) -> Tuple[torch.Tensor, int]:
|
||||
from torch._C._distributed_c10d import _SymmetricMemory
|
||||
|
||||
# a GroupCoordinator names the allocation by its cpu_group, as
|
||||
# CustomAllReduceV2 does; a torch process group names it itself
|
||||
pg = getattr(group, "cpu_group", group)
|
||||
buf = _SymmetricMemory.empty_strided_p2p(
|
||||
(shape[0] * shape[1],),
|
||||
[1],
|
||||
dtype,
|
||||
torch.device("cuda", torch.cuda.current_device()),
|
||||
pg.group_name,
|
||||
)
|
||||
mc_ptr = int(_SymmetricMemory.rendezvous(buf).multicast_ptr)
|
||||
return buf.view(shape), mc_ptr
|
||||
|
||||
|
||||
class NVLinkVocabGather(VocabGather):
|
||||
"""The NVLink collectives on CustomAllReduceV2's multicast plane.
|
||||
|
||||
A slice that fits one slot of the push plane takes the push kernel into a
|
||||
fresh tensor; a larger one that fits ``pull_out`` takes the pull kernel into
|
||||
that symmetric-memory output, which is reused every call; anything else goes
|
||||
to ``fallback``, the NCCL ring. Both kernels gather along the row axis, so
|
||||
the ranks come back stacked and are transposed into place. ``pull_out`` is
|
||||
allocated here: the allocation is collective and captured graphs keep its
|
||||
address.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
ca_comm,
|
||||
group,
|
||||
local_width: int,
|
||||
dtype: torch.dtype,
|
||||
symm_rows: int,
|
||||
fallback: VocabGather,
|
||||
) -> None:
|
||||
self.comm = ca_comm.obj
|
||||
self.world_size = int(group.world_size)
|
||||
self.slot_bytes = int(ca_comm.max_push_size)
|
||||
self.fallback = fallback
|
||||
self.pull_out: Optional[torch.Tensor] = None
|
||||
self.pull_mc_ptr = 0
|
||||
if symm_rows > 0 and self.comm.pull is not None:
|
||||
out, mc_ptr = _alloc_symm(
|
||||
group, (self.world_size * symm_rows, local_width), dtype
|
||||
)
|
||||
if mc_ptr != 0:
|
||||
self.pull_out, self.pull_mc_ptr = out, mc_ptr
|
||||
logger.info(
|
||||
"NVLink vocab gather: pull output %s (%d MB)",
|
||||
tuple(out.shape),
|
||||
out.numel() * out.element_size() >> 20,
|
||||
)
|
||||
else:
|
||||
logger.warning("NVLink vocab gather: no multicast alias, pull path off")
|
||||
|
||||
def __call__(self, local: torch.Tensor) -> torch.Tensor:
|
||||
rows = local.shape[0]
|
||||
if local.nbytes <= self.slot_bytes:
|
||||
return self._push(local)
|
||||
total_rows = self.world_size * rows
|
||||
if self.pull_out is not None and total_rows <= self.pull_out.shape[0]:
|
||||
return self._pull(local, self.pull_out[:total_rows])
|
||||
return self.fallback(local)
|
||||
|
||||
def gather_stacked(self, local: torch.Tensor) -> torch.Tensor:
|
||||
# Compact argmax partials need rank-major output, no symmetric pull buffer.
|
||||
if (
|
||||
local.is_contiguous()
|
||||
and local.shape[1] * local.element_size() % 16 == 0
|
||||
and local.nbytes <= self.slot_bytes
|
||||
):
|
||||
return self._push_stacked(local)
|
||||
return self.fallback.gather_stacked(local)
|
||||
|
||||
def _push(self, local: torch.Tensor) -> torch.Tensor:
|
||||
return self._unstack(self._push_stacked(local))
|
||||
|
||||
def _push_stacked(self, local: torch.Tensor) -> torch.Tensor:
|
||||
from sglang.kernels.ops.communication import nvlink_comm
|
||||
|
||||
rows, width = local.shape
|
||||
gathered = torch.empty(
|
||||
(self.world_size * rows, width), dtype=local.dtype, device=local.device
|
||||
)
|
||||
nvlink_comm.all_gather_push(self.comm, local, gathered)
|
||||
return gathered
|
||||
|
||||
def _pull(self, local: torch.Tensor, out: torch.Tensor) -> torch.Tensor:
|
||||
from sglang.kernels.ops.communication import nvlink_comm
|
||||
|
||||
nvlink_comm.all_gather_pull(self.comm, local, out, out_mc_ptr=self.pull_mc_ptr)
|
||||
full = self._unstack(out)
|
||||
# the transpose copies except at one row, where it would alias the
|
||||
# shared buffer that the next call overwrites
|
||||
return full.clone() if full.data_ptr() == out.data_ptr() else full
|
||||
|
||||
def _unstack(self, gathered: torch.Tensor) -> torch.Tensor:
|
||||
rows = gathered.shape[0] // self.world_size
|
||||
width = gathered.shape[1]
|
||||
if rows == 1:
|
||||
return gathered.view(1, self.world_size * width)
|
||||
return (
|
||||
gathered.view(self.world_size, rows, width)
|
||||
.transpose(0, 1)
|
||||
.reshape(rows, self.world_size * width)
|
||||
)
|
||||
|
||||
|
||||
def _nvlink_ca_comm(group, *, local_width: int, dtype: torch.dtype):
|
||||
ca_comm = getattr(group, "ca_comm", None)
|
||||
if ca_comm is None or getattr(ca_comm, "disabled", True):
|
||||
return None
|
||||
comm = getattr(ca_comm, "obj", None)
|
||||
if comm is None or not getattr(ca_comm, "has_multicast", False):
|
||||
return None
|
||||
if comm.push is None or comm.world_size != group.world_size:
|
||||
return None
|
||||
# the kernels move 16-byte vectors along the row
|
||||
if local_width % (128 // torch.finfo(dtype).bits) != 0:
|
||||
return None
|
||||
return ca_comm
|
||||
|
||||
|
||||
def _default_symm_rows() -> int:
|
||||
try:
|
||||
from sglang.srt.runtime_context import get_exec, get_schedule
|
||||
|
||||
return int(
|
||||
get_schedule().max_running_requests
|
||||
or get_exec().graph.cuda_graph_config.decode.max_bs
|
||||
or 0
|
||||
)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
|
||||
def make_vocab_gather(
|
||||
group,
|
||||
*,
|
||||
local_width: int,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
prefer_nvlink: bool = True,
|
||||
symm_rows: Optional[int] = None,
|
||||
) -> VocabGather:
|
||||
"""The gather for ``group``: local for a group of one, NVLink when the
|
||||
group's custom all-reduce has a multicast plane (and ``prefer_nvlink``),
|
||||
the NCCL ring otherwise. ``symm_rows`` is the row capacity of the NVLink
|
||||
gather's symmetric-memory output for slices past the push slot; None sizes
|
||||
it for the server's largest batch, 0 leaves those slices to NCCL."""
|
||||
if group is None or group.world_size == 1:
|
||||
return LocalVocabGather()
|
||||
nccl = NcclVocabGather(group)
|
||||
if not prefer_nvlink:
|
||||
return nccl
|
||||
ca_comm = _nvlink_ca_comm(group, local_width=local_width, dtype=dtype)
|
||||
if ca_comm is None:
|
||||
return nccl
|
||||
return NVLinkVocabGather(
|
||||
ca_comm=ca_comm,
|
||||
group=group,
|
||||
local_width=local_width,
|
||||
dtype=dtype,
|
||||
symm_rows=_default_symm_rows() if symm_rows is None else symm_rows,
|
||||
fallback=nccl,
|
||||
)
|
||||
Reference in New Issue
Block a user