minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[kernel] Split the custom all-reduce communicator into push/pull planes (#35735)

Browse Source
This commit is contained in:
DarkSharpness
2026-08-26 17:40:28 +08:00
committed by GitHub
parent 58ecbba0bd
commit 689ade69d1
28 changed files with 1881 additions and 1580 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/distributed/communicator.cuh
-117
View File
@@ -1,117 +0,0 @@
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/distributed/communicator.cuh>
#include <tvm/ffi/extra/stl.h>
#include <tvm/ffi/reflection/registry.h>
#include <cstdint>
#include <map>
#include <optional>
#include <string>
#include <vector>
namespace sglang {
namespace host::distributed {
inline CommunicatorObj::CommunicatorObj(
const uint32_t rank,
const uint32_t world_size,
std::vector<TensorView> push_workspaces,
std::vector<TensorView> pull_workspaces,
std::vector<TensorView> pull_semaphores,
TensorView push_counter,
const std::optional<int64_t> pull_mc_workspace_ptr) {
this->rank = rank;
this->world_size = world_size;
RuntimeCheck(1 < world_size && world_size <= kMaxWorldSize, "Invalid world size: ", world_size);
RuntimeCheck(rank < world_size, "Invalid rank: ", rank);
RuntimeCheck(push_workspaces.size() == world_size, "Bad push workspace count");
RuntimeCheck(pull_workspaces.size() == world_size, "Bad pull workspace count");
RuntimeCheck(pull_semaphores.size() == world_size, "Bad pull semaphore count");
// Shared symbolic sizes / device enforce consistency across ranks; the
// matchers also require contiguity (no strides given) and uint8 dtype.
auto push_bytes = SymbolicSize{"push_bytes"};
auto pull_bytes = SymbolicSize{"pull_bytes"};
auto num_pull_blocks = SymbolicSize{"num_pull_blocks"};
auto num_push_blocks = SymbolicSize{"num_push_blocks"};
auto device = SymbolicDevice{};
device.set_options<kDLCUDA>();
for (uint32_t i = 0; i < world_size; ++i) {
TensorMatcher({2 * world_size, push_bytes}).with_dtype<uint8_t>().with_device(device).verify(push_workspaces[i]);
TensorMatcher({pull_bytes}) //
.with_dtype<uint8_t>()
.with_device(device)
.verify(pull_workspaces[i]);
TensorMatcher({num_pull_blocks, static_cast<int64_t>(sizeof(Semaphore))})
.with_dtype<uint8_t>()
.with_device(device)
.verify(pull_semaphores[i]);
this->push_workspaces[i] = static_cast<uint8_t*>(push_workspaces[i].data_ptr());
this->pull_workspaces[i] = static_cast<uint8_t*>(pull_workspaces[i].data_ptr());
this->pull_semaphores[i] = static_cast<Semaphore*>(pull_semaphores[i].data_ptr());
}
TensorMatcher({num_push_blocks, static_cast<int64_t>(sizeof(Counter))})
.with_dtype<uint8_t>()
.with_device(device)
.verify(push_counter);
RuntimeCheck(push_bytes.unwrap() > 0 && pull_bytes.unwrap() > 0, "Workspace sizes must be positive");
if (pull_mc_workspace_ptr.has_value()) {
this->pull_mc_workspace = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(pull_mc_workspace_ptr.value()));
} else {
this->pull_mc_workspace = nullptr;
}
// push config
this->push_counter = static_cast<Counter*>(push_counter.data_ptr());
this->push_bytes = push_bytes.unwrap();
this->num_push_blocks = static_cast<uint32_t>(num_push_blocks.unwrap());
// pull config
this->pull_bytes = pull_bytes.unwrap();
this->num_pull_blocks = static_cast<uint32_t>(num_pull_blocks.unwrap());
this->num_multicast_blocks = this->num_pull_blocks;
this->total_pull_blocks = this->num_pull_blocks;
}
inline void CommunicatorObj::config(std::map<std::string, uint32_t> config) {
for (const auto& [key, value] : config) {
if (key == "num_pull_blocks") {
RuntimeCheck(value > 0 && value <= total_pull_blocks, "Invalid number of pull blocks: ", value);
this->num_pull_blocks = value;
} else if (key == "num_multicast_blocks") {
RuntimeCheck(value > 0 && value <= total_pull_blocks, "Invalid number of multicast blocks: ", value);
this->num_multicast_blocks = value;
} else {
RuntimeCheck(false, "Unknown config key: ", key);
}
}
}
} // namespace host::distributed
inline void register_communicator() {
namespace refl = tvm::ffi::reflection;
using Class = host::distributed::CommunicatorObj;
using TensorView = tvm::ffi::TensorView;
refl::ObjectDef<Class>()
.def(
refl::init<
uint32_t,
uint32_t,
std::vector<TensorView>,
std::vector<TensorView>,
std::vector<TensorView>,
TensorView,
std::optional<int64_t>>(),
"__init__")
.def_ro("world_size", &Class::world_size)
.def_ro("rank", &Class::rank)
.def("_config", &Class::config);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh
+286 -522
View File
@@ -9,12 +9,20 @@
// shard in place so every workspace ends up holding the full result.
//
// Unlike the previous implementation, the kernels carry no storage or IPC
// logic: all pointers arrive via `CommunicatorObj` (owned by Python) and the
// per-call `AllReduceParams`.
// logic: all pointers arrive via the communication planes (owned by Python)
// and the per-call params. The push and pull families take disjoint params,
// so neither carries the other's pointer table into its grid constants.
//
// The pull family reduces over the pull plane's workspaces, which exist
// because this kernel's callers hand it plain tensors: the host stages the
// input in and copies the result back out. Callers that already allocate
// from symmetric memory (the K3 fused paths) reduce in place instead and
// borrow the plane only to barrier on.
#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>
@@ -22,475 +30,235 @@
#include <sgl_kernel/distributed/communicator.cuh>
#include <tvm/ffi/extra/stl.h>
#include <tvm/ffi/object.h>
#include <algorithm>
#include <bit>
#include <cstdint>
#include <cstring>
#include <string>
#include <variant>
namespace sglang {
using device::distributed::Counter, device::distributed::Semaphore;
using device::distributed::PullWorkSpace, device::distributed::PushWorkSpace;
using host::distributed::CommunicatorRef;
inline constexpr uint32_t kMaxWorldSize = device::distributed::kMaxWorldSize;
enum class PullMode {
Graph,
Eager,
Multicast, // also eager
};
template <typename T>
struct fp_trait {};
template <>
struct fp_trait<bf16_t> {
using type = uint16_t;
[[maybe_unused]]
static constexpr uint16_t pos_zero = 0x0000u;
[[maybe_unused]]
static constexpr uint16_t neg_zero = 0x8000u;
};
template <>
struct fp_trait<fp16_t> {
using type = uint16_t;
[[maybe_unused]]
static constexpr uint16_t pos_zero = 0x0000u;
[[maybe_unused]]
static constexpr uint16_t neg_zero = 0x8000u;
};
template <>
struct fp_trait<float> {
using type = uint32_t;
[[maybe_unused]]
static constexpr uint32_t pos_zero = 0x00000000u;
[[maybe_unused]]
static constexpr uint32_t neg_zero = 0x80000000u;
};
template <typename DType>
SGL_DEVICE void clear_pos_zero(DType& val) {
using Trait = fp_trait<DType>;
const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
}
template <typename DType>
SGL_DEVICE bool is_pos_zero(const DType& val) {
using Trait = fp_trait<DType>;
const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
return *ptr == Trait::pos_zero;
}
template <typename DType>
SGL_DEVICE DType get_pos_zero() {
using Trait = fp_trait<DType>;
const auto value = Trait::pos_zero;
return *reinterpret_cast<const DType*>(&value);
}
template <typename T2, size_t N, size_t M>
SGL_DEVICE auto reduce(device::AlignedVector<T2, N> (&vec)[M]) -> device::AlignedVector<T2, N> {
fp32x2_t acc_vec[N];
#pragma unroll
for (size_t i = 0; i < M; ++i) {
#pragma unroll
for (size_t j = 0; j < N; ++j) {
const auto [x, y] = device::cast<fp32x2_t>(vec[i][j]);
auto& [acc_x, acc_y] = acc_vec[j];
acc_x = i == 0 ? x : acc_x + x;
acc_y = i == 0 ? y : acc_y + y;
}
}
device::AlignedVector<T2, N> out_vec;
#pragma unroll
for (size_t j = 0; j < N; ++j) {
out_vec[j] = device::cast<T2>(acc_vec[j]);
}
return out_vec;
}
template <typename V>
SGL_DEVICE void ld_global_16B(V& x, const void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
addr = static_cast<const uint8_t*>(addr) + vec_offset * sizeof(V);
uint4 val;
asm volatile("ld.global.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(addr));
x = *reinterpret_cast<const V*>(&val);
}
template <typename V>
SGL_DEVICE void st_global_16B(const V& x, void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
const uint4 val = *reinterpret_cast<const uint4*>(&x);
addr = static_cast<uint8_t*>(addr) + vec_offset * sizeof(V);
asm volatile("st.global.v4.b32 [%4], {%0, %1, %2, %3};"
: //
: "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
}
template <typename V>
SGL_DEVICE void ld_relaxed_16B(V& x, const void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
addr = static_cast<const uint8_t*>(addr) + vec_offset * sizeof(V);
uint4 val;
asm volatile("ld.relaxed.sys.global.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(addr));
x = *reinterpret_cast<const V*>(&val);
}
template <typename V>
SGL_DEVICE void st_relaxed_16B(const V& x, void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
const uint4 val = *reinterpret_cast<const uint4*>(&x);
addr = static_cast<uint8_t*>(addr) + vec_offset * sizeof(V);
asm volatile("st.relaxed.sys.global.v4.b32 [%4], {%0, %1, %2, %3};"
: //
: "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
}
template <typename V>
SGL_DEVICE void ld_multimem_16B(V& x, const void* mc_addr, int64_t vec_offset) {
#if SGL_ARCH_HOPPER_OR_GREATER
static_assert(alignof(V) == 16 && sizeof(V) == 16);
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];"
: "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w)
: "l"(mc_addr));
x = *reinterpret_cast<const V*>(&val);
} else {
// Packed f16x2/bf16x2 results live in b32 registers ("=r"); .acc::f32 only
// raises the accumulation precision, not the result register type — ptxas
// 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];"
: "=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];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(mc_addr));
}
x = *reinterpret_cast<const V*>(&val);
}
#else
assert(false && "multimem load is only supported on Hopper or later architecture");
#endif
}
template <typename V>
SGL_DEVICE void st_multimem_16B(const V& x, void* mc_addr, int64_t vec_offset) {
#if SGL_ARCH_HOPPER_OR_GREATER
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};"
:
: "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w), "l"(mc_addr));
#else
assert(false && "multimem store is only supported on Hopper or later architecture");
#endif
}
template <uint32_t kWorldSize>
struct AllReduceParams {
struct AllReducePushParams {
const void* __restrict__ input;
void* __restrict__ output;
uint32_t num_elements;
uint32_t num_vecs;
uint32_t rank;
PushWorkSpace<kWorldSize> ws;
};
template <uint32_t kWorldSize>
struct AllReducePullParams {
void* __restrict__ output;
uint32_t num_vecs;
uint32_t rank;
void* const* __restrict__ graph_params;
uint8_t* pull_workspaces[kWorldSize]; // must be symmetric memory
uint8_t* push_workspaces[kWorldSize]; // must be symmetric memory
Semaphore* pull_semaphores[kWorldSize]; // must be symmetric memory
Counter* push_counter;
uint8_t* pull_mc_workspace; // must be a multicast address
int64_t push_buffer_stride; // per-buffer bytes; each rank holds 2 * world_size buffers
PullWorkSpace<kWorldSize> ws;
};
template <typename T, uint32_t kWorldSize, bool kUsePDL>
struct AllReducePushImpl {
private:
using T2 = packed_t<T>;
/// NOTE: force 16B load/store to reduce register pressure
static constexpr uint32_t kVecSize = 16 / sizeof(T2);
static constexpr uint32_t kElemsPerVec = 16 / sizeof(T);
using vec_t = device::AlignedVector<T2, kVecSize>;
static_assert(kWorldSize <= kMaxWorldSize);
static SGL_DEVICE bool sync_enter_push(const AllReduceParams<kWorldSize>& params) {
device::PDLWaitPrimary<kUsePDL>();
return (params.push_counter[blockIdx.x].get() % 2) != 0;
}
static SGL_DEVICE void sync_exit_push(const AllReduceParams<kWorldSize>& params) {
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
if (threadIdx.x == 0) {
params.push_counter[blockIdx.x].inc(1); // NOTE: u32 overflow is safe under mod 2
}
}
static SGL_DEVICE void push_impl(uint32_t num_vecs, void* (&data)[kWorldSize], const void* src) {
const auto num_threads = blockDim.x * gridDim.x;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
#pragma unroll
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
ld_global_16B(vec, src, vid);
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
clear_pos_zero(vec[j].x);
clear_pos_zero(vec[j].y);
}
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
st_relaxed_16B(vec, data[i], vid);
}
}
}
static SGL_DEVICE void poll_impl(uint32_t num_vecs, void* (&data)[kWorldSize], void* out) {
// need polling to ensure data is ready
const auto num_threads = blockDim.x * gridDim.x;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
// pos_zero-filled vec we write back after consuming each slot, so the
// double-buffered phase comes back around with the "slot empty" marker
// re-established.
vec_t pos_zero_vec;
{
const auto z = get_pos_zero<T>();
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
pos_zero_vec[j].x = z;
pos_zero_vec[j].y = z;
}
}
#pragma unroll
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec[kWorldSize];
do {
bool has_zero = false;
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
ld_relaxed_16B(vec[i], data[i], vid);
}
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
#pragma unroll
for (uint32_t j = 0; j < kVecSize; ++j) {
has_zero |= is_pos_zero(vec[i][j].x);
has_zero |= is_pos_zero(vec[i][j].y);
}
}
if (!has_zero) break;
} while (true);
const auto out_vec = reduce(vec);
st_global_16B(out_vec, out, vid);
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
st_global_16B(pos_zero_vec, data[i], vid);
}
}
}
/// `vec_offset` is a *vector index* and must stay folded into the base pointers
/// here: a typed 32-bit bias becomes one widening multiply-add off a
/// constant-bank base, which ptxas keeps on the uniform datapath, so the whole
/// peer table lives in uniform registers. Biasing `vid` at each access instead,
/// or using a 64-bit byte bias (the uniform datapath has no 64-bit add), makes
/// all `kWorldSize` addresses thread-varying and costs 2shot ~12 registers per
/// thread and ~10% throughput.
template <typename V, uint32_t kWorldSize, bool kUseGraph>
struct LoadStoreImpl {
public:
static SGL_DEVICE void forward_1shot(const AllReduceParams<kWorldSize>& params) {
// push local data to peer ranks, then reduce locally
const auto phase = sync_enter_push(params);
const auto r = params.rank;
const auto num_vecs = device::div_ceil(params.num_elements, kElemsPerVec);
const auto stride_bytes = params.push_buffer_stride;
const auto phase_stride_bytes = phase * stride_bytes * kWorldSize;
static constexpr uint32_t size() {
return kWorldSize;
}
SGL_DEVICE LoadStoreImpl(const AllReducePullParams<kWorldSize>& params, uint32_t vec_offset = 0) {
if constexpr (kUseGraph) {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
m_data[i] = reinterpret_cast<V*>(params.graph_params[i]) + vec_offset;
}
} else {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
m_data[i] = reinterpret_cast<V*>(params.ws.workspaces[i]) + vec_offset;
}
}
}
// push to peer
void* push_buf[kWorldSize];
SGL_DEVICE void load_reduce(V& vec, uint32_t vid) const {
V vecs[kWorldSize];
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
push_buf[i] = params.push_workspaces[i] + r * stride_bytes + phase_stride_bytes;
vecs[i].load(m_data[i], vid);
}
push_impl(num_vecs, push_buf, params.input);
vec = device::reduce_vec(vecs);
}
// poll from local
void* poll_buf[kWorldSize];
SGL_DEVICE void store_multi(const V& val, uint32_t vid) const {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
poll_buf[i] = params.push_workspaces[r] + i * stride_bytes + phase_stride_bytes;
val.store(m_data[i], vid);
}
poll_impl(num_vecs, poll_buf, params.output);
sync_exit_push(params);
}
};
template <typename T, uint32_t kWorldSize, PullMode kMode, bool kUsePDL>
struct AllReducePullImpl {
private:
using T2 = packed_t<T>;
static constexpr uint32_t kVecSize = 16 / sizeof(T2);
static constexpr uint32_t kElemsPerVec = 16 / sizeof(T);
using vec_t = device::AlignedVector<T2, kVecSize>;
static_assert(kWorldSize <= kMaxWorldSize);
template <bool kFence>
static SGL_DEVICE uint32_t sync_enter_pull(const AllReduceParams<kWorldSize>& params) {
uint32_t current_counter_val = 0;
if (const auto tx = threadIdx.x; tx < kWorldSize) {
device::PDLWaitPrimary<kUsePDL>();
const auto bx = blockIdx.x;
const auto semaphore = &params.pull_semaphores[tx][bx];
const auto counter = semaphore->counter_ptr();
const auto current = tx == params.rank ? counter->inc(2 * kWorldSize) : 0;
current_counter_val = current;
if constexpr (kFence) {
semaphore->put_release();
} else {
semaphore->put_relaxed();
}
if (tx == params.rank) {
if constexpr (kFence) {
while (semaphore->get_acquire() - current < kWorldSize)
;
} else {
while (semaphore->get_relaxed() - current < kWorldSize)
;
}
}
}
__syncthreads();
return current_counter_val + kWorldSize;
}
template <bool kFence>
static SGL_DEVICE void sync_exit_pull(const AllReduceParams<kWorldSize>& params, uint32_t current) {
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
if (const auto tx = threadIdx.x; tx < kWorldSize) {
const auto bx = blockIdx.x;
const auto semaphore = &params.pull_semaphores[tx][bx];
if constexpr (kFence) {
semaphore->put_release();
} else {
semaphore->put_relaxed();
}
if (tx == params.rank) {
if constexpr (kFence) {
while (semaphore->get_acquire() - current < kWorldSize)
;
} else {
while (semaphore->get_relaxed() - current < kWorldSize)
;
}
}
}
}
template <bool kIs2shot>
static SGL_DEVICE void reduce_impl(
uint32_t num_vecs, //
[[maybe_unused]] void* (&data)[kWorldSize],
[[maybe_unused]] void* out,
[[maybe_unused]] void* mc_addr) {
const auto num_threads = blockDim.x * gridDim.x;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
if constexpr (kMode == PullMode::Multicast) {
vec_t out_vec;
ld_multimem_16B(out_vec, mc_addr, vid);
if constexpr (kIs2shot) {
// inplace write to workspace for 2-shot all reduce
st_multimem_16B(out_vec, mc_addr, vid);
} else {
// write to output for 1-shot all reduce
out_vec.store(out, vid);
}
} else {
vec_t vec[kWorldSize];
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
vec[i].load(data[i], vid);
}
const auto out_vec = reduce(vec);
if constexpr (kIs2shot) {
// inplace write to buffer for 2-shot all reduce
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
out_vec.store(data[i], vid);
}
} else {
// write to output for 1-shot all reduce
out_vec.store(out, vid);
}
}
}
}
public:
static SGL_DEVICE void forward_1shot(const AllReduceParams<kWorldSize>& params) {
const auto total_num_vecs = device::div_ceil(params.num_elements, kElemsPerVec);
void* data[kWorldSize];
if constexpr (kMode == PullMode::Graph) {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
data[i] = params.graph_params[i];
}
} else if constexpr (kMode == PullMode::Eager) {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
data[i] = params.pull_workspaces[i];
}
}
const auto counter = sync_enter_pull<false>(params);
reduce_impl<false>(total_num_vecs, data, params.output, params.pull_mc_workspace);
sync_exit_pull<false>(params, counter);
}
static SGL_DEVICE void forward_2shot(const AllReduceParams<kWorldSize>& params) {
const auto total_num_vecs = device::div_ceil(params.num_elements, kElemsPerVec);
const auto avg_vecs = total_num_vecs / kWorldSize;
const auto rem_vecs = total_num_vecs % kWorldSize;
// usually, hidden size is a multiple of 1024, so 1024 / 8 = 128 is typically 128-bytes aligned
const auto local_vec_bias = avg_vecs * params.rank + min(params.rank, rem_vecs);
const auto local_num_vecs = avg_vecs + (params.rank < rem_vecs ? 1 : 0);
void* data[kWorldSize];
if constexpr (kMode == PullMode::Graph) {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
data[i] = reinterpret_cast<vec_t*>(params.graph_params[i]) + local_vec_bias;
}
} else if constexpr (kMode == PullMode::Eager) {
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
data[i] = reinterpret_cast<vec_t*>(params.pull_workspaces[i]) + local_vec_bias;
}
}
const auto counter = sync_enter_pull<false>(params);
const auto mc_addr = reinterpret_cast<vec_t*>(params.pull_mc_workspace) + local_vec_bias;
reduce_impl<true>(local_num_vecs, data, params.output, mc_addr);
sync_exit_pull<true>(params, counter);
}
V* m_data[kWorldSize];
};
template <typename Impl, uint32_t kWorldSize, int kShot>
__global__ __launch_bounds__(1024, 1) //
void all_reduce_kernel(const __grid_constant__ AllReduceParams<kWorldSize> params) {
static_assert(kShot == 1 || kShot == 2, "invalid shot");
if constexpr (kShot == 1) {
return Impl::forward_1shot(params);
} else {
return Impl::forward_2shot(params);
template <typename V, uint32_t kWorldSize, bool kUseGraph>
struct MultiCastImpl {
public:
static_assert(kUseGraph == false);
static constexpr uint32_t size() {
return kWorldSize;
}
SGL_DEVICE MultiCastImpl(const AllReducePullParams<kWorldSize>& params, uint32_t vec_offset = 0)
: m_multicast_ptr(reinterpret_cast<V*>(params.ws.mc_workspace) + vec_offset) {}
SGL_DEVICE void load_reduce(V& vec, uint32_t vid) const {
device::ptx::ld_multimem_16B(vec, m_multicast_ptr, vid);
}
SGL_DEVICE void store_multi(const V& val, uint32_t vid) const {
return device::ptx::st_multimem_16B(val, m_multicast_ptr, vid);
}
private:
V* m_multicast_ptr;
};
#define ALL_REDUCE_KERNEL __global__ __launch_bounds__(1024, 1)
template <typename Impl, typename T, uint32_t kWorldSize, bool kUsePDL>
ALL_REDUCE_KERNEL void all_reduce_1shot_push_kernel(const __grid_constant__ AllReducePushParams<kWorldSize> params) {
using namespace device;
constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2);
using vec_t = AlignedVector<packed_t<T>, kVecSize>;
using Lamport = distributed::LamportTrait<T, kVecSize * 2, /*kAtom=*/4>;
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;
PDLWaitPrimary<kUsePDL>();
const auto epoch = distributed::PushEpoch<kWorldSize>{params.ws};
// push to peer
void* push_ptrs[kWorldSize];
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
push_ptrs[i] = epoch.slot_ptr(/*dst=*/i, /*src=*/r);
}
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
vec.load(params.input, vid);
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);
}
}
// poll from local
void* poll_ptrs[kWorldSize];
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
poll_ptrs[i] = epoch.slot_ptr(/*dst=*/r, /*src=*/i);
}
vec_t pos_zero_vec;
Lamport::fill_pos_zero(pos_zero_vec.data());
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t 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) {
has_zero |= Lamport::has_pos_zero(vec[i].data());
}
if (!has_zero) break;
} while (true);
const auto out_vec = reduce_vec(vec);
out_vec.store(params.output, vid);
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
ptx::st_global_16B(pos_zero_vec, poll_ptrs[i], vid);
}
}
PDLTriggerSecondary<kUsePDL>();
__syncthreads();
epoch.flip();
}
template <typename Impl, typename T, uint32_t kWorldSize, bool kUsePDL>
ALL_REDUCE_KERNEL void all_reduce_1shot_pull_kernel(const __grid_constant__ AllReducePullParams<kWorldSize> params) {
using namespace device;
constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2);
using vec_t = AlignedVector<packed_t<T>, kVecSize>;
const auto num_vecs = params.num_vecs;
const auto impl = Impl{params};
PDLWaitPrimary<kUsePDL>();
const auto barrier = distributed::Barrier<kWorldSize>{
params.ws.semaphores.data(),
params.rank,
/*num_arrives=*/2,
};
barrier.arrive_relaxed(/*n=*/0);
__syncthreads();
const auto num_threads = blockDim.x * gridDim.x;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
impl.load_reduce(vec, vid);
vec.store(params.output, vid);
}
PDLTriggerSecondary<kUsePDL>();
__syncthreads();
barrier.arrive_relaxed(/*n=*/1);
}
template <typename Impl, typename T, uint32_t kWorldSize, bool kUsePDL>
ALL_REDUCE_KERNEL void all_reduce_2shot_pull_kernel(const __grid_constant__ AllReducePullParams<kWorldSize> params) {
using namespace device;
constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2);
using vec_t = AlignedVector<packed_t<T>, kVecSize>;
const auto num_total_vecs = params.num_vecs;
const auto avg_vecs = num_total_vecs / kWorldSize;
const auto rem_vecs = num_total_vecs % kWorldSize;
const auto num_vecs = avg_vecs + (params.rank < rem_vecs ? 1 : 0);
const auto vec_offset = params.rank * avg_vecs + min(params.rank, rem_vecs);
const auto impl = Impl{params, vec_offset};
PDLWaitPrimary<kUsePDL>();
const auto barrier = distributed::Barrier<kWorldSize>{
params.ws.semaphores.data(),
params.rank,
/*num_arrives=*/2,
};
barrier.arrive_relaxed(/*n=*/0);
__syncthreads();
const auto num_threads = blockDim.x * gridDim.x;
const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
impl.load_reduce(vec, vid);
impl.store_multi(vec, vid);
}
PDLTriggerSecondary<kUsePDL>();
__syncthreads();
barrier.arrive_rel_acq(/*n=*/1);
}
template <uint32_t N>
@@ -507,15 +275,11 @@ __global__ void memcpy_kernel(void* __restrict__ dst, const void* __restrict__ s
}
}
// Pick the smallest block size whose grid still fits in one wave; the kernels
// are grid-stride so any choice is correct, this only tunes occupancy.
[[maybe_unused]]
uint32_t choose_block_size(uint32_t num_threads) {
inline auto choose_block_size(uint32_t num_threads) -> uint32_t {
static const uint32_t kNumSM = [] {
int device = 0, sm_count = 0;
host::RuntimeDeviceCheck(cudaGetDevice(&device));
host::RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device));
return static_cast<uint32_t>(sm_count);
int device = 0;
CHECK_CUDA(cudaGetDevice(&device));
return host::runtime::get_sm_count(device);
}();
for (const uint32_t block_size : {128u, 256u, 512u}) {
if (host::div_ceil(num_threads, block_size) <= kNumSM) return block_size;
@@ -528,68 +292,67 @@ struct AllReduceKernel {
private:
using Tensor = tvm::ffi::Tensor;
using TensorView = tvm::ffi::TensorView;
template <int kShot, PullMode kPullMode>
static constexpr auto kernel_pull =
all_reduce_kernel<AllReducePullImpl<T, kWorldSize, kPullMode, kUsePDL>, kWorldSize, kShot>;
template <int kShot>
static constexpr auto kernel_push = all_reduce_kernel<AllReducePushImpl<T, kWorldSize, kUsePDL>, kWorldSize, kShot>;
using PushParams = AllReducePushParams<kWorldSize>;
using PullParams = AllReducePullParams<kWorldSize>;
using vec_t = device::AlignedVector<packed_t<T>, 16 / (sizeof(T) * 2)>;
public:
static Tensor run(CommunicatorRef ref, Tensor in_, std::string algo, std::variant<TensorView, bool> pull_arg) {
static Tensor
run(const CommunicatorRef comm_ref,
const Tensor in,
const std::string algo,
const tvm::ffi::Optional<TensorView> graph_params_opt,
const bool use_multicast) {
using namespace host;
const auto& data = *ref.get();
RuntimeCheck(algo == "1shot_pull" || algo == "2shot_pull" || algo == "1shot_push", "Invalid algo: ", algo);
RuntimeCheck(data.world_size == kWorldSize, "Mismatch world size");
RuntimeCheck(in_.IsContiguous(), "Input tensor must be contiguous");
RuntimeCheck(is_type<T>(in_.dtype()), "Input dtype mismatch");
RuntimeCheck(in_.device().device_type == kDLCUDA, "Only CUDA device is supported");
RuntimeCheck(std::bit_cast<intptr_t>(in_.data_ptr()) % 16 == 0, "Input pointer is not properly aligned");
const auto num_elems_int64 = in_.numel();
const auto& comm = *comm_ref.get();
CHECK_HOST(algo == "1shot_pull" || algo == "2shot_pull" || algo == "1shot_push") << algo;
CHECK_HOST(comm.get_world_size() == kWorldSize) << comm.get_world_size();
CHECK_HOST(in.IsContiguous() && is_type<T>(in.dtype()) && in.device().device_type == kDLCUDA);
const auto num_elems_int64 = in.numel();
const auto num_elems = static_cast<uint32_t>(num_elems_int64);
RuntimeCheck(static_cast<int64_t>(num_elems) == num_elems_int64, "Number of items exceeds 4G limit");
const bool use_graph = std::holds_alternative<TensorView>(pull_arg);
const auto graph_ptr = use_graph ? std::get<TensorView>(pull_arg).data_ptr() : nullptr;
const bool inplace = use_graph && algo == "2shot_pull";
Tensor out = inplace ? in_ : ffi::empty_like(in_);
AllReduceParams<kWorldSize> params{
.input = in_.data_ptr(),
.output = out.data_ptr(),
.num_elements = num_elems,
.rank = data.rank,
.graph_params = static_cast<void* const*>(graph_ptr),
.pull_workspaces = {},
.push_workspaces = {},
.pull_semaphores = {},
.push_counter = data.push_counter,
.pull_mc_workspace = data.pull_mc_workspace,
.push_buffer_stride = data.push_bytes,
};
for (uint32_t i = 0; i < kWorldSize; ++i) {
params.pull_workspaces[i] = data.pull_workspaces[i];
params.push_workspaces[i] = data.push_workspaces[i];
params.pull_semaphores[i] = data.pull_semaphores[i];
}
const int64_t nbytes = num_elems_int64 * sizeof(T);
RuntimeCheck(nbytes % 16 == 0, "Input bytes must be a multiple of 16, got: ", nbytes);
const uint32_t num_vecs = num_elems / (16 / sizeof(T));
const auto stream = LaunchKernel::resolve_device(in_.device());
const auto nbytes = static_cast<int64_t>(num_elems_int64 * sizeof(T));
CHECK_HOST(static_cast<int64_t>(num_elems) == num_elems_int64) << num_elems_int64;
CHECK_HOST(reinterpret_cast<intptr_t>(in.data_ptr()) % 16 == 0 && nbytes % 16 == 0);
const auto num_vecs = static_cast<uint32_t>(num_elems / (16 / sizeof(T)));
const auto stream = LaunchKernel::resolve_device(in.device());
const auto use_graph = graph_params_opt.has_value();
if (algo == "1shot_push") {
RuntimeCheck(!use_graph, "Push mode doesn't have graph mode optimization");
RuntimeCheck(nbytes <= data.push_bytes, "Input size ", nbytes, " exceeds push workspace size ", data.push_bytes);
CHECK_HOST(!use_graph && !use_multicast);
const auto& push = comm.get_push_obj();
Tensor out = ffi::empty_like(in);
const auto params = PushParams{
.input = in.data_ptr(),
.output = out.data_ptr(),
.num_vecs = num_vecs,
.rank = push.rank,
.ws = push.get_workspace<kWorldSize>(nbytes),
};
using Impl = LoadStoreImpl<vec_t, kWorldSize, /*kUseGraph=*/false>;
const auto kernel = all_reduce_1shot_push_kernel<Impl, T, kWorldSize, kUsePDL>;
// the grid is bound to the counter array and must stay constant
const uint32_t num_blocks = data.num_push_blocks;
LaunchKernel(num_blocks, choose_block_size(num_vecs), stream) //
.enable_pdl(kUsePDL)(kernel_push<1>, params);
LaunchKernel(push.num_blocks, choose_block_size(num_vecs), stream) //
.enable_pdl(kUsePDL)(kernel, params);
return out;
}
using enum PullMode;
RuntimeCheck(nbytes <= data.pull_bytes, "Input size ", nbytes, " exceeds pull workspace size ", data.pull_bytes);
const auto pull_mode = use_graph ? Graph : std::get<bool>(pull_arg) ? Multicast : Eager;
RuntimeCheck(pull_mode != Multicast || data.pull_mc_workspace != nullptr, "Multicast requires an mc workspace");
const auto& pull = comm.get_pull_obj();
const auto graph_params = use_graph ? graph_params_opt.value().data_ptr() : nullptr;
// only 2shot pull + graph mode forces inplace implementation
const auto is_inplace = use_graph && algo == "2shot_pull";
Tensor out = is_inplace ? in : ffi::empty_like(in);
const auto params = PullParams{
.output = out.data_ptr(),
.num_vecs = num_vecs,
.rank = pull.rank,
.graph_params = static_cast<void* const*>(graph_params),
// Graph mode reduces over the caller's own registered buffers and only
// barriers on this plane; the eager modes stage the input through it.
.ws = pull.get_workspace<kWorldSize>(use_graph ? 0 : nbytes),
};
const auto& ws = params.ws;
CHECK_HOST(!use_multicast || ws.mc_workspace != nullptr) << "multicast needs a plane with an mc workspace";
const uint32_t num_blocks = data.num_pull_blocks;
const auto cuda_memcpy = [&](void* dst, const void* src) {
if constexpr (SGL_ARCH_HOPPER_OR_GREATER) { // PDL memcpy is faster
// based on micro benchmark, only enable when batch size is small + aligned
@@ -605,48 +368,49 @@ struct AllReduceKernel {
}
}
// safe fallback to cudaMemcpyAsync for large size or older architecture
RuntimeDeviceCheck(cudaMemcpyAsync(dst, src, nbytes, cudaMemcpyDeviceToDevice, stream));
CHECK_CUDA(cudaMemcpyAsync(dst, src, nbytes, cudaMemcpyDeviceToDevice, stream));
};
const uint32_t num_blocks = comm.get_pull_blocks();
const auto local_workspace = ws.workspaces[pull.rank];
const auto local_workspace = data.pull_workspaces[data.rank];
using LS = LoadStoreImpl<vec_t, kWorldSize, /*kUseGraph=*/false>;
using LS_GRAPH = LoadStoreImpl<vec_t, kWorldSize, /*kUseGraph=*/true>;
using MC = MultiCastImpl<vec_t, kWorldSize, /*kUseGraph=*/false>;
if (algo == "1shot_pull") {
// first copy to workspace
if (!use_graph) cuda_memcpy(local_workspace, in_.data_ptr());
const auto kernel = (pull_mode == Graph) ? kernel_pull<1, Graph>
: pull_mode == Eager ? kernel_pull<1, Eager>
: kernel_pull<1, Multicast>;
// first copy to the workspace
CHECK_HOST(!use_multicast);
if (!use_graph) cuda_memcpy(local_workspace, in.data_ptr());
const auto kernel = use_graph ? all_reduce_1shot_pull_kernel<LS_GRAPH, T, kWorldSize, kUsePDL>
: all_reduce_1shot_pull_kernel<LS, T, kWorldSize, kUsePDL>;
// then launch kernel to reduce and write to output
LaunchKernel(num_blocks, choose_block_size(num_vecs), stream) //
.enable_pdl(kUsePDL)(kernel, params);
} else /* 2shot_pull */ {
} else /* algo == "2shot_pull" */ {
const uint32_t avg_vecs = div_ceil(num_vecs, kWorldSize);
// first copy to workspace
if (!use_graph) cuda_memcpy(local_workspace, in_.data_ptr());
// then launch kernel to reduce in workspace
const auto kernel = (pull_mode == Graph) ? kernel_pull<2, Graph>
: pull_mode == Eager ? kernel_pull<2, Eager>
: kernel_pull<2, Multicast>;
if (pull_mode == Multicast) {
const auto max_blocks = data.num_multicast_blocks;
// first copy to the workspace
if (!use_graph) cuda_memcpy(local_workspace, in.data_ptr());
// then launch kernel to reduce in the workspace
if (use_multicast) {
CHECK_HOST(!use_graph);
const auto kernel = all_reduce_2shot_pull_kernel<MC, T, kWorldSize, kUsePDL>;
// NOTE: too much traffic will degrade performance in multicast impl
constexpr uint32_t kMulticastNumThreads = 512u;
// NOTE: too much traffic will degrade performance in multicast
LaunchKernel(std::min(num_blocks, max_blocks), kMulticastNumThreads, stream)
const auto num_blocks = comm.get_pull_multicast_blocks();
LaunchKernel(num_blocks, kMulticastNumThreads, stream) //
.enable_pdl(kUsePDL)(kernel, params);
} else {
const auto kernel = use_graph ? all_reduce_2shot_pull_kernel<LS_GRAPH, T, kWorldSize, kUsePDL>
: all_reduce_2shot_pull_kernel<LS, T, kWorldSize, kUsePDL>;
LaunchKernel(num_blocks, choose_block_size(avg_vecs), stream) //
.enable_pdl(kUsePDL)(kernel, params);
}
// finally copy from workspace to output
// finally copy from the workspace to output
if (!use_graph) cuda_memcpy(out.data_ptr(), local_workspace);
}
return out;
}
};
template <typename T, uint32_t kWorldSize, bool kUsePDL>
tvm::ffi::Tensor custom_all_reduce(
CommunicatorRef comm, tvm::ffi::Tensor input, std::string algo, std::variant<tvm::ffi::TensorView, bool> pull_arg) {
return AllReduceKernel<T, kWorldSize, kUsePDL>::run(comm, input, algo, pull_arg);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/registry.cuh
+146
View File
@@ -0,0 +1,146 @@
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/distributed/communicator.cuh>
#include <tvm/ffi/extra/stl.h>
#include <tvm/ffi/reflection/registry.h>
namespace sglang {
namespace host::distributed {
BasePlane::BasePlane(uint32_t rank, uint32_t world_size) : rank(rank), world_size(world_size) {
CHECK_HOST(1 < world_size && world_size <= kMaxWorldSize) << "Invalid world size: " << world_size;
CHECK_HOST(rank < world_size) << "Invalid rank " << rank << " for world size " << world_size;
}
PushPlaneObj::PushPlaneObj(
uint32_t rank,
uint32_t world_size,
std::vector<TensorView> workspaces, // world_size * [2 * world_size][slot_bytes]
TensorView counter, // [num_blocks]
intptr_t mc_workspace_ptr)
: BasePlane(rank, world_size), //
num_blocks(),
slot_bytes(),
counter(),
workspaces{},
mc_workspace() {
CHECK_HOST(workspaces.size() == world_size) << "Bad push workspace count";
// Shared symbolic sizes and device enforce consistency across ranks; the
// matchers also require contiguity (no strides given) and uint8 dtype.
auto N = SymbolicSize{"slot_bytes"};
auto M = SymbolicSize{"num_blocks"};
auto device_sym = SymbolicDevice{};
device_sym.set_options<kDLCUDA>();
for (uint32_t i = 0; i < world_size; ++i) {
TensorMatcher({2 * world_size, N}) //
.with_dtype<uint8_t>()
.with_device(device_sym)
.verify(workspaces[i]);
}
TensorMatcher({M, static_cast<int64_t>(sizeof(Counter))}) //
.with_dtype<uint8_t>()
.with_device(device_sym)
.verify(counter);
CHECK_HOST(N.unwrap() > 0 && M.unwrap() > 0) << "A push plane needs a non-empty workspace and counter";
// only set the value safely after the symbolic size has been verified
this->num_blocks = static_cast<uint32_t>(M.unwrap());
this->slot_bytes = N.unwrap();
this->counter = static_cast<Counter*>(counter.data_ptr());
for (uint32_t i = 0; i < world_size; ++i) {
this->workspaces[i] = static_cast<uint8_t*>(workspaces[i].data_ptr());
}
this->mc_workspace = reinterpret_cast<uint8_t*>(mc_workspace_ptr);
}
PullPlaneObj::PullPlaneObj(
uint32_t rank,
uint32_t world_size,
std::vector<TensorView> workspaces, // world_size * [num_bytes]
std::vector<TensorView> semaphores, // world_size * [num_blocks]
intptr_t mc_workspace_ptr,
intptr_t mc_semaphore_ptr)
: BasePlane(rank, world_size), //
num_blocks(),
num_bytes(),
semaphores{},
workspaces{},
mc_semaphore(),
mc_workspace() {
CHECK_HOST(workspaces.size() == world_size) << "Bad pull workspace count";
CHECK_HOST(semaphores.size() == world_size) << "Bad pull semaphore count";
// Either half may be empty (a 0-element tensor); the halves a caller does
// own still have to agree on size and device across ranks.
auto N = SymbolicSize{"num_bytes"};
auto M = SymbolicSize{"num_blocks"};
auto device_sym = SymbolicDevice{};
device_sym.set_options<kDLCUDA>();
for (uint32_t i = 0; i < world_size; ++i) {
TensorMatcher({N}) //
.with_dtype<uint8_t>()
.with_device(device_sym)
.verify(workspaces[i]);
TensorMatcher({M, static_cast<int64_t>(sizeof(Semaphore))}) //
.with_dtype<uint8_t>()
.with_device(device_sym)
.verify(semaphores[i]);
}
CHECK_HOST(N.unwrap() > 0 || M.unwrap() > 0) << "A pull plane with neither workspaces nor semaphores is useless";
// only set the value safely after the symbolic size has been verified
this->num_blocks = static_cast<uint32_t>(M.unwrap());
this->num_bytes = N.unwrap();
for (uint32_t i = 0; i < world_size; ++i) {
this->semaphores[i] = static_cast<Semaphore*>(semaphores[i].data_ptr());
this->workspaces[i] = static_cast<uint8_t*>(workspaces[i].data_ptr());
}
this->mc_semaphore = reinterpret_cast<Semaphore*>(mc_semaphore_ptr);
this->mc_workspace = reinterpret_cast<uint8_t*>(mc_workspace_ptr);
}
CommunicatorObj::CommunicatorObj(Optional<PushPlaneRef> push, Optional<PullPlaneRef> pull)
: m_push(std::move(push)), m_pull(std::move(pull)) {
CHECK_HOST(m_push.has_value() || m_pull.has_value()) << "A communicator needs at least one plane";
if (m_push.has_value() && m_pull.has_value()) {
const auto& push_obj = *m_push.value().get();
const auto& pull_obj = *m_pull.value().get();
CHECK_HOST(push_obj.rank == pull_obj.rank && push_obj.world_size == pull_obj.world_size)
<< "Push and pull planes disagree on (rank, world_size)";
}
}
} // namespace host::distributed
inline void register_communicator() {
namespace refl = tvm::ffi::reflection;
namespace dist = host::distributed;
using TensorView = tvm::ffi::TensorView;
using Tensors = std::vector<TensorView>;
refl::ObjectDef<dist::PushPlaneObj>()
.def(refl::init<uint32_t, uint32_t, Tensors, TensorView, intptr_t>(), "__init__")
.def_ro("rank", &dist::PushPlaneObj::rank)
.def_ro("world_size", &dist::PushPlaneObj::world_size)
.def_ro("num_blocks", &dist::PushPlaneObj::num_blocks)
.def_ro("slot_bytes", &dist::PushPlaneObj::slot_bytes);
refl::ObjectDef<dist::PullPlaneObj>()
.def(refl::init<uint32_t, uint32_t, Tensors, Tensors, intptr_t, intptr_t>(), "__init__")
.def_ro("rank", &dist::PullPlaneObj::rank)
.def_ro("world_size", &dist::PullPlaneObj::world_size)
.def_ro("num_blocks", &dist::PullPlaneObj::num_blocks)
.def_ro("num_bytes", &dist::PullPlaneObj::num_bytes);
refl::ObjectDef<dist::CommunicatorObj>()
.def(refl::init<dist::Optional<dist::PushPlaneRef>, dist::Optional<dist::PullPlaneRef>>(), "__init__")
.def("get_rank", &dist::CommunicatorObj::get_rank)
.def("get_world_size", &dist::CommunicatorObj::get_world_size)
.def("get_push", &dist::CommunicatorObj::get_push)
.def("get_pull", &dist::CommunicatorObj::get_pull)
.def("set_pull_blocks", &dist::CommunicatorObj::set_pull_blocks)
.def("set_pull_multicast_blocks", &dist::CommunicatorObj::set_pull_multicast_blocks);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh
+10 -26
View File
@@ -36,23 +36,6 @@ struct ParallelQKNormParams {
uint32_t num_clean_up_count = 0;
};
template <typename T>
SGL_DEVICE void ld_global_volatile_8B(T& x, const void* addr, int64_t offset) {
static_assert(alignof(T) == 8 && sizeof(T) == 8);
addr = device::pointer::offset<T>(addr, offset);
uint2 val;
asm volatile("ld.volatile.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr));
x = *reinterpret_cast<const T*>(&val);
}
template <typename T>
SGL_DEVICE void st_global_volatile_8B(const T& x, void* addr, int64_t offset) {
static_assert(alignof(T) == 8 && sizeof(T) == 8);
const uint2 val = *reinterpret_cast<const uint2*>(&x);
addr = device::pointer::offset<T>(addr, offset);
asm volatile("st.volatile.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr));
}
[[maybe_unused]]
SGL_DEVICE float sync_float(float x) {
return __shfl_sync(0xffffffffu, x, 0);
@@ -193,10 +176,10 @@ __global__ __launch_bounds__(Trait::kBlockSize, Trait::kOccupancy) void parallel
sum_q_k[0] = sum_q + eps;
sum_q_k[1] = sum_k + eps;
const auto push_ptr = pointer::offset(buffer[tx], epoch_offset);
st_global_volatile_8B(sum_q_k, push_ptr, i * kNumGPU + rank);
ptx::st_relaxed_8B(sum_q_k, push_ptr, i * kNumGPU + rank);
const auto poll_ptr = pointer::offset(buffer[rank], epoch_offset);
while (true) {
ld_global_volatile_8B(sum_q_k, poll_ptr, i * kNumGPU + tx);
ptx::ld_relaxed_8B(sum_q_k, poll_ptr, i * kNumGPU + tx);
if (sum_q_k[0] != 0.0f && sum_q_k[1] != 0.0f) break;
}
constexpr uint32_t kActiveMask = (1 << kNumGPU) - 1;
@@ -246,6 +229,7 @@ struct FusedParallelQKNormAcrossHead {
const float eps // passed in unscaled
) {
using namespace host;
const auto& push = comm.get_push_obj();
constexpr auto Q = Trait::kLocalQDim;
constexpr auto K = Trait::kLocalKDim;
auto N = SymbolicSize{"num_tokens"};
@@ -274,7 +258,7 @@ struct FusedParallelQKNormAcrossHead {
// use at most `world_size` blocks to clean up,
// this is based on the observation that occupancy is usually linear
// with respect to the world size
const auto max_num_blocks = comm.num_push_blocks;
const auto max_num_blocks = push.num_blocks;
const bool need_clean = num_tokens < max_num_blocks;
const auto num_clean = need_clean ? (max_num_blocks - num_tokens) : 0;
const auto num_blocks = need_clean ? num_tokens + div_ceil(num_clean, Trait::kBlockSize) //
@@ -283,7 +267,7 @@ struct FusedParallelQKNormAcrossHead {
RuntimeCheck(num_blocks <= max_num_blocks, "internal error");
ParallelQKNormParams params;
for (uint32_t i = 0; i < kNumGPU; ++i) {
params.buffer[i] = comm.push_workspaces[i];
params.buffer[i] = push.workspaces[i];
}
params.q_ptr = q.data_ptr();
params.k_ptr = k.data_ptr();
@@ -292,21 +276,21 @@ struct FusedParallelQKNormAcrossHead {
params.q_stride_bytes = q.stride(0) * sizeof(DType);
params.k_stride_bytes = k.stride(0) * sizeof(DType);
params.eps = eps / kNumGPU; // scale down eps by number of GPUs
params.rank = comm.rank;
params.rank = push.rank;
params.num_tokens = num_tokens;
params.epoch_bytes = static_cast<uint32_t>(comm.push_bytes);
params.epoch_bytes = static_cast<uint32_t>(push.slot_bytes);
params.num_clean_up_count = num_clean;
const auto needed_buffer_bytes = static_cast<int64_t>(num_tokens) * 2 * sizeof(float);
RuntimeCheck(comm.world_size == kNumGPU, "Number of GPUs mismatch");
RuntimeCheck(push.world_size == kNumGPU, "Number of GPUs mismatch");
RuntimeCheck(std::bit_cast<intptr_t>(params.q_ptr) % 16 == 0, "q pointer is not properly aligned");
RuntimeCheck(std::bit_cast<intptr_t>(params.k_ptr) % 16 == 0, "k pointer is not properly aligned");
RuntimeCheck(std::bit_cast<intptr_t>(params.q_weight) % 16 == 0, "q_weight pointer is not properly aligned");
RuntimeCheck(std::bit_cast<intptr_t>(params.k_weight) % 16 == 0, "k_weight pointer is not properly aligned");
RuntimeCheck(needed_buffer_bytes <= comm.push_bytes, "Push buffer is too small");
RuntimeCheck(needed_buffer_bytes <= push.slot_bytes, "Push buffer is too small");
LaunchKernel(num_blocks, num_threads, device) //
.enable_pdl(kUsePDL)(kernel, params, comm.push_counter);
.enable_pdl(kUsePDL)(kernel, params, push.counter);
}
static uint32_t get_max_occupancy() {
python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh
+3 -2
View File
@@ -15,6 +15,7 @@ limitations under the License.
#pragma once
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.h>
@@ -160,7 +161,7 @@ void launch_sm120_fp8_blockwise_scaled_mm(
}
size_t workspace_size = gemm_op.get_workspace_size(args);
auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device());
auto workspace_tensor = host::ffi::alloc_workspace_tensor(workspace_size, a.device());
void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr();
auto init_status = gemm_op.initialize(args, workspace, stream);
@@ -379,7 +380,7 @@ void launch_sm120_fp8_blockwise_scaled_mm_swapab(
}
size_t workspace_size = gemm_op.get_workspace_size(args);
auto workspace_tensor = alloc_workspace_tensor(workspace_size, a.device());
auto workspace_tensor = host::ffi::alloc_workspace_tensor(workspace_size, a.device());
void* workspace = (workspace_size == 0) ? nullptr : workspace_tensor.data_ptr();
auto init_status = gemm_op.initialize(args, workspace, stream);
python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh
+62 -93
View File
@@ -9,9 +9,12 @@
#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/container/tensor.h>
#include <tvm/ffi/extra/stl.h>
#include "../../distributed/custom_all_reduce.cuh"
#include <algorithm>
#include <array>
#include <cfloat>
@@ -22,14 +25,14 @@
namespace sglang {
namespace ptx {
namespace device::ptx {
// ---- bulk 1D TMA (PTX ISA §9.7.9.25) ---------------------------------------
// global -> shared::cluster, completed by an smem mbarrier. Arm `bar` with
// `mbar_arrive_expect_tx(bar, bytes)` before issuing; `bytes` and both
// endpoints must be 16-byte aligned.
static SGL_DEVICE void cp_async_bulk_1d_load(void* smem_dst, const void* gmem_src, uint32_t bytes, uint64_t* bar) {
SGL_DEVICE void cp_async_bulk_1d_load(void* smem_dst, const void* gmem_src, uint32_t bytes, uint64_t* bar) {
asm volatile(
"cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes"
" [%0], [%1], %2, [%3];" ::"r"(to_shared(smem_dst)),
@@ -41,21 +44,21 @@ static SGL_DEVICE void cp_async_bulk_1d_load(void* smem_dst, const void* gmem_sr
// Publish mbarrier initialization from the generic proxy before an async engine
// uses the barrier.
static SGL_DEVICE void fence_mbarrier_init() {
SGL_DEVICE void fence_mbarrier_init() {
asm volatile("fence.mbarrier_init.release.cluster;");
}
// ---- warp / warp-group sync (PTX ISA §9.7.4, §9.7.12.6, §9.7.13) -----------
// Partial-CTA rendezvous. `id` must be in [1, 15]; barrier 0 is reserved for
// the full-CTA barrier behind __syncthreads().
static SGL_DEVICE void named_barrier_sync(uint32_t id, uint32_t num_threads) {
// Partial-CTA rendezvous. `id` must be in [1, 15]; pull 0 is reserved for
// the full-CTA pull behind __syncthreads().
SGL_DEVICE void named_barrier_sync(uint32_t id, uint32_t num_threads) {
asm volatile("bar.sync %0, %1;" ::"r"(id), "r"(num_threads) : "memory");
}
// True on exactly one lane of the issuing warp — guards single-issuer sites
// (mbar init, TMA issue, MMA issue, TMEM alloc) without gating on lane_id.
static SGL_DEVICE bool elect_one() {
SGL_DEVICE bool elect_one() {
uint32_t pred;
asm volatile(
"{\n\t.reg .pred p;\n\t"
@@ -81,14 +84,14 @@ static SGL_DEVICE bool elect_one() {
// source. For a symmetric cap, `__launch_bounds__(NUM_THREADS, 1)` is cleaner
// and measured faster on B100/B300.
template <int N>
static SGL_DEVICE void setmaxnreg_dec() {
SGL_DEVICE void setmaxnreg_dec() {
static_assert(N >= 24 && N <= 256, "setmaxnreg N must be in [24, 256]");
static_assert((N & 7) == 0, "setmaxnreg N must be a multiple of 8");
asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" ::"n"(N));
}
template <int N>
static SGL_DEVICE void setmaxnreg_inc() {
SGL_DEVICE void setmaxnreg_inc() {
static_assert(N >= 24 && N <= 256, "setmaxnreg N must be in [24, 256]");
static_assert((N & 7) == 0, "setmaxnreg N must be a multiple of 8");
asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" ::"n"(N));
@@ -103,23 +106,23 @@ static SGL_DEVICE void setmaxnreg_inc() {
// Each warp can only touch its own 32-lane TMEM band (§9.7.16.8.1): warp 0 ->
// lanes 0-31, warp 1 -> 32-63, and so on. Use `tcgen05_wait_st` /
// `tcgen05_wait_ld` before consuming the other side of a store / drain.
static SGL_DEVICE void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) {
SGL_DEVICE void tcgen05_alloc(uint32_t smem_addr_for_taddr, uint32_t n_cols) {
asm volatile(
"tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" ::"r"(smem_addr_for_taddr), "r"(n_cols));
}
static SGL_DEVICE void tcgen05_dealloc(uint32_t taddr, uint32_t n_cols) {
SGL_DEVICE void tcgen05_dealloc(uint32_t taddr, uint32_t n_cols) {
asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" ::"r"(taddr), "r"(n_cols));
}
static SGL_DEVICE void tcgen05_relinquish() {
SGL_DEVICE void tcgen05_relinquish() {
asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;");
}
// .32x32b.x8: 8 b32 per lane = 8 TMEM columns. Per-lane 8 FP32 -> 4 bf16x2
// packs = one int4, the natural fit for a BF16 epilogue moving a column band
// with 16-byte smem accesses.
static SGL_DEVICE void tcgen05_ld_32x32b_x8(
SGL_DEVICE void tcgen05_ld_32x32b_x8(
uint32_t taddr,
uint32_t& r0,
uint32_t& r1,
@@ -136,11 +139,11 @@ static SGL_DEVICE void tcgen05_ld_32x32b_x8(
: "r"(taddr));
}
static SGL_DEVICE void tcgen05_ld_32x32b_x8(uint32_t taddr, uint32_t* dst) {
SGL_DEVICE void tcgen05_ld_32x32b_x8(uint32_t taddr, uint32_t* dst) {
tcgen05_ld_32x32b_x8(taddr, dst[0], dst[1], dst[2], dst[3], dst[4], dst[5], dst[6], dst[7]);
}
static SGL_DEVICE void tcgen05_st_32x32b_x8(uint32_t taddr, const uint32_t* src) {
SGL_DEVICE void tcgen05_st_32x32b_x8(uint32_t taddr, const uint32_t* src) {
asm volatile(
"tcgen05.st.sync.aligned.32x32b.x8.b32 "
" [%8], {%0, %1, %2, %3, %4, %5, %6, %7};"
@@ -156,11 +159,11 @@ static SGL_DEVICE void tcgen05_st_32x32b_x8(uint32_t taddr, const uint32_t* src)
"r"(taddr));
}
static SGL_DEVICE void tcgen05_wait_st() {
SGL_DEVICE void tcgen05_wait_st() {
asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory");
}
} // namespace ptx
} // namespace device::ptx
struct AttnResTMAParams {
const bf16_t* __restrict__ prefix_sum; // [T, H]
@@ -179,7 +182,7 @@ struct AttnResTMAParams {
const bf16_t* residual;
bf16_t* prefix_out;
device::distributed::Semaphore* sem_local;
uint8_t* sem_mc;
device::distributed::Semaphore* sem_mc;
uint8_t* output_mc;
uint32_t world_size;
uint32_t rank;
@@ -194,7 +197,7 @@ struct KimiK3AttnResTrait {
static constexpr int64_t kDim = kDim_;
static constexpr int64_t kTile = 1024; // one warp-group-wide 16B sweep
static constexpr uint32_t kNumRows = kNumBankRows_; // bank rows; +1 prefix row
static constexpr uint32_t kChunkRows = kChunkRows_; // rows per chunk (one barrier pair per chunk)
static constexpr uint32_t kChunkRows = kChunkRows_; // rows per chunk (one pull pair per chunk)
// Chunk slots in the smem ring. Frozen at 2 (double buffering): 1 stalls
// the producer behind the consumers (~10% slower), >2 gains nothing and
// costs smem at small T.
@@ -224,7 +227,7 @@ struct KimiK3AttnResTrait {
// TMEM: per group, kTmemColsPerGroup columns of cw then of ow.
static constexpr uint32_t kTmemColsPerGroup = 32;
static constexpr uint32_t kTmemCols = 2 * kNumGroups * kTmemColsPerGroup;
static constexpr uint32_t kConsumerBarId = 1; // barrier 0 stays __syncthreads'
static constexpr uint32_t kConsumerBarId = 1; // pull 0 stays __syncthreads'
static_assert(kDim % kTile == 0, "kDim must be a whole number of tiles");
static_assert(kTile == kGroupThreads * kVecElems, "a tile is one group-wide 16B sweep");
@@ -300,7 +303,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
if (global_chunks >= kNumStages) {
ptx::mbar_wait_parity(&smem->bar_free[slot], phase ^ 1);
}
// One barrier per chunk; each row still gets its own bulk copy.
// One pull per chunk; each row still gets its own bulk copy.
ptx::mbar_arrive_expect_tx(&smem->bar_full[slot], an * kRowBytes);
#pragma unroll
for (uint32_t r = 0; r < an; ++r) {
@@ -544,7 +547,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
if (params.output_mc != nullptr) {
const auto global_token = static_cast<int64_t>(params.rank) * params.num_tokens + token;
const auto global_vid = global_token * (kDim / kVecElems) + row_vid;
st_multimem_16B(out_vec, params.output_mc, global_vid);
ptx::st_multimem_16B(out_vec, params.output_mc, global_vid);
} else {
out_vec.store(out_ptr, row_vid);
}
@@ -566,27 +569,6 @@ __global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
Trait::forward(params, reinterpret_cast<typename Trait::Smem*>(smem_raw));
}
SGL_DEVICE uint32_t* attn_res_sem_mc_flag(uint8_t* sem_mc, uint32_t block) {
static_assert(sizeof(device::distributed::Semaphore) == 128);
return reinterpret_cast<uint32_t*>(sem_mc + block * sizeof(device::distributed::Semaphore));
}
SGL_DEVICE void attn_res_sem_arrive_relaxed(uint32_t* flag) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
#else
assert(false && "multimem red requires Hopper or later");
#endif
}
SGL_DEVICE void attn_res_sem_arrive_release(uint32_t* flag) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
#else
assert(false && "multimem red requires Hopper or later");
#endif
}
// Fused NVLS pull RS + local residual + attention-residual aggregation.
// The entry/exit barriers make local o_proj writes visible before the
// producer's multimem reduction and preserve the shared pull-semaphore
@@ -594,15 +576,16 @@ SGL_DEVICE void attn_res_sem_arrive_release(uint32_t* flag) {
template <typename Trait, uint32_t kOccupancy>
__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
attn_res_fused_pull_rs_kernel(const __grid_constant__ AttnResTMAParams params) {
// The window base lives in shared memory, not in the barrier object: this
// kernel's body is register-budgeted (see kConsumerRegs / setmaxnreg), so
// keeping the object live across Trait::forward would spill.
__shared__ uint32_t exit_base;
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size);
exit_base = reserved + params.world_size;
{
const auto barrier =
device::distributed::McBarrier(params.sem_local, params.sem_mc, params.world_size, /*num_arrives=*/2);
device::PDLWaitPrimary<true>();
attn_res_sem_arrive_relaxed(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_relaxed() - reserved < params.world_size)
;
barrier.arrive_relaxed(/*n=*/0);
if (threadIdx.x == 0) exit_base = barrier.window() + params.world_size;
}
__syncthreads();
@@ -619,7 +602,7 @@ __global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
params.residual == nullptr ? nullptr : params.residual + static_cast<int64_t>(token) * Trait::kDim;
for (uint32_t vid = threadIdx.x; vid < kRowVecs; vid += blockDim.x) {
pull_vec_t vec;
ld_multimem_16B(vec, input_mc, vid);
device::ptx::ld_multimem_16B(vec, input_mc, vid);
if (residual != nullptr) {
pull_vec_t res;
res.load(residual, vid);
@@ -638,12 +621,7 @@ __global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
Trait::forward(params, reinterpret_cast<typename Trait::Smem*>(smem_raw));
__syncthreads();
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
attn_res_sem_arrive_release(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_acquire() - exit_base < params.world_size)
;
}
device::distributed::McBarrier::arrive_at<true>(params.sem_local, params.sem_mc, params.world_size, exit_base);
}
// Local attention-residual aggregation + direct AG epilogue. The consumer
@@ -653,14 +631,12 @@ __global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
template <typename Trait, uint32_t kOccupancy>
__global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
attn_res_fused_direct_ag_kernel(const __grid_constant__ AttnResTMAParams params) {
__shared__ uint32_t exit_base;
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size);
exit_base = reserved + params.world_size;
attn_res_sem_arrive_relaxed(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_relaxed() - reserved < params.world_size)
;
__shared__ uint32_t exit_base; // see the note in the pull-RS kernel above
{
const auto barrier =
device::distributed::McBarrier(params.sem_local, params.sem_mc, params.world_size, /*num_arrives=*/2);
barrier.arrive_relaxed(/*n=*/0);
if (threadIdx.x == 0) exit_base = barrier.window() + params.world_size;
}
__syncthreads();
@@ -668,12 +644,7 @@ __global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
Trait::forward(params, reinterpret_cast<typename Trait::Smem*>(smem_raw));
__syncthreads();
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
attn_res_sem_arrive_release(attn_res_sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_acquire() - exit_base < params.world_size)
;
}
device::distributed::McBarrier::arrive_at<true>(params.sem_local, params.sem_mc, params.world_size, exit_base);
}
using host::distributed::CommunicatorRef;
@@ -792,10 +763,9 @@ struct AttnResFusedTmaKernel {
int64_t nvb,
double eps,
int64_t input_mc_ptr,
int64_t sem_mc_ptr,
int64_t max_blocks) {
using namespace host;
const auto& data = *ref.get();
const auto& pull = ref.get()->get_pull_obj();
auto GT_ = SymbolicSize{"global_tokens"};
auto T_ = SymbolicSize{"local_tokens"};
auto H_ = SymbolicSize{"hidden_size"};
@@ -815,12 +785,12 @@ struct AttnResFusedTmaKernel {
const auto num_tokens = static_cast<int64_t>(T_.unwrap());
const auto H = static_cast<int64_t>(H_.unwrap());
const auto NB = static_cast<int64_t>(NB_.unwrap());
RuntimeCheck(data.world_size > 1, "fused pull RS requires world_size > 1");
RuntimeCheck(global_tokens == num_tokens * data.world_size, "global tokens must equal local tokens * world size");
RuntimeCheck(pull.world_size > 1, "fused pull RS requires world_size > 1");
RuntimeCheck(global_tokens == num_tokens * pull.world_size, "global tokens must equal local tokens * world size");
RuntimeCheck(H == kDim, "fused pull RS: H must be ", kDim, ", got ", H);
RuntimeCheck(1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, "fused pull RS: invalid nvb=", nvb, " NB=", NB);
RuntimeCheck(input_mc_ptr != 0, "fused pull RS requires multicast input");
RuntimeCheck(sem_mc_ptr != 0, "fused pull RS requires multicast semaphores");
RuntimeCheck(pull.mc_semaphore != nullptr, "fused pull RS requires a multicast-capable pull plane");
RuntimeCheck(max_blocks > 0, "fused pull RS requires max_blocks > 0");
if (num_tokens == 0) return;
@@ -837,7 +807,7 @@ struct AttnResFusedTmaKernel {
{static_cast<int64_t>(num_sm) * kOccupancy,
num_tokens,
max_blocks,
static_cast<int64_t>(data.num_pull_blocks)});
static_cast<int64_t>(ref.get()->get_pull_blocks())});
const auto local_elems = num_tokens * H;
const auto params = AttnResTMAParams{
.prefix_sum = static_cast<const bf16_t*>(prefix_out.data_ptr()),
@@ -847,14 +817,14 @@ struct AttnResFusedTmaKernel {
.out = static_cast<bf16_t*>(out.data_ptr()),
.prefix_dst = nullptr,
.input_mc = reinterpret_cast<const uint8_t*>(static_cast<uintptr_t>(input_mc_ptr)) +
data.rank * local_elems * sizeof(bf16_t),
pull.rank * local_elems * sizeof(bf16_t),
.residual = residual.has_value() ? static_cast<const bf16_t*>(residual.value().data_ptr()) : nullptr,
.prefix_out = static_cast<bf16_t*>(prefix_out.data_ptr()),
.sem_local = data.pull_semaphores[data.rank],
.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr)),
.sem_local = pull.semaphores[pull.rank],
.sem_mc = pull.mc_semaphore,
.output_mc = nullptr,
.world_size = data.world_size,
.rank = data.rank,
.world_size = pull.world_size,
.rank = pull.rank,
.stride_bm = NB * H,
.eps = static_cast<float>(eps),
.num_tokens = static_cast<uint32_t>(num_tokens),
@@ -872,11 +842,10 @@ struct AttnResFusedTmaKernel {
int64_t nvb,
double eps,
int64_t output_mc_ptr,
int64_t sem_mc_ptr,
int64_t max_blocks,
bool write_prefix) {
using namespace host;
const auto& data = *ref.get();
const auto& pull = ref.get()->get_pull_obj();
auto T_ = SymbolicSize{"local_tokens"};
auto GT_ = SymbolicSize{"global_tokens"};
auto H_ = SymbolicSize{"hidden_size"};
@@ -893,13 +862,13 @@ struct AttnResFusedTmaKernel {
const auto global_tokens = static_cast<int64_t>(GT_.unwrap());
const auto H = static_cast<int64_t>(H_.unwrap());
const auto NB = static_cast<int64_t>(NB_.unwrap());
RuntimeCheck(data.world_size > 1, "fused direct AG requires world_size > 1");
RuntimeCheck(global_tokens == num_tokens * data.world_size, "global tokens must equal local tokens * world size");
RuntimeCheck(pull.world_size > 1, "fused direct AG requires world_size > 1");
RuntimeCheck(global_tokens == num_tokens * pull.world_size, "global tokens must equal local tokens * world size");
RuntimeCheck(H == kDim, "fused direct AG: H must be ", kDim, ", got ", H);
RuntimeCheck(1 <= nvb && nvb <= kMaxBankRows && nvb <= NB, "fused direct AG: invalid nvb=", nvb, " NB=", NB);
RuntimeCheck(!write_prefix || nvb < NB, "fused direct AG: write_prefix targets bank row nvb, needs nvb < NB");
RuntimeCheck(output_mc_ptr != 0, "fused direct AG requires multicast output");
RuntimeCheck(sem_mc_ptr != 0, "fused direct AG requires multicast semaphores");
RuntimeCheck(pull.mc_semaphore != nullptr, "fused direct AG requires a multicast-capable pull plane");
RuntimeCheck(max_blocks > 0, "fused direct AG requires max_blocks > 0");
if (num_tokens == 0) return;
@@ -915,7 +884,7 @@ struct AttnResFusedTmaKernel {
{static_cast<int64_t>(num_sm) * kOccupancy,
num_tokens,
max_blocks,
static_cast<int64_t>(data.num_pull_blocks)});
static_cast<int64_t>(ref.get()->get_pull_blocks())});
const auto params = AttnResTMAParams{
.prefix_sum = static_cast<const bf16_t*>(prefix_sum.data_ptr()),
.bank = static_cast<const bf16_t*>(bank.data_ptr()),
@@ -926,11 +895,11 @@ struct AttnResFusedTmaKernel {
.input_mc = nullptr,
.residual = nullptr,
.prefix_out = nullptr,
.sem_local = data.pull_semaphores[data.rank],
.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr)),
.sem_local = pull.semaphores[pull.rank],
.sem_mc = pull.mc_semaphore,
.output_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(output_mc_ptr)),
.world_size = data.world_size,
.rank = data.rank,
.world_size = pull.world_size,
.rank = pull.rank,
.stride_bm = NB * H,
.eps = static_cast<float>(eps),
.num_tokens = static_cast<uint32_t>(num_tokens),
python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh
+124 -159
View File
@@ -30,7 +30,10 @@
// the attn-res prefix sum — or absent) or the RMSNorm epilogue over the
// latent of the K3 latent|shared MoE buffer (*_norm variants).
//
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/math.cuh>
#include <sgl_kernel/runtime.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <sgl_kernel/warp.cuh>
@@ -38,15 +41,23 @@
#include <cooperative_groups.h>
// TODO: remove dependency on the custom_all_reduce, move out common utilities
#include "../../distributed/custom_all_reduce.cuh"
#include "ptx_sys.cuh"
#include <sgl_kernel/distributed/communicator.cuh>
#include <sgl_kernel/distributed/ptx.cuh>
#include <tvm/ffi/extra/stl.h>
namespace sglang {
// Same shape as gemm_ag / gemm_ar: pull the ptx_sys helpers in by name so the
// call sites below stay unqualified.
using device::distributed::multimem_red_add_relaxed;
using device::distributed::multimem_red_add_release;
// The shared vocabulary this file is written in, pulled in by name so the call
// sites below stay unqualified: see include/sgl_kernel/distributed/ptx.cuh and
// .../communicator.cuh.
using device::distributed::Counter;
using device::distributed::Semaphore;
using host::distributed::CommunicatorRef;
/// The 16 B staging vector, viewed as the 4 u32 words the lamport marker
/// protocol tests. See LamportTrait in distributed/communicator.cuh.
using Lamport = device::distributed::LamportTrait<bf16_t, 8, /*kAtom=*/4>;
struct FusionParams {
uint8_t* input; // tensor pointer (in place)
@@ -99,17 +110,11 @@ __global__ __launch_bounds__(1024, 1) void all_reduce_push_res_kernel(const __gr
const auto poll_ptr = params.push_ws_local + phase_stride_bytes;
// stage 1: multicast-push local data, remapping all-zero bf16x2 pairs
static_assert(fp_trait<bf16_t>::pos_zero == 0, "the empty marker is all-zero bits");
constexpr uint32_t kNegZeroPair = 0x8000u; // {-0.0, +0.0}: sum-neutral, non-zero
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
ld_global_16B(vec, params.input, vid);
auto& bits = *reinterpret_cast<uint4*>(&vec);
if (bits.x == 0) bits.x = kNegZeroPair;
if (bits.y == 0) bits.y = kNegZeroPair;
if (bits.z == 0) bits.z = kNegZeroPair;
if (bits.w == 0) bits.w = kNegZeroPair;
st_multimem_16B(vec, push_ptr, vid);
device::ptx::ld_global_16B(vec, params.input, vid);
Lamport::clear_pos_zero(vec.data());
device::ptx::st_multimem_16B(vec, push_ptr, vid);
}
// launch pdl early for low latency case
@@ -118,7 +123,7 @@ __global__ __launch_bounds__(1024, 1) void all_reduce_push_res_kernel(const __gr
// stage 2: poll all slots, reduce (+ residual), write back in place,
// re-establish the empty markers for the next same-phase round
vec_t zero_vec;
zero_vec.fill(bf16x2_t{get_pos_zero<bf16_t>(), get_pos_zero<bf16_t>()});
Lamport::fill_pos_zero(zero_vec.data());
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec[kWorldSize + kHasResidual];
if constexpr (kHasResidual) vec[kWorldSize].load(params.residual, vid);
@@ -126,22 +131,18 @@ __global__ __launch_bounds__(1024, 1) void all_reduce_push_res_kernel(const __gr
bool has_zero = false;
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid);
// the producer remapped all-zero pairs, so a written u32 is never
// 0: u32 == 0 <=> the 4B atom still holds the empty marker
const auto bits = *reinterpret_cast<const uint4*>(&vec[i]);
has_zero |= bits.x == 0;
has_zero |= bits.y == 0;
has_zero |= bits.z == 0;
has_zero |= bits.w == 0;
device::ptx::ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid);
// the producer remapped all-zero pairs, so a written atom is never 0:
// atom == 0 <=> the slot still holds the empty marker
has_zero |= Lamport::has_pos_zero(vec[i].data());
}
if (!has_zero) break;
} while (true);
const auto out_vec = reduce(vec); // fp32 accumulation over 8(+1) inputs
st_global_16B(out_vec, params.input, vid);
const auto out_vec = device::reduce_vec(vec); // fp32 accumulation over 8(+1) inputs
device::ptx::st_global_16B(out_vec, params.input, vid);
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid);
device::ptx::st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid);
}
}
@@ -285,20 +286,15 @@ __global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClus
// stage 1: multicast staging (grid-stride); kFinalize computes each vector
// in place of the load
static_assert(fp_trait<bf16_t>::pos_zero == 0, "the empty marker is all-zero bits");
for (auto vid = global_tid; vid < num_vecs; vid += num_threads) {
vec_t vec;
if constexpr (kFinalize) {
vec = finalize_vec(params, vid);
} else {
ld_global_16B(vec, params.input, vid);
ptx::ld_global_16B(vec, params.input, vid);
}
auto& bits = *reinterpret_cast<uint4*>(&vec);
if (bits.x == 0) bits.x = fp_trait<bf16_t>::neg_zero;
if (bits.y == 0) bits.y = fp_trait<bf16_t>::neg_zero;
if (bits.z == 0) bits.z = fp_trait<bf16_t>::neg_zero;
if (bits.w == 0) bits.w = fp_trait<bf16_t>::neg_zero;
st_multimem_16B(vec, push_ptr, vid);
Lamport::clear_pos_zero(vec.data());
ptx::st_multimem_16B(vec, push_ptr, vid);
}
// stage 2: one row per cluster pass (the bumper cluster owns no rows)
@@ -307,7 +303,7 @@ __global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClus
vec_t w;
w.load(params.norm_weight, cluster_rank * kBlockSize + tx);
vec_t zero_vec;
zero_vec.fill(bf16x2_t{get_pos_zero<bf16_t>(), get_pos_zero<bf16_t>()});
Lamport::fill_pos_zero(zero_vec.data());
__shared__ alignas(8) float smem_raw[2][kClusterSize][kNumWarps];
uint32_t parity = 0;
@@ -321,12 +317,8 @@ __global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClus
bool has_zero = false;
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid);
const auto bits = *reinterpret_cast<const uint4*>(&vec[i]);
has_zero |= bits.x == 0;
has_zero |= bits.y == 0;
has_zero |= bits.z == 0;
has_zero |= bits.w == 0;
ptx::ld_relaxed_16B(vec[i], poll_ptr + i * stride_bytes, vid);
has_zero |= Lamport::has_pos_zero(vec[i].data());
}
if (!has_zero) break;
} while (true);
@@ -369,13 +361,13 @@ __global__ __launch_bounds__(kNormRowVecs / kClusterSize) __cluster_dims__(kClus
out_vec[j] = cast<bf16x2_t>(fp32x2_t{a * norm_factor * wa, b * norm_factor * wb});
}
} else {
out_vec = reduce(vec);
out_vec = device::reduce_vec(vec);
}
st_global_16B(out_vec, params.input, vid);
ptx::st_global_16B(out_vec, params.input, vid);
#pragma unroll
for (uint32_t i = 0; i < kWorldSize; ++i) {
st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid);
ptx::st_global_16B(zero_vec, poll_ptr + i * stride_bytes, vid);
}
}
@@ -394,7 +386,7 @@ struct PullParams {
uint8_t* input_mc; // multicast VA of the symmetric input
const uint8_t* residual; // may be null (compile-time kHasResidual selects)
Semaphore* sem_local; // this rank's v2 pull semaphores (poll side)
uint8_t* sem_mc; // multicast VA of the pull-semaphore region
Semaphore* sem_mc; // multicast VA of the pull-semaphore region
uint32_t rank;
uint32_t world_size;
uint32_t num_vecs; // 16B vectors
@@ -413,53 +405,6 @@ struct PullParams {
// (same aggregate effect: every rank's flag gains world_size arrivals per
// phase). Identical memory effects per call, so both kernel families share
// the slots freely (single-stream calls are serialized).
//
// The multicast alias of Semaphore::m_flag (the struct's first member):
SGL_DEVICE uint32_t* pull_sem_mc_flag(uint8_t* sem_mc, uint32_t block) {
static_assert(sizeof(Semaphore) == 128);
return reinterpret_cast<uint32_t*>(sem_mc + block * sizeof(Semaphore));
}
// enter barrier (relaxed): reserve this call's flag window, signal arrival
// with one multicast red, poll the local flag until all world_size arrivals
// landed — every rank's producer has finished writing the input. The
// reservation atomicAdd sits BEFORE the PDL wait: it is safe there (the
// previous same-slot call's reservation completed at ITS enter, which
// precedes its launch_dependents and hence this kernel's start, so windows
// are handed out in stream order) and it keeps the RMW latency off the
// post-wait critical path. The red must stay AFTER the wait — it asserts
// the producer grid has flushed. Returns the window base for the exit
// barrier — meaningful in thread 0 only, the sole barrier poller.
template <bool kUsePDL>
SGL_DEVICE uint32_t pull_barrier_enter(const PullParams& params) {
uint32_t current = 0;
if (threadIdx.x == 0) {
const auto semaphore = &params.sem_local[blockIdx.x];
const auto reserved = semaphore->counter_ptr()->inc(2 * params.world_size);
current = reserved + params.world_size;
device::PDLWaitPrimary<kUsePDL>();
multimem_red_add_relaxed(pull_sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_relaxed() - reserved < params.world_size)
;
}
__syncthreads();
return current;
}
// exit barrier (release/acquire): every peer has finished reading my buffer
// (and, for 2shot, its broadcast into it is visible) before my next kernel
// may touch it. Mirrors AllReducePullImpl::sync_exit_pull<true>.
template <bool kUsePDL>
SGL_DEVICE void pull_barrier_exit(const PullParams& params, uint32_t current) {
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
if (threadIdx.x == 0) {
const auto semaphore = &params.sem_local[blockIdx.x];
multimem_red_add_release(pull_sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_acquire() - current < params.world_size)
;
}
}
// One pipelined pass at width kWidth (kWidth multimem loads in flight per
// thread), then recurse to kWidth/2 for the remainder, down to a plain
@@ -475,7 +420,7 @@ pull_reduce_pass(uint32_t& vid, const uint32_t num_vecs, const uint32_t step, ui
vec_t vec[kWidth];
#pragma unroll
for (uint32_t u = 0; u < kWidth; ++u) {
ld_multimem_16B(vec[u], mc_ptr, vid + u * step);
device::ptx::ld_multimem_16B(vec[u], mc_ptr, vid + u * step);
}
if constexpr (kHasResidual) {
#pragma unroll
@@ -490,7 +435,7 @@ pull_reduce_pass(uint32_t& vid, const uint32_t num_vecs, const uint32_t step, ui
}
#pragma unroll
for (uint32_t u = 0; u < kWidth; ++u) {
st_multimem_16B(vec[u], mc_ptr, vid + u * step);
device::ptx::st_multimem_16B(vec[u], mc_ptr, vid + u * step);
}
}
if constexpr (kWidth > 1) {
@@ -505,7 +450,14 @@ __launch_bounds__(kPullBlockSize, 1) void all_reduce_pull_res_kernel(const __gri
const auto tx = threadIdx.x;
const auto bx = blockIdx.x;
const auto barrier_window = pull_barrier_enter<kUsePDL>(params);
// Reserve the window before the PDL wait, signal after it: the reservation's
// RMW latency stays off the post-wait critical path, while the signal must
// follow the wait because it asserts the producer grid has flushed.
const auto barrier =
device::distributed::McBarrier(params.sem_local, params.sem_mc, params.world_size, /*num_arrives=*/2);
device::PDLWaitPrimary<kUsePDL>();
barrier.arrive_relaxed(/*n=*/0);
__syncthreads();
// this rank's shard of the 16B-vector range
const auto r = params.rank;
@@ -524,7 +476,11 @@ __launch_bounds__(kPullBlockSize, 1) void all_reduce_pull_res_kernel(const __gri
auto vid = bx * kPullBlockSize + tx;
pull_reduce_pass<kUnroll, kHasResidual>(vid, num_vecs, step, mc_ptr, res_ptr);
pull_barrier_exit<kUsePDL>(params, barrier_window);
// exit barrier: every peer has finished reading my buffer (and, for 2shot,
// its broadcast into it is visible) before my next kernel may touch it.
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
barrier.arrive_rel_acq(/*n=*/1);
}
// Fused RMSNorm over the latent of the K3 latent|shared MoE buffer: the
@@ -544,7 +500,14 @@ __launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid
const auto tx = threadIdx.x;
const auto bx = blockIdx.x;
const auto barrier_window = pull_barrier_enter<kUsePDL>(params);
// Reserve the window before the PDL wait, signal after it: the reservation's
// RMW latency stays off the post-wait critical path, while the signal must
// follow the wait because it asserts the producer grid has flushed.
const auto barrier =
device::distributed::McBarrier(params.sem_local, params.sem_mc, params.world_size, /*num_arrives=*/2);
device::PDLWaitPrimary<kUsePDL>();
barrier.arrive_relaxed(/*n=*/0);
__syncthreads();
// this rank's shard of the row range
const auto num_rows = params.num_vecs / kNormRowVecs;
@@ -568,7 +531,7 @@ __launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid
vec_t vec[kUnroll];
#pragma unroll
for (uint32_t u = 0; u < kUnroll; ++u) {
if (u < cnt) ld_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
if (u < cnt) ptx::ld_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
}
// norm rows: push this warp's partial sum of squares to smem; non-norm
// rows (the shared 2/3) don't wait for the barrier — store right away
@@ -587,7 +550,7 @@ __launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid
sum_of_squares = warp::reduce_sum(sum_of_squares);
if (lane == 0) sm[u][warp] = sum_of_squares;
} else {
st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
ptx::st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
}
}
__syncthreads();
@@ -606,14 +569,28 @@ __launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid
const auto [wa, wb] = cast<fp32x2_t>(wvec[j]);
vec[u][j] = cast<bf16x2_t>(fp32x2_t{a * norm_factor * wa, b * norm_factor * wb});
}
st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
ptx::st_multimem_16B(vec[u], params.input_mc, vid0 + u * kNormRowVecs);
}
}
pull_barrier_exit<kUsePDL>(params, barrier_window);
// exit barrier: every peer has finished reading my buffer (and, for 2shot,
// its broadcast into it is visible) before my next kernel may touch it.
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
barrier.arrive_rel_acq(/*n=*/1);
}
// Host entry points
inline auto choose_block_size(uint32_t num_threads) -> 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, 512u}) {
if (host::div_ceil(num_threads, block_size) <= kNumSM) return block_size;
}
return 1024u;
}
template <uint32_t kWorldSize, bool kUsePDL>
struct AllReduceFusionKernel {
@@ -624,8 +601,9 @@ struct AllReduceFusionKernel {
static constexpr auto res_push_kernel = all_reduce_push_res_kernel<kWorldSize, kHasResidual, kUsePDL>;
static FusionParams
make_params(const host::distributed::CommunicatorObj& data, TensorView input, std::optional<TensorView> residual) {
make_params(const host::distributed::CommunicatorObj& comm, TensorView input, std::optional<TensorView> residual) {
using namespace host;
const auto& push = comm.get_push_obj();
SymbolicSize N = {"num_elements"};
SymbolicDevice device;
device.set_options<kDLCUDA>();
@@ -642,16 +620,17 @@ struct AllReduceFusionKernel {
.verify(input);
}
const auto num_elems = N.unwrap();
CHECK_HOST(data.world_size == kWorldSize);
CHECK_HOST(push.world_size == kWorldSize);
CHECK_HOST(num_elems > 0 && num_elems % 8 == 0);
CHECK_HOST(push.mc_workspace != nullptr) << "the fused push needs a multicast-capable push plane";
FusionParams params{};
params.input = static_cast<uint8_t*>(input.data_ptr());
params.residual = residual.has_value() ? static_cast<const uint8_t*>(residual.value().data_ptr()) : nullptr;
params.push_ws_mc = nullptr;
params.push_ws_local = data.push_workspaces[data.rank];
params.push_counter = data.push_counter;
params.push_buffer_stride = data.push_bytes;
params.rank = data.rank;
params.push_ws_mc = push.mc_workspace;
params.push_ws_local = push.workspaces[push.rank];
params.push_counter = push.counter;
params.push_buffer_stride = push.slot_bytes;
params.rank = push.rank;
params.num_vecs = static_cast<uint32_t>(num_elems / 8);
return params;
}
@@ -661,13 +640,13 @@ struct AllReduceFusionKernel {
/// (K3 uses [N latent rows | 2N shared rows] with num_norm_rows = N, or a
/// latent-only [N, 3584] tensor with num_norm_rows = N).
static FusionParams make_params_norm(
const host::distributed::CommunicatorObj& data,
const host::distributed::CommunicatorObj& comm,
TensorView input,
TensorView weight,
float eps,
int64_t num_norm_rows) {
using namespace host;
auto params = make_params(data, input, std::nullopt);
auto params = make_params(comm, input, std::nullopt);
SymbolicDevice device;
device.set_options<kDLCUDA>();
TensorMatcher({kNormDim}).with_dtype<bf16_t>().with_device<kDLCUDA>(device).verify(weight);
@@ -679,18 +658,18 @@ struct AllReduceFusionKernel {
params.norm_weight = static_cast<const uint8_t*>(weight.data_ptr());
params.norm_eps = static_cast<float>(eps);
params.num_norm_rows = static_cast<uint32_t>(num_norm_rows);
params.num_push_counters = data.num_push_blocks;
params.num_push_counters = comm.get_push_obj().num_blocks;
return params;
}
// Shared pull validation; the reduce is in place on the symmetric input.
static PullParams make_pull_params(
const host::distributed::CommunicatorObj& data,
const host::distributed::CommunicatorObj& comm,
TensorView input,
std::optional<TensorView> residual,
int64_t input_mc_ptr,
int64_t sem_mc_ptr) {
int64_t input_mc_ptr) {
using namespace host;
const auto& pull = comm.get_pull_obj();
SymbolicSize N = {"num_elements"};
SymbolicDevice device;
device.set_options<kDLCUDA>();
@@ -707,20 +686,20 @@ struct AllReduceFusionKernel {
.verify(input);
}
const auto num_elems = N.unwrap();
CHECK_HOST(data.world_size == kWorldSize);
CHECK_HOST(pull.world_size == kWorldSize);
CHECK_HOST(num_elems > 0 && num_elems % 8 == 0) << "numel must be a positive multiple of 8, got " << num_elems;
// headroom below 2^32 so the unrolled loop's `vid + (kUnroll-1)*step`
// arithmetic can never wrap around u32
CHECK_HOST(num_elems / 8 < (int64_t(1) << 31)) << "numel exceeds the 16B-vector limit";
CHECK_HOST(input_mc_ptr != 0) << "pull requires the input's multicast address";
CHECK_HOST(sem_mc_ptr != 0) << "pull requires the semaphores' multicast address";
CHECK_HOST(pull.mc_semaphore != nullptr) << "pull requires a multicast-capable pull plane";
PullParams params{};
params.input_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(input_mc_ptr));
params.residual = residual.has_value() ? static_cast<const uint8_t*>(residual.value().data_ptr()) : nullptr;
params.sem_local = data.pull_semaphores[data.rank];
params.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr));
params.rank = data.rank;
params.world_size = data.world_size;
params.sem_local = pull.semaphores[pull.rank];
params.sem_mc = pull.mc_semaphore;
params.rank = pull.rank;
params.world_size = pull.world_size;
params.num_vecs = static_cast<uint32_t>(num_elems / 8);
return params;
}
@@ -765,33 +744,26 @@ struct AllReduceFusionKernel {
}
public:
static void push_res(CommunicatorRef ref, TensorView input, std::optional<TensorView> residual, int64_t ws_mc_base) {
const auto& data = *ref.get();
auto params = make_params(data, input, residual);
CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace";
static void push_res(CommunicatorRef ref, TensorView input, std::optional<TensorView> residual) {
const auto& push = ref.get()->get_push_obj();
auto params = make_params(*ref.get(), input, residual);
const int64_t nbytes = int64_t(params.num_vecs) * 16;
CHECK_HOST(nbytes <= data.push_bytes)
<< "input size " << nbytes << " exceeds push workspace size " << data.push_bytes;
params.push_ws_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ws_mc_base));
CHECK_HOST(nbytes <= push.slot_bytes) << "input size " << nbytes << " exceeds push slot size " << push.slot_bytes;
const auto kernel = residual.has_value() ? res_push_kernel<true> : res_push_kernel<false>;
host::LaunchKernel(data.num_push_blocks, choose_block_size(params.num_vecs), input.device())
host::LaunchKernel(push.num_blocks, choose_block_size(params.num_vecs), input.device())
.enable_pdl(kUsePDL)(kernel, params);
}
static void push_norm(
CommunicatorRef ref, TensorView input, TensorView weight, float eps, int64_t num_norm_rows, int64_t ws_mc_base) {
static void push_norm(CommunicatorRef ref, TensorView input, TensorView weight, float eps, int64_t num_norm_rows) {
constexpr auto kClusterSize = 7;
const auto& data = *ref.get();
auto params = make_params_norm(data, input, weight, eps, num_norm_rows);
CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace";
const auto& push = ref.get()->get_push_obj();
auto params = make_params_norm(*ref.get(), input, weight, eps, num_norm_rows);
const int64_t nbytes = int64_t(params.num_vecs) * 16;
CHECK_HOST(nbytes <= data.push_bytes)
<< "input size " << nbytes << " exceeds push workspace size " << data.push_bytes;
params.push_ws_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ws_mc_base));
CHECK_HOST(nbytes <= push.slot_bytes) << "input size " << nbytes << " exceeds push slot size " << push.slot_bytes;
const auto num_rows = params.num_vecs / kNormRowVecs;
constexpr uint32_t kMaxClusters = 96;
const auto num_row_clusters = std::max<uint32_t>(std::min(num_rows, kMaxClusters), 1);
CHECK_HOST(num_row_clusters < data.num_push_blocks);
CHECK_HOST(num_row_clusters < push.num_blocks);
host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, input.device())
.enable_pdl(kUsePDL)(all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL>, params);
}
@@ -807,13 +779,12 @@ struct AllReduceFusionKernel {
TensorView permuted_idx,
TensorView expert_weights,
TensorView weight,
float eps,
int64_t ws_mc_base) {
float eps) {
using namespace host;
constexpr auto kClusterSize = 7;
const auto& data = *ref.get();
const auto& push = ref.get()->get_push_obj();
// every row of the latent-only output is normed
auto params = make_params_norm(data, out, weight, eps, out.size(0) / kNormDim);
auto params = make_params_norm(*ref.get(), out, weight, eps, out.size(0) / kNormDim);
const auto num_tokens = params.num_vecs / kNormRowVecs;
auto P = SymbolicSize{"num_permuted_rows"};
@@ -829,39 +800,35 @@ struct AllReduceFusionKernel {
TensorMatcher({TK}).with_dtype<int32_t>().with_device<kDLCUDA>(device).verify(permuted_idx);
CHECK_HOST(K.unwrap() == kFinTopK) << "finalize_push_norm is specialized for top_k = " << kFinTopK;
CHECK_HOST(ws_mc_base != 0) << "push requires a multicast-capable workspace";
const int64_t nbytes = int64_t(params.num_vecs) * 16;
CHECK_HOST(nbytes <= data.push_bytes)
<< "output size " << nbytes << " exceeds push workspace size " << data.push_bytes;
params.push_ws_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(ws_mc_base));
CHECK_HOST(nbytes <= push.slot_bytes) << "output size " << nbytes << " exceeds push slot size " << push.slot_bytes;
params.fin_gemm2 = static_cast<const uint8_t*>(gemm2_out.data_ptr());
params.fin_idx = static_cast<const uint8_t*>(permuted_idx.data_ptr());
params.fin_weights = static_cast<const uint8_t*>(expert_weights.data_ptr());
constexpr uint32_t kMaxClusters = 96;
const auto num_row_clusters = std::max<uint32_t>(std::min(num_tokens, kMaxClusters), 1);
CHECK_HOST(num_row_clusters < data.num_push_blocks);
CHECK_HOST(num_row_clusters < push.num_blocks);
host::LaunchKernel((num_row_clusters + 1) * kClusterSize, kNormRowVecs / kClusterSize, out.device())
.enable_pdl(kUsePDL)(
all_reduce_push_norm_cluster_kernel<kWorldSize, kClusterSize, kUsePDL, /*kFinalize=*/true>, params);
}
/// Low-SM NVLS pull (+ optional residual): in-place reduce-scatter +
/// broadcast on the symmetric input. `sem_mc_ptr` is the multicast VA of
/// the v2 pull-semaphore region; num_blocks — which must be uniform
/// across ranks per call — is clamped to the semaphore capacity.
/// broadcast on the symmetric input, whose multicast VA the caller passes
/// (it varies per call, unlike the barrier plane's own multicast base).
/// num_blocks -- which must be uniform across ranks per call -- is clamped to
/// the barrier plane's capacity.
static void pull_res(
CommunicatorRef ref,
TensorView input,
std::optional<TensorView> residual,
int64_t input_mc_ptr,
int64_t sem_mc_ptr,
int64_t num_blocks,
int64_t unroll) {
const auto& data = *ref.get();
const auto params = make_pull_params(data, input, residual, input_mc_ptr, sem_mc_ptr);
const auto params = make_pull_params(*ref.get(), input, residual, input_mc_ptr);
CHECK_HOST(num_blocks >= 1) << "invalid num_blocks: " << num_blocks;
num_blocks = std::min<int64_t>(num_blocks, data.num_pull_blocks);
num_blocks = std::min<int64_t>(num_blocks, ref.get()->get_pull_blocks());
if (residual.has_value()) {
launch_pull_res<true>(params, num_blocks, unroll, input.device());
} else {
@@ -880,13 +847,11 @@ struct AllReduceFusionKernel {
double eps,
int64_t num_norm_rows,
int64_t input_mc_ptr,
int64_t sem_mc_ptr,
int64_t num_blocks,
int64_t unroll) {
const auto& data = *ref.get();
auto params = make_pull_params(data, input, std::nullopt, input_mc_ptr, sem_mc_ptr);
auto params = make_pull_params(*ref.get(), input, std::nullopt, input_mc_ptr);
CHECK_HOST(num_blocks >= 1) << "invalid num_blocks: " << num_blocks;
num_blocks = std::min<int64_t>(num_blocks, data.num_pull_blocks);
num_blocks = std::min<int64_t>(num_blocks, ref.get()->get_pull_blocks());
using namespace host;
SymbolicDevice device;
device.set_options<kDLCUDA>();
python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh
+23 -34
View File
@@ -11,7 +11,6 @@
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/stl.h>
#include "ptx_sys.cuh"
#include <array>
#include <cstdint>
#include <optional>
@@ -22,7 +21,6 @@ namespace sglang {
namespace gemm_ag {
using device::distributed::Counter;
using device::distributed::multimem_store_relaxed;
constexpr uint32_t kWorld = 8; // TP world size
constexpr uint32_t kVecSize = 32 / sizeof(bf16_t); // 16 bf16 per 32B vector
@@ -34,7 +32,7 @@ constexpr uint32_t kSpinVec = 16 / sizeof(bf16_t); // 8 bf16 (16B) per consumer
struct ProducerParams {
uint8_t* ws_mc; // multicast VA of the push workspace base
Counter* counter; // per-block phase counters (READ only here)
uint32_t half_bytes; // bytes per phase half (world_size * push_bytes)
uint32_t half_bytes; // bytes per phase half (world_size * slot_bytes)
uint32_t rank;
};
@@ -73,7 +71,7 @@ __global__ __launch_bounds__(K / kVecSize) void gemm_ag_gemv_kernel(
// Every push-workspace consumer flips the WHOLE counter array each round
// (each has a tail loop up to num_counters), so all counters hold the same
// phase at this point and counter[0] is equivalent to counter[bx]. Reading a
// single counter is what frees the producer grid from num_push_blocks.
// single counter is what frees the producer grid from the counter array size.
const uint32_t phase = params.counter[0].get() & 1;
vec_t input_vec[M];
@@ -117,7 +115,7 @@ __global__ __launch_bounds__(K / kVecSize) void gemm_ag_gemv_kernel(
const uint32_t elem = (params.rank * M + m) * kNLocal + bx * N_SPLIT + n;
const auto base = reinterpret_cast<bf16_t*>(params.ws_mc + phase * params.half_bytes);
const auto dst = reinterpret_cast<uint32_t*>(base + elem);
multimem_store_relaxed(dst, bits);
ptx::multimem_store_relaxed(dst, bits);
}
PDLTriggerSecondary<kUsePDL>();
}
@@ -127,8 +125,8 @@ __global__ __launch_bounds__(K / kVecSize) void gemm_ag_gemv_kernel(
struct ConsumerParams {
uint8_t* ws_local; // LOCAL VA of the push workspace base (poll + reset)
Counter* counter; // per-block phase counters (read + flip)
uint32_t num_counters; // full counter array size (num_push_blocks)
uint32_t half_bytes; // bytes per phase half (world_size * push_bytes)
uint32_t num_counters; // full counter array size (PushPlane::num_blocks)
uint32_t half_bytes; // bytes per phase half (world_size * slot_bytes)
const bf16_t* b; // [M, N]
const bf16_t* c; // may be null
bf16_t* out; // [M, N]
@@ -175,10 +173,7 @@ __global__ void spin_add3_kernel(const __grid_constant__ ConsumerParams params)
// spin until all 4 packed pairs of the vector have landed
uint4 raw;
do {
asm volatile("ld.relaxed.gpu.global.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(raw.x), "=r"(raw.y), "=r"(raw.z), "=r"(raw.w)
: "l"(src)
: "memory");
ptx::ld_relaxed_16B(raw, src, 0);
} while (raw.x == 0 || raw.y == 0 || raw.z == 0 || raw.w == 0);
const auto& gathered = *reinterpret_cast<const vec_t*>(&raw);
vec_t out_vec;
@@ -215,7 +210,7 @@ struct GEMMAGKernel {
// Standalone GEMV at the same shape: 4.15 / 3.20 / 2.56 us, and 16 -> 4.42 us,
// so the trend is monotonic and 2 is the floor (the epilogue stores column
// pairs, so N_SPLIT must stay even). 8 used to be the largest grid that fit
// the old kNumProducerBlocks <= num_push_blocks bound; the producer now reads
// the old kNumProducerBlocks <= push.num_blocks bound; the producer now reads
// a single phase counter, so the grid is free and 112 blocks did not even
// fill one per SM.
static constexpr uint32_t N_SPLIT = 2;
@@ -233,15 +228,9 @@ struct GEMMAGKernel {
static constexpr auto kGemvTable = make_table(std::make_index_sequence<kMaxM>{});
static void
run(CommunicatorRef ref,
TensorView x,
TensorView weight,
TensorView b,
std::optional<TensorView> c,
TensorView out,
intptr_t ws_mc_base) {
run(CommunicatorRef ref, TensorView x, TensorView weight, TensorView b, std::optional<TensorView> c, TensorView out) {
using namespace host;
const auto& data = *ref.get();
const auto& push = ref.get()->get_push_obj();
auto M = SymbolicSize{"num_tokens"};
auto device = SymbolicDevice{};
@@ -255,19 +244,19 @@ struct GEMMAGKernel {
TensorMatcher({M, N}).with_dtype<bf16_t>().with_device(device).verify(out);
const auto num_tokens = static_cast<uint32_t>(M.unwrap());
CHECK_HOST(num_tokens >= 1 && num_tokens <= kMaxM);
CHECK_HOST(data.world_size == gemm_ag::kWorld) << "the kernel is compiled for TP" << gemm_ag::kWorld;
CHECK_HOST(ws_mc_base != 0) << "requires a multicast-capable workspace";
CHECK_HOST(int64_t(num_tokens) * kNLocal * 2 <= data.push_bytes)
<< "staging slice exceeds the push slot size " << data.push_bytes;
CHECK_HOST(push.world_size == gemm_ag::kWorld) << "the kernel is compiled for TP" << gemm_ag::kWorld;
CHECK_HOST(push.mc_workspace != nullptr) << "requires a multicast-capable push plane";
CHECK_HOST(int64_t(num_tokens) * kNLocal * 2 <= push.slot_bytes)
<< "staging slice exceeds the push slot size " << push.slot_bytes;
// The producer grid is no longer bound to the counter array: it reads only
// counter[0] (see gemm_ag_gemv_kernel). The consumer grid still is.
CHECK_HOST(data.num_push_blocks > 0) << "no push blocks available";
CHECK_HOST(push.num_blocks > 0) << "no push blocks available";
// producer: GEMV
const auto producer_params = gemm_ag::ProducerParams{
.ws_mc = reinterpret_cast<uint8_t*>(ws_mc_base),
.counter = data.push_counter,
.half_bytes = static_cast<uint32_t>(data.push_bytes * data.world_size),
.rank = data.rank,
.ws_mc = push.mc_workspace,
.counter = push.counter,
.half_bytes = static_cast<uint32_t>(push.slot_bytes * push.world_size),
.rank = push.rank,
};
LaunchKernel(kNumProducerBlocks, kGemvBlock, device.unwrap())
.enable_pdl(kUsePDL)(
@@ -278,10 +267,10 @@ struct GEMMAGKernel {
// consumer: spin + add3
const auto consumer_params = gemm_ag::ConsumerParams{
.ws_local = data.push_workspaces[data.rank],
.counter = data.push_counter,
.num_counters = data.num_push_blocks,
.half_bytes = static_cast<uint32_t>(data.push_bytes * data.world_size),
.ws_local = push.workspaces[push.rank],
.counter = push.counter,
.num_counters = push.num_blocks,
.half_bytes = static_cast<uint32_t>(push.slot_bytes * push.world_size),
.b = static_cast<const bf16_t*>(b.data_ptr()),
.c = c.has_value() ? static_cast<const bf16_t*>(c.value().data_ptr()) : nullptr,
.out = static_cast<bf16_t*>(out.data_ptr()),
@@ -289,7 +278,7 @@ struct GEMMAGKernel {
};
const auto num_vecs = num_tokens * N / gemm_ag::kSpinVec;
const auto num_consumers = host::div_ceil(num_vecs, gemm_ag::kSpinBlock);
CHECK_HOST(num_consumers + 1 <= data.num_push_blocks);
CHECK_HOST(num_consumers + 1 <= push.num_blocks);
// use last block to clean up the counter
const auto num_consumer_blocks = num_consumers + 1;
using gemm_ag::spin_add3_kernel;
python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh
+84 -84
View File
@@ -17,7 +17,6 @@
#include <cute/arch/cluster_sm90.hpp>
#include <cutlass/cuda_host_adapter.hpp>
#include "ptx_sys.cuh"
#include <cstdint>
#include <cstdio>
#include <cstdlib>
@@ -34,7 +33,7 @@
namespace sglang {
namespace ptx {
namespace device::ptx {
// ---- generic → shared address conversion (PTX ISA §10.4) --------------------
@@ -343,7 +342,7 @@ static SGL_DEVICE void cp_async_bulk_tensor_2d_load_multicast_cg1(
: "memory");
}
} // namespace ptx
} // namespace device::ptx
namespace swz {
@@ -465,11 +464,6 @@ __device__ __forceinline__ int2 group_n_swizzle(int linear, int crank, int clust
namespace oproj_ar {
using device::distributed::atomic_add_acq_rel_gpu;
using device::distributed::fence_release_sys;
using device::distributed::load_acquire_sys;
using device::distributed::red_add_relaxed_sys;
// ---------------------------------------------------------------- constants
#ifndef OPROJ_N // output dim (columns of W / of out). The
#define OPROJ_N 7168 // default is the Kimi-K3 o_proj shape; other
@@ -639,19 +633,19 @@ __global__ void __launch_bounds__(kThreads) oproj_ar_kernel(
uint64_t* fullb = reinterpret_cast<uint64_t*>(a_st + size_t(S) * kStA);
uint64_t* emptyb = fullb + S;
const uint32_t crank = C > 1 ? ptx::cluster_cta_rank() : 0;
const uint32_t crank = C > 1 ? device::ptx::cluster_cta_rank() : 0;
{
if (tid == 0) {
ptx::prefetch_tensormap(&w_map);
ptx::prefetch_tensormap(&x_map);
device::ptx::prefetch_tensormap(&w_map);
device::ptx::prefetch_tensormap(&x_map);
#pragma unroll
for (int s = 0; s < S; ++s) {
ptx::mbar_init(fullb + s, 1);
ptx::mbar_init(emptyb + s, crank == 0 ? kCWarps * C : kCWarps);
device::ptx::mbar_init(fullb + s, 1);
device::ptx::mbar_init(emptyb + s, crank == 0 ? kCWarps * C : kCWarps);
}
}
__syncthreads();
if constexpr (C > 1) ptx::cluster_sync_rel_acq();
if constexpr (C > 1) device::ptx::cluster_sync_rel_acq();
if (warp == kCWarps) {
// ---- producer: the whole K stream, one thread -----------------
@@ -663,28 +657,28 @@ __global__ void __launch_bounds__(kThreads) oproj_ar_kernel(
for (int j = 0; j < KSTEPS; ++j) {
const int slot = j % S;
const int jj = (j + phase) % KSTEPS;
if (j >= S) ptx::mbar_wait_parity(emptyb + slot, ((j - S) / S) & 1);
ptx::mbar_arrive_expect_tx(fullb + slot, kStB + kStA);
if (j >= S) device::ptx::mbar_wait_parity(emptyb + slot, ((j - S) / S) & 1);
device::ptx::mbar_arrive_expect_tx(fullb + slot, kStB + kStA);
#pragma unroll
for (int c = 0; c < CH; ++c) {
ptx::cp_async_bulk_tensor_2d_load(
ptx::to_shared(b_st + size_t(slot) * kStB + c * kBBytes),
device::ptx::cp_async_bulk_tensor_2d_load(
device::ptx::to_shared(b_st + size_t(slot) * kStB + c * kBBytes),
&w_map,
(jj * CH + c) * kBK,
strip.t0 * 8,
fullb + slot);
if constexpr (C > 1) {
if (crank == 0)
ptx::cp_async_bulk_tensor_2d_load_multicast_cg1(
ptx::to_shared(a_st + size_t(slot) * kStA + c * kABytes),
device::ptx::cp_async_bulk_tensor_2d_load_multicast_cg1(
device::ptx::to_shared(a_st + size_t(slot) * kStA + c * kABytes),
&x_map,
(jj * CH + c) * kBK,
0,
fullb + slot,
uint16_t((1u << C) - 1));
} else {
ptx::cp_async_bulk_tensor_2d_load(
ptx::to_shared(a_st + size_t(slot) * kStA + c * kABytes),
device::ptx::cp_async_bulk_tensor_2d_load(
device::ptx::to_shared(a_st + size_t(slot) * kStA + c * kABytes),
&x_map,
(jj * CH + c) * kBK,
0,
@@ -702,36 +696,36 @@ __global__ void __launch_bounds__(kThreads) oproj_ar_kernel(
const int a_row = lane & 15, a_ka = lane >> 4;
for (int s = 0; s < KSTEPS; ++s) {
const int slot = s % S;
ptx::mbar_wait_parity(fullb + slot, (s / S) & 1);
device::ptx::mbar_wait_parity(fullb + slot, (s / S) & 1);
#pragma unroll
for (int c = 0; c < CH; ++c) {
const uint32_t b_base = ptx::to_shared(b_st + size_t(slot) * kStB + c * kBBytes);
const uint32_t b_base = device::ptx::to_shared(b_st + size_t(slot) * kStB + c * kBBytes);
#pragma unroll
for (int k16 = 0; k16 < kBK / 16; ++k16) {
uint32_t b0, b1;
ptx::ldmatrix_x2_b16(
device::ptx::ldmatrix_x2_b16(
b_base + uint32_t(b_row) * (kBK * 2) + swz::smem_col_128b_bf16(b_row, (k16 * 2 + b_ka) * 8) * 2,
b0,
b1);
#pragma unroll
for (int mt = 0; mt < MT; ++mt) {
uint32_t a0, a1, a2, a3;
ptx::ldmatrix_x4_b16(
ptx::to_shared(
device::ptx::ldmatrix_x4_b16(
device::ptx::to_shared(
a_st + size_t(slot) * kStA + c * kABytes + uint32_t(mt * 16 + a_row) * (kBK * 2) +
swz::smem_col_128b_bf16(a_row, (k16 * 2 + a_ka) * 8) * 2),
a0,
a1,
a2,
a3);
ptx::mma_m16n8k16_bf16f32(acc[mt], a0, a1, a2, a3, b0, b1);
device::ptx::mma_m16n8k16_bf16f32(acc[mt], a0, a1, a2, a3, b0, b1);
}
}
}
if (lane == 0) {
ptx::mbar_arrive(emptyb + slot);
device::ptx::mbar_arrive(emptyb + slot);
if constexpr (C > 1)
if (crank != 0) ptx::mbar_arrive_cluster_release(emptyb + slot, 0);
if (crank != 0) device::ptx::mbar_arrive_cluster_release(emptyb + slot, 0);
}
}
}
@@ -751,7 +745,7 @@ __global__ void __launch_bounds__(kThreads) oproj_ar_kernel(
if (tid == 0 && epoch >= 2) {
const uint32_t e2 = epoch - 2;
const uint32_t tgt = (e2 / kRing + 1) * R;
while (load_acquire_sys(done_local + size_t(e2 % kRing) * kMaxCTA) < tgt) {
while (device::ptx::load_acquire_sys(done_local + size_t(e2 % kRing) * kMaxCTA) < tgt) {
}
}
__syncthreads();
@@ -792,19 +786,19 @@ __global__ void __launch_bounds__(kThreads) oproj_ar_kernel(
// ---- boundary (gather + flag + per-CTA local-replica spin) -------------
__syncthreads();
if (tid == 0) {
const uint32_t old = atomic_add_acq_rel_gpu(gather_fam + ring, 1);
const uint32_t old = device::ptx::atomic_add_acq_rel_gpu(gather_fam + ring, 1);
if (old + 1 == gath_target) {
// ONE completing-CTA fence: publishes every CTA's pushes via the
// acq_rel gather chain. Measured 0.9 us better than per-CTA
// fences at bs1 (11.04 vs 11.97) — the parallel-drain theory lost.
fence_release_sys();
device::ptx::fence_release_sys();
{
#pragma unroll
for (int r = 0; r < R; ++r)
red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + foff) + ring, 1);
device::ptx::red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + foff) + ring, 1);
}
}
while (load_acquire_sys(flag_local) < flag_target) {
while (device::ptx::load_acquire_sys(flag_local) < flag_target) {
}
}
__syncthreads();
@@ -844,7 +838,7 @@ __global__ void __launch_bounds__(kThreads) oproj_ar_kernel(
{
#pragma unroll
for (int r = 0; r < R; ++r)
red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + doff2) + off, 1);
device::ptx::red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + doff2) + off, 1);
}
}
}
@@ -913,7 +907,7 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
: 2 * kD3BN <= 256 ? 256
: 512;
extern __shared__ __align__(1024) uint8_t smem_buf[];
const uint32_t smem_base = ptx::to_shared(smem_buf);
const uint32_t smem_base = device::ptx::to_shared(smem_buf);
constexpr uint32_t kSmemAOff = 0, kSmemBOff = kD3NS * kD3ABytes;
__shared__ __align__(8) uint64_t tma_mbars[12];
@@ -940,23 +934,29 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
const uint32_t gath_target = wrap * gridDim.x;
const size_t sb = slot_off(parity, prm.my_rank, R);
if (warp_id == 0 && ptx::elect_one()) {
if (warp_id == 0 && device::ptx::elect_one()) {
for (int i = 0; i < kD3NS; ++i) {
ptx::mbar_init(&tma_mbars[i], 1);
ptx::mbar_init(&mma_mbars[i], 1);
device::ptx::mbar_init(&tma_mbars[i], 1);
device::ptx::mbar_init(&mma_mbars[i], 1);
}
for (int i = 0; i < 2; ++i) {
ptx::mbar_init(&mainloop_mbars[i], 1);
ptx::mbar_init(&epi_mbars[i], 4 * 32);
device::ptx::mbar_init(&mainloop_mbars[i], 1);
device::ptx::mbar_init(&epi_mbars[i], 4 * 32);
}
} else if (warp_id == 1) {
ptx::tcgen05_alloc(ptx::to_shared(s_taddr), kTmemCols);
device::ptx::tcgen05_alloc(device::ptx::to_shared(s_taddr), kTmemCols);
}
__syncthreads();
const uint32_t taddr = s_taddr[0];
constexpr uint32_t i_desc = ptx::mma_inst_desc_f16(
kD3BM, kD3BN, ptx::F16Type::BF16, ptx::F16Type::BF16, ptx::DType::F32, ptx::Major::K, ptx::Major::K);
constexpr uint32_t i_desc = device::ptx::mma_inst_desc_f16(
kD3BM,
kD3BN,
device::ptx::F16Type::BF16,
device::ptx::F16Type::BF16,
device::ptx::DType::F32,
device::ptx::Major::K,
device::ptx::Major::K);
auto tile_mn = [&](int linear) -> int2 {
return dense_gemm_mainloop::group_n_swizzle<1, kD3GroupN>(linear, 0, kGridM, kGridN);
};
@@ -964,13 +964,13 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
// Prefetch the first ring of input-independent weight stages before the
// PDL dependency. SWAP changes which operand slot contains W, but never
// changes which tensor map is safe to touch here.
if (warp_id == 0 && ptx::elect_one()) {
if (warp_id == 0 && device::ptx::elect_one()) {
constexpr int kPrefetch = kIters < kD3NS ? kIters : kD3NS;
const int2 mn = tile_mn(bid);
#pragma unroll
for (int k = 0; k < kPrefetch; ++k) {
ptx::mbar_arrive_expect_tx(&tma_mbars[k], kD3ABytes + kD3BBytes);
ptx::cp_async_bulk_tensor_2d_load(
device::ptx::mbar_arrive_expect_tx(&tma_mbars[k], kD3ABytes + kD3BBytes);
device::ptx::cp_async_bulk_tensor_2d_load(
smem_base + (SWAP ? kSmemAOff + k * kD3ABytes : kSmemBOff + k * kD3BBytes),
&w_tmap,
k * kD3BK,
@@ -987,26 +987,26 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
if (tid == 0 && epoch >= 2) {
const uint32_t e2 = epoch - 2;
const uint32_t tgt = (e2 / kRing + 1) * R;
while (load_acquire_sys(done_local + size_t(e2 % kRing) * kMaxCTA) < tgt) {
while (device::ptx::load_acquire_sys(done_local + size_t(e2 % kRing) * kMaxCTA) < tgt) {
}
}
__syncthreads();
}
if (warp_id == 0 && ptx::elect_one()) {
if (warp_id == 0 && device::ptx::elect_one()) {
// TMA issuer (simple persistent) — verbatim dense_1cta fp8out shape
int stage = 0, mma_phase = 1;
for (int t = bid; t < kTiles; t += num_bids) {
const int2 mn = tile_mn(t);
for (int k = 0; k < kIters; ++k) {
ptx::mbar_wait_parity(&mma_mbars[stage], mma_phase);
device::ptx::mbar_wait_parity(&mma_mbars[stage], mma_phase);
constexpr bool kDropA = false;
const bool prefetched = (t == bid && k < kD3NS);
if (!prefetched) ptx::mbar_arrive_expect_tx(&tma_mbars[stage], (kDropA ? 0 : kD3ABytes) + kD3BBytes);
if (!prefetched) device::ptx::mbar_arrive_expect_tx(&tma_mbars[stage], (kDropA ? 0 : kD3ABytes) + kD3BBytes);
if constexpr (!kDropA) {
// In SWAP, A is the prefetched weight; otherwise A is x.
if (!prefetched || !SWAP)
ptx::cp_async_bulk_tensor_2d_load(
device::ptx::cp_async_bulk_tensor_2d_load(
smem_base + kSmemAOff + stage * kD3ABytes,
SWAP ? &w_tmap : &x_tmap,
k * kD3BK,
@@ -1015,7 +1015,7 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
}
// In SWAP, B is x; otherwise B is the prefetched weight.
if (!prefetched || SWAP)
ptx::cp_async_bulk_tensor_2d_load(
device::ptx::cp_async_bulk_tensor_2d_load(
smem_base + kSmemBOff + stage * kD3BBytes,
SWAP ? &x_tmap : &w_tmap,
k * kD3BK,
@@ -1027,30 +1027,30 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
}
}
}
} else if (warp_id == 1 && ptx::elect_one()) {
} else if (warp_id == 1 && device::ptx::elect_one()) {
// MMA issuer with 2-stage TMEM ping-pong
int stage = 0, tma_phase = 0, ml_stage = 0, epi_phase = 1;
for (int t = bid; t < kTiles; t += num_bids) {
ptx::mbar_wait_parity(&epi_mbars[ml_stage], epi_phase);
device::ptx::mbar_wait_parity(&epi_mbars[ml_stage], epi_phase);
const uint32_t tmem_d = taddr + uint32_t(ml_stage) * kD3BN;
for (int k = 0; k < kIters; ++k) {
ptx::mbar_wait_parity(&tma_mbars[stage], tma_phase);
ptx::tcgen05_fence_after_thread_sync();
device::ptx::mbar_wait_parity(&tma_mbars[stage], tma_phase);
device::ptx::tcgen05_fence_after_thread_sync();
const uint32_t a_smem = smem_base + kSmemAOff + stage * kD3ABytes;
const uint32_t b_smem = smem_base + kSmemBOff + stage * kD3BBytes;
#pragma unroll
for (int k2 = 0; k2 < kD3KPer; ++k2) {
const uint64_t da = ptx::mma_smem_desc_k_major<uint16_t, kD3BK, 128>(a_smem + uint32_t(k2) * 32);
const uint64_t db = ptx::mma_smem_desc_k_major<uint16_t, kD3BK, 128>(b_smem + uint32_t(k2) * 32);
ptx::tcgen05_mma_f16(tmem_d, da, db, i_desc, (k == 0 && k2 == 0) ? 0u : 1u);
const uint64_t da = device::ptx::mma_smem_desc_k_major<uint16_t, kD3BK, 128>(a_smem + uint32_t(k2) * 32);
const uint64_t db = device::ptx::mma_smem_desc_k_major<uint16_t, kD3BK, 128>(b_smem + uint32_t(k2) * 32);
device::ptx::tcgen05_mma_f16(tmem_d, da, db, i_desc, (k == 0 && k2 == 0) ? 0u : 1u);
}
ptx::tcgen05_commit_arrive(&mma_mbars[stage]);
device::ptx::tcgen05_commit_arrive(&mma_mbars[stage]);
if (++stage == kD3NS) {
stage = 0;
tma_phase ^= 1;
}
}
ptx::tcgen05_commit_arrive(&mainloop_mbars[ml_stage]);
device::ptx::tcgen05_commit_arrive(&mainloop_mbars[ml_stage]);
ml_stage ^= 1;
if (ml_stage == 0) epi_phase ^= 1;
}
@@ -1061,21 +1061,21 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
int ml_stage = 0, ml_phase = 0;
for (int t = bid; t < kTiles; t += num_bids) {
const int2 mn = tile_mn(t);
ptx::mbar_wait_parity(&mainloop_mbars[ml_stage], ml_phase);
ptx::tcgen05_fence_after_thread_sync();
device::ptx::mbar_wait_parity(&mainloop_mbars[ml_stage], ml_phase);
device::ptx::tcgen05_fence_after_thread_sync();
const uint32_t tmem_d_base = taddr + uint32_t(ml_stage) * kD3BN;
const int row = mn.x * kD3BM + epi_warp * 32 + lane_id; // SWAP: n
#pragma unroll 4
for (int nb = 0; nb < kNBlk; ++nb) {
const uint32_t taddr_n = tmem_d_base + uint32_t(nb) * 8 + taddr_lane;
uint32_t r0, r1, r2, r3, r4, r5, r6, r7;
ptx::tcgen05_ld_32x32b_x8(taddr_n, r0, r1, r2, r3, r4, r5, r6, r7);
ptx::tcgen05_wait_ld();
device::ptx::tcgen05_ld_32x32b_x8(taddr_n, r0, r1, r2, r3, r4, r5, r6, r7);
device::ptx::tcgen05_wait_ld();
uint4 v;
v.x = ptx::cvt_pack_f32x2_to<ptx::bf16>(__int_as_float(r1), __int_as_float(r0));
v.y = ptx::cvt_pack_f32x2_to<ptx::bf16>(__int_as_float(r3), __int_as_float(r2));
v.z = ptx::cvt_pack_f32x2_to<ptx::bf16>(__int_as_float(r5), __int_as_float(r4));
v.w = ptx::cvt_pack_f32x2_to<ptx::bf16>(__int_as_float(r7), __int_as_float(r6));
v.x = device::ptx::cvt_pack_f32x2_to<device::ptx::bf16>(__int_as_float(r1), __int_as_float(r0));
v.y = device::ptx::cvt_pack_f32x2_to<device::ptx::bf16>(__int_as_float(r3), __int_as_float(r2));
v.z = device::ptx::cvt_pack_f32x2_to<device::ptx::bf16>(__int_as_float(r5), __int_as_float(r4));
v.w = device::ptx::cvt_pack_f32x2_to<device::ptx::bf16>(__int_as_float(r7), __int_as_float(r6));
if (SWAP || row < M) { // SWAP masks pad m-cols below
{
const size_t off = sb + d3_slot_off(t, nb, epi_warp, lane_id, kNBlk);
@@ -1093,7 +1093,7 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
}
}
}
(void)ptx::mbar_arrive(&epi_mbars[ml_stage]);
(void)device::ptx::mbar_arrive(&epi_mbars[ml_stage]);
ml_stage ^= 1;
if (ml_stage == 0) ml_phase ^= 1;
}
@@ -1101,22 +1101,22 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
__syncthreads();
device::PDLTriggerSecondary<true>();
if (warp_id == 1) {
ptx::tcgen05_dealloc(taddr, kTmemCols);
ptx::tcgen05_relinquish();
device::ptx::tcgen05_dealloc(taddr, kTmemCols);
device::ptx::tcgen05_relinquish();
}
// ---- boundary (fam rings) — verbatim member-1 contract ----------------
if (tid == 0) {
const uint32_t old = atomic_add_acq_rel_gpu(gather_fam + ring, 1);
const uint32_t old = device::ptx::atomic_add_acq_rel_gpu(gather_fam + ring, 1);
if (old + 1 == gath_target) {
fence_release_sys();
device::ptx::fence_release_sys();
{
#pragma unroll
for (int r = 0; r < R; ++r)
red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + foff) + ring, 1);
device::ptx::red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + foff) + ring, 1);
}
}
while (load_acquire_sys(flag_local) < flag_target) {
while (device::ptx::load_acquire_sys(flag_local) < flag_target) {
}
}
__syncthreads();
@@ -1156,16 +1156,16 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
__syncthreads();
if (tid == 0) {
uint32_t* const g2 = gather_fam + kRing; // AG gather ring
const uint32_t old = atomic_add_acq_rel_gpu(g2 + ring, 1);
const uint32_t old = device::ptx::atomic_add_acq_rel_gpu(g2 + ring, 1);
if (old + 1 == gath_target) {
fence_release_sys();
device::ptx::fence_release_sys();
{
#pragma unroll
for (int r = 0; r < R; ++r)
red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + foff + 256) + ring, 1);
device::ptx::red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + foff + 256) + ring, 1);
}
}
while (load_acquire_sys(reinterpret_cast<uint32_t*>(prm.uc_base[prm.my_rank] + foff + 256) + ring) <
while (device::ptx::load_acquire_sys(reinterpret_cast<uint32_t*>(prm.uc_base[prm.my_rank] + foff + 256) + ring) <
flag_target) {
}
}
@@ -1200,7 +1200,7 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
{
#pragma unroll
for (int r = 0; r < R; ++r)
red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + doff2) + off, 1);
device::ptx::red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + doff2) + off, 1);
}
}
return;
@@ -1254,7 +1254,7 @@ __global__ void __launch_bounds__(kD3Threads) oproj_dense_ar_kernel(
{
#pragma unroll
for (int r = 0; r < R; ++r)
red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + doff2) + off, 1);
device::ptx::red_add_relaxed_sys(reinterpret_cast<uint32_t*>(prm.uc_base[r] + doff2) + off, 1);
}
}
}
python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh
-72
View File
@@ -1,72 +0,0 @@
#pragma once
// System-scope PTX the Kimi K3 collectives need and no shared sglang header
// wraps. These live beside their only consumers (gemm_ar / gemm_ag) rather than
// in distributed/communicator.cuh: that header's Semaphore does not use any of
// them, so putting them there would grow a shared header for one caller's
// benefit.
//
// The `device::distributed` namespace is deliberate -- it is where the rest of
// the collective vocabulary lives, so call sites and `using` declarations read
// the same whichever header supplied the symbol.
#include <sgl_kernel/utils.cuh>
#include <cstdint>
namespace sglang {
namespace device::distributed {
// Peer-visible flag increment. `.sys` scope, relaxed: ordering is established by
// the surrounding fence, not by this store.
SGL_DEVICE void red_add_relaxed_sys(uint32_t* ptr, uint32_t val) {
asm volatile("red.relaxed.sys.global.add.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
}
// Acquire load of a peer-written flag: everything the writer released before
// its matching store is visible to this thread afterwards.
SGL_DEVICE uint32_t load_acquire_sys(const uint32_t* ptr) {
uint32_t val;
asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(ptr) : "memory");
return val;
}
// Publishes every prior write to system scope. Pair with a relaxed flag store so
// a peer's acquire load of that flag also observes the payload.
SGL_DEVICE void fence_release_sys() {
asm volatile("fence.release.sys;" ::: "memory");
}
// Device-scope arrival counter. acq_rel so the winner of the count also observes
// the losers' payload writes.
SGL_DEVICE uint32_t atomic_add_acq_rel_gpu(uint32_t* ptr, uint32_t val) {
uint32_t old;
asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(ptr), "r"(val) : "memory");
return old;
}
// One store fanned out to every rank in the multicast team.
SGL_DEVICE void multimem_store_relaxed(uint32_t* ptr, uint32_t val) {
asm volatile("multimem.st.relaxed.sys.global.b32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
}
SGL_DEVICE void multimem_red_add_relaxed(uint32_t* mc_flag) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(mc_flag) : "memory");
#else
assert(false && "multimem red is only supported on Hopper or later architecture");
#endif
}
SGL_DEVICE void multimem_red_add_release(uint32_t* mc_flag) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(mc_flag) : "memory");
#else
assert(false && "multimem red is only supported on Hopper or later architecture");
#endif
}
} // namespace device::distributed
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh
+104 -143
View File
@@ -15,10 +15,15 @@
// phase counters. A bumper block advances counters outside the tuned work
// grid, so calls remain protocol-compatible with all-reduce and gemm_ag.
#include <sgl_kernel/tensor.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include "../../distributed/custom_all_reduce.cuh"
#include <sgl_kernel/distributed/communicator.cuh>
#include <sgl_kernel/distributed/ptx.cuh>
#include <tvm/ffi/extra/stl.h>
namespace sglang {
@@ -35,7 +40,7 @@ struct Params {
uint8_t* push_ws_mc;
Counter* counter;
Semaphore* sem_local;
uint8_t* sem_mc;
Semaphore* sem_mc;
uint8_t* input_mc;
uint8_t* output_mc;
int64_t stride_bytes;
@@ -45,48 +50,15 @@ struct Params {
uint32_t residual_is_local;
};
template <typename Vec>
SGL_DEVICE void make_nonzero(Vec& vec) {
constexpr uint32_t kNegZeroPair = 0x8000u;
auto& bits = *reinterpret_cast<uint4*>(&vec);
if (bits.x == 0) bits.x = kNegZeroPair;
if (bits.y == 0) bits.y = kNegZeroPair;
if (bits.z == 0) bits.z = kNegZeroPair;
if (bits.w == 0) bits.w = kNegZeroPair;
}
SGL_DEVICE uint32_t* sem_mc_flag(uint8_t* sem_mc, uint32_t block) {
static_assert(sizeof(Semaphore) == 128);
return reinterpret_cast<uint32_t*>(sem_mc + block * sizeof(Semaphore));
}
SGL_DEVICE void sem_arrive_relaxed(uint32_t* flag) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
#else
assert(false && "multimem red requires Hopper or later");
#endif
}
SGL_DEVICE void sem_arrive_release(uint32_t* flag) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.release.sys.global.add.u32 [%0], 1;" ::"l"(flag) : "memory");
#else
assert(false && "multimem red requires Hopper or later");
#endif
}
/// The 16 B staging vector, viewed as the 4 u32 words the lamport marker
/// protocol tests. See LamportTrait in distributed/communicator.cuh.
using Lamport = device::distributed::LamportTrait<bf16_t, 8, /*kAtom=*/4>;
template <typename Vec>
SGL_DEVICE bool has_empty_marker(const Vec& vec) {
const auto bits = *reinterpret_cast<const uint4*>(&vec);
return bits.x == 0 || bits.y == 0 || bits.z == 0 || bits.w == 0;
}
template <typename Vec>
SGL_DEVICE Vec zero_vec() {
Vec zero;
zero.fill(bf16x2_t{get_pos_zero<bf16_t>(), get_pos_zero<bf16_t>()});
return zero;
SGL_DEVICE Vec empty_vec() {
Vec vec;
Lamport::fill_pos_zero(vec.data());
return vec;
}
template <uint32_t kWorldSize>
@@ -123,9 +95,9 @@ __global__ void reduce_scatter_res_kernel(const __grid_constant__ Params params)
const uint32_t dst_rank = vid / params.local_vecs;
const uint32_t local_vid = vid - dst_rank * params.local_vecs;
vec_t vec;
ld_global_16B(vec, params.input, vid);
make_nonzero(vec);
st_relaxed_16B(vec, params.push_workspaces[dst_rank] + producer_offset, local_vid);
device::ptx::ld_global_16B(vec, params.input, vid);
Lamport::clear_pos_zero(vec.data());
device::ptx::st_relaxed_16B(vec, params.push_workspaces[dst_rank] + producer_offset, local_vid);
}
device::PDLTriggerSecondary<kUsePDL>();
@@ -133,26 +105,26 @@ __global__ void reduce_scatter_res_kernel(const __grid_constant__ Params params)
// Poll this rank's shard from all producer slots and reduce locally.
const auto poll_base = params.push_workspaces[params.rank] + phase_offset;
const auto residual_base = params.residual + (params.residual_is_local ? 0 : params.rank * params.local_vecs * 16);
const auto zero = zero_vec<vec_t>();
const auto zero = empty_vec<vec_t>();
for (uint32_t vid = tid; vid < params.local_vecs; vid += num_threads) {
vec_t vec[kWorldSize + kHasResidual];
if constexpr (kHasResidual) {
ld_global_16B(vec[kWorldSize], residual_base, vid);
device::ptx::ld_global_16B(vec[kWorldSize], residual_base, vid);
}
do {
bool empty = false;
#pragma unroll
for (uint32_t rank = 0; rank < kWorldSize; ++rank) {
ld_relaxed_16B(vec[rank], poll_base + rank * params.stride_bytes, vid);
empty |= has_empty_marker(vec[rank]);
device::ptx::ld_relaxed_16B(vec[rank], poll_base + rank * params.stride_bytes, vid);
empty |= Lamport::has_pos_zero(vec[rank].data());
}
if (!empty) break;
} while (true);
const auto out = reduce(vec);
st_global_16B(out, params.output, vid);
const auto out = device::reduce_vec(vec);
device::ptx::st_global_16B(out, params.output, vid);
#pragma unroll
for (uint32_t rank = 0; rank < kWorldSize; ++rank) {
st_global_16B(zero, poll_base + rank * params.stride_bytes, vid);
device::ptx::st_global_16B(zero, poll_base + rank * params.stride_bytes, vid);
}
}
@@ -180,25 +152,25 @@ __global__ void all_gather_kernel(const __grid_constant__ Params params) {
// One multicast store places this rank's shard in the same slot on peers.
for (uint32_t vid = tid; vid < params.local_vecs; vid += num_threads) {
vec_t vec;
ld_global_16B(vec, params.input, vid);
make_nonzero(vec);
st_multimem_16B(vec, params.push_ws_mc + producer_offset, vid);
device::ptx::ld_global_16B(vec, params.input, vid);
Lamport::clear_pos_zero(vec.data());
device::ptx::st_multimem_16B(vec, params.push_ws_mc + producer_offset, vid);
}
device::PDLTriggerSecondary<kUsePDL>();
const auto poll_base = params.push_workspaces[params.rank] + phase_offset;
const auto zero = zero_vec<vec_t>();
const auto zero = empty_vec<vec_t>();
for (uint32_t vid = tid; vid < kWorldSize * params.local_vecs; vid += num_threads) {
const uint32_t src_rank = vid / params.local_vecs;
const uint32_t local_vid = vid - src_rank * params.local_vecs;
const auto src = poll_base + src_rank * params.stride_bytes;
vec_t vec;
do {
ld_relaxed_16B(vec, src, local_vid);
} while (has_empty_marker(vec));
st_global_16B(vec, params.output, vid);
st_global_16B(zero, src, local_vid);
device::ptx::ld_relaxed_16B(vec, src, local_vid);
} while (Lamport::has_pos_zero(vec.data()));
device::ptx::st_global_16B(vec, params.output, vid);
device::ptx::st_global_16B(zero, src, local_vid);
}
__syncthreads();
@@ -214,16 +186,10 @@ template <uint32_t kWorldSize, bool kUsePDL>
__global__ void all_gather_direct_kernel(const __grid_constant__ Params params) {
using vec_t = device::AlignedVector<bf16x2_t, 4>; // 16 B
uint32_t exit_base = 0;
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize);
exit_base = reserved + kWorldSize;
device::PDLWaitPrimary<kUsePDL>();
sem_arrive_relaxed(sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_relaxed() - reserved < kWorldSize)
;
}
// Reserve the window before the PDL wait, signal after it.
const auto barrier = device::distributed::McBarrier(params.sem_local, params.sem_mc, kWorldSize, /*num_arrives=*/2);
device::PDLWaitPrimary<kUsePDL>();
barrier.arrive_relaxed(/*n=*/0);
__syncthreads();
const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -231,18 +197,13 @@ __global__ void all_gather_direct_kernel(const __grid_constant__ Params params)
const uint32_t dst_bias = params.rank * params.local_vecs;
for (uint32_t vid = tid; vid < params.local_vecs; vid += step) {
vec_t vec;
ld_global_16B(vec, params.input, vid);
st_multimem_16B(vec, params.output_mc, dst_bias + vid);
device::ptx::ld_global_16B(vec, params.input, vid);
device::ptx::st_multimem_16B(vec, params.output_mc, dst_bias + vid);
}
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
sem_arrive_release(sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_acquire() - exit_base < kWorldSize)
;
}
barrier.arrive_rel_acq(/*n=*/1);
}
// NVLS pull variant: o_proj writes its TP-partial result into multicast-bound
@@ -253,16 +214,10 @@ __global__ void reduce_scatter_pull_kernel(const __grid_constant__ Params params
using vec_t = device::AlignedVector<bf16x2_t, 4>; // 16 B
using SumOp = device::ReductionTrait<device::ReductionOp::SUM, bf16x2_t>;
uint32_t exit_base = 0;
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
const auto reserved = semaphore->counter_ptr()->inc(2 * kWorldSize);
exit_base = reserved + kWorldSize;
device::PDLWaitPrimary<kUsePDL>();
sem_arrive_relaxed(sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_relaxed() - reserved < kWorldSize)
;
}
// Reserve the window before the PDL wait, signal after it.
const auto barrier = device::distributed::McBarrier(params.sem_local, params.sem_mc, kWorldSize, /*num_arrives=*/2);
device::PDLWaitPrimary<kUsePDL>();
barrier.arrive_relaxed(/*n=*/0);
__syncthreads();
const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -272,26 +227,21 @@ __global__ void reduce_scatter_pull_kernel(const __grid_constant__ Params params
kHasResidual ? params.residual + (params.residual_is_local ? 0 : params.rank * params.local_vecs * 16) : nullptr;
for (uint32_t vid = tid; vid < params.local_vecs; vid += step) {
vec_t vec;
ld_multimem_16B(vec, input_mc, vid);
device::ptx::ld_multimem_16B(vec, input_mc, vid);
if constexpr (kHasResidual) {
vec_t res;
ld_global_16B(res, residual, vid);
device::ptx::ld_global_16B(res, residual, vid);
#pragma unroll
for (uint32_t j = 0; j < 4; ++j) {
vec[j] = SumOp::reduce(vec[j], res[j]);
}
}
st_global_16B(vec, params.output, vid);
device::ptx::st_global_16B(vec, params.output, vid);
}
device::PDLTriggerSecondary<kUsePDL>();
__syncthreads();
if (threadIdx.x == 0) {
auto* semaphore = &params.sem_local[blockIdx.x];
sem_arrive_release(sem_mc_flag(params.sem_mc, blockIdx.x));
while (semaphore->get_acquire() - exit_base < kWorldSize)
;
}
barrier.arrive_rel_acq(/*n=*/1);
}
} // namespace sp_collective
@@ -301,13 +251,16 @@ template <uint32_t kWorldSize, bool kUsePDL>
struct SPCollectiveKernel {
using TensorView = tvm::ffi::TensorView;
/// Tensor validation shared by every variant. Plane pointers are bound
/// separately: the push variants stage through the push plane and never
/// touch a semaphore, the pull variants barrier on the pull plane's
/// semaphores and never touch the push workspace.
static sp_collective::Params make_params(
const host::distributed::CommunicatorObj& data,
const host::distributed::CommunicatorObj& comm,
TensorView input,
TensorView output,
std::optional<TensorView> residual,
bool residual_is_local,
int64_t ws_mc_base) {
bool residual_is_local) {
using namespace host;
auto input_elems = SymbolicSize{"input_elems"};
auto local_elems = SymbolicSize{"local_elems"};
@@ -322,42 +275,59 @@ struct SPCollectiveKernel {
TensorMatcher({input_elems}).with_dtype<bf16_t>().with_device(device).verify(residual.value());
}
}
CHECK_HOST(data.world_size == kWorldSize);
CHECK_HOST(comm.get_world_size() == kWorldSize);
CHECK_HOST(local_elems.unwrap() > 0);
CHECK_HOST(input_elems.unwrap() == local_elems.unwrap() * kWorldSize);
CHECK_HOST(local_elems.unwrap() % 8 == 0) << "local shard bytes must be 16B aligned";
CHECK_HOST(local_elems.unwrap() * sizeof(bf16_t) <= data.push_bytes) << "local shard exceeds a push slot";
sp_collective::Params params{
return sp_collective::Params{
.input = static_cast<const uint8_t*>(input.data_ptr()),
.output = static_cast<uint8_t*>(output.data_ptr()),
.residual = residual.has_value() ? static_cast<const uint8_t*>(residual.value().data_ptr()) : nullptr,
.push_workspaces = {},
.push_ws_mc = reinterpret_cast<uint8_t*>(ws_mc_base),
.counter = data.push_counter,
.sem_local = data.pull_semaphores[data.rank],
.push_ws_mc = nullptr,
.counter = nullptr,
.sem_local = nullptr,
.sem_mc = nullptr,
.input_mc = nullptr,
.output_mc = nullptr,
.stride_bytes = data.push_bytes,
.num_counters = data.num_push_blocks,
.rank = data.rank,
.stride_bytes = 0,
.num_counters = 0,
.rank = comm.get_rank(),
.local_vecs = static_cast<uint32_t>(local_elems.unwrap() * sizeof(bf16_t) / 16),
.residual_is_local = static_cast<uint32_t>(residual_is_local),
};
for (uint32_t i = 0; i < kWorldSize; ++i) {
params.push_workspaces[i] = data.push_workspaces[i];
}
return params;
}
static void check_launch(const host::distributed::CommunicatorObj& data, int64_t num_blocks, int64_t block_size) {
CHECK_HOST(num_blocks > 0 && num_blocks < data.num_push_blocks);
static void bind_push(sp_collective::Params& params, const host::distributed::PushPlaneObj& push) {
CHECK_HOST(int64_t(params.local_vecs) * 16 <= push.slot_bytes) << "local shard exceeds a push slot";
for (uint32_t i = 0; i < kWorldSize; ++i) {
params.push_workspaces[i] = push.workspaces[i];
}
params.push_ws_mc = push.mc_workspace;
params.counter = push.counter;
params.stride_bytes = push.slot_bytes;
params.num_counters = push.num_blocks;
}
static void bind_pull(sp_collective::Params& params, const host::distributed::PullPlaneObj& pull) {
CHECK_HOST(pull.mc_semaphore != nullptr) << "the pull path needs a multicast-capable pull plane";
params.sem_local = pull.semaphores[pull.rank];
params.sem_mc = pull.mc_semaphore;
}
static void check_push_launch(const host::distributed::PushPlaneObj& push, int64_t num_blocks, int64_t block_size) {
CHECK_HOST(num_blocks > 0 && num_blocks < push.num_blocks);
// The RS reduction keeps one 16B vector per producer in registers.
// 1024-thread CTAs exceed the GB300 launch resource limit.
CHECK_HOST(block_size >= 32 && block_size <= 512 && block_size % 32 == 0);
}
static void check_pull_launch(uint32_t max_blocks, int64_t num_blocks, int64_t block_size) {
CHECK_HOST(num_blocks > 0 && num_blocks <= max_blocks);
CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0);
}
static void reduce_scatter_res(
CommunicatorRef ref,
TensorView input,
@@ -366,26 +336,23 @@ struct SPCollectiveKernel {
bool residual_is_local,
int64_t num_blocks,
int64_t block_size) {
const auto& data = *ref.get();
check_launch(data, num_blocks, block_size);
auto params = make_params(data, input, output, residual, residual_is_local, 0);
const auto& push = ref.get()->get_push_obj();
check_push_launch(push, num_blocks, block_size);
auto params = make_params(*ref.get(), input, output, residual, residual_is_local);
bind_push(params, push);
const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_res_kernel<kWorldSize, true, kUsePDL>
: sp_collective::reduce_scatter_res_kernel<kWorldSize, false, kUsePDL>;
host::LaunchKernel(num_blocks + 1, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params);
}
static void all_gather(
CommunicatorRef ref,
TensorView input,
TensorView output,
int64_t ws_mc_base,
int64_t num_blocks,
int64_t block_size) {
const auto& data = *ref.get();
CHECK_HOST(ws_mc_base != 0) << "all-gather requires multicast workspace";
check_launch(data, num_blocks, block_size);
static void
all_gather(CommunicatorRef ref, TensorView input, TensorView output, int64_t num_blocks, int64_t block_size) {
const auto& push = ref.get()->get_push_obj();
check_push_launch(push, num_blocks, block_size);
// Reuse the RS matcher by swapping input/output roles conceptually.
auto params = make_params(data, output, input, std::nullopt, false, ws_mc_base);
auto params = make_params(*ref.get(), output, input, std::nullopt, false);
bind_push(params, push);
CHECK_HOST(params.push_ws_mc != nullptr) << "all-gather requires a multicast-capable push plane";
params.input = static_cast<const uint8_t*>(input.data_ptr());
params.output = static_cast<uint8_t*>(output.data_ptr());
host::LaunchKernel(num_blocks + 1, block_size, input.device())
@@ -397,19 +364,16 @@ struct SPCollectiveKernel {
TensorView input,
TensorView output,
int64_t output_mc_ptr,
int64_t sem_mc_ptr,
int64_t num_blocks,
int64_t block_size) {
const auto& data = *ref.get();
const auto& pull = ref.get()->get_pull_obj();
CHECK_HOST(output_mc_ptr != 0) << "direct all-gather needs symmetric output";
CHECK_HOST(sem_mc_ptr != 0) << "direct all-gather needs multicast semaphores";
CHECK_HOST(num_blocks > 0 && num_blocks <= data.num_pull_blocks);
CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0);
auto params = make_params(data, output, input, std::nullopt, false, 0);
check_pull_launch(ref.get()->get_pull_blocks(), num_blocks, block_size);
auto params = make_params(*ref.get(), output, input, std::nullopt, false);
bind_pull(params, pull);
params.input = static_cast<const uint8_t*>(input.data_ptr());
params.output = static_cast<uint8_t*>(output.data_ptr());
params.output_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(output_mc_ptr));
params.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr));
host::LaunchKernel(num_blocks, block_size, input.device())
.enable_pdl(kUsePDL)(sp_collective::all_gather_direct_kernel<kWorldSize, kUsePDL>, params);
}
@@ -421,17 +385,14 @@ struct SPCollectiveKernel {
std::optional<TensorView> residual,
bool residual_is_local,
int64_t input_mc_ptr,
int64_t sem_mc_ptr,
int64_t num_blocks,
int64_t block_size) {
const auto& data = *ref.get();
const auto& pull = ref.get()->get_pull_obj();
CHECK_HOST(input_mc_ptr != 0) << "pull RS needs symmetric input";
CHECK_HOST(sem_mc_ptr != 0) << "pull RS needs multicast semaphores";
CHECK_HOST(num_blocks > 0 && num_blocks <= data.num_pull_blocks);
CHECK_HOST(block_size >= 32 && block_size <= 1024 && block_size % 32 == 0);
auto params = make_params(data, input, output, residual, residual_is_local, 0);
check_pull_launch(ref.get()->get_pull_blocks(), num_blocks, block_size);
auto params = make_params(*ref.get(), input, output, residual, residual_is_local);
bind_pull(params, pull);
params.input_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(input_mc_ptr));
params.sem_mc = reinterpret_cast<uint8_t*>(static_cast<uintptr_t>(sem_mc_ptr));
const auto kernel = residual.has_value() ? sp_collective::reduce_scatter_pull_kernel<kWorldSize, true, kUsePDL>
: sp_collective::reduce_scatter_pull_kernel<kWorldSize, false, kUsePDL>;
host::LaunchKernel(num_blocks, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params);
python/sglang/kernels/jit/include/sgl_kernel/distributed/communicator.cuh
+443 -51
View File
@@ -1,16 +1,21 @@
#pragma once
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/utils.h>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/distributed/ptx.cuh>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/optional.h>
#include <algorithm>
#include <array>
#include <cstdint>
#include <map>
#include <limits>
#include <optional>
#include <string>
#include <type_traits>
#include <vector>
namespace sglang {
@@ -36,6 +41,8 @@ struct Counter {
uint32_t m_counter;
};
/// One block's arrival slot: a flag peers increment plus a phase counter.
/// Padded to a cache line so neighbouring blocks never share one.
struct alignas(128) Semaphore {
public:
Semaphore(const Semaphore&) = delete;
@@ -43,82 +50,467 @@ struct alignas(128) Semaphore {
return &m_counter;
}
SGL_DEVICE uint32_t get_relaxed() const {
uint32_t val;
asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag) : "memory");
return val;
return ptx::load_relaxed_sys(&m_flag);
}
SGL_DEVICE void put_relaxed() {
asm volatile("red.relaxed.sys.global.add.u32 [%0], 1;" : : "l"(&m_flag) : "memory");
ptx::red_add_relaxed_sys(&m_flag, 1);
}
SGL_DEVICE void put_relaxed_multicast() {
ptx::multimem_red_add_relaxed(&m_flag, 1);
}
SGL_DEVICE uint32_t get_acquire() const {
uint32_t val;
asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag) : "memory");
return val;
return ptx::load_acquire_sys(&m_flag);
}
SGL_DEVICE void put_release() {
asm volatile("red.release.sys.global.add.u32 [%0], 1;" : : "l"(&m_flag) : "memory");
ptx::red_add_release_sys(&m_flag, 1);
}
SGL_DEVICE void put_release_multicast() {
ptx::multimem_red_add_release(&m_flag, 1);
}
private:
uint32_t m_flag;
Counter m_counter;
};
static_assert(sizeof(Semaphore) == 128, "must match _SEMAPHORE_BYTES in custom_all_reduce_v2.py");
/// Kernel-facing slice of a push plane: `[offset, offset + size)` within every
/// slot, with `slot_bytes` as the stride from one slot to the next. Trivially
/// copyable, so it drops straight into a `__grid_constant__` params struct.
template <uint32_t kWorldSize>
struct PushWorkSpace {
std::array<uint8_t*, kWorldSize> workspaces;
Counter* counter; // NOTE: this is a local tensor, so no mc
uint8_t* mc_workspace;
uint32_t slot_bytes;
};
/// Kernel-facing slice of a pull plane. The semaphores are indexed by block,
/// not by byte, so they pass through the slice unchanged.
template <uint32_t kWorldSize>
struct PullWorkSpace {
std::array<Semaphore*, kWorldSize> semaphores;
std::array<uint8_t*, kWorldSize> workspaces;
Semaphore* mc_semaphore;
uint8_t* mc_workspace;
};
/// Bit patterns of the lamport "slot empty" marker. A producer rewrites any
/// +0.0 in its payload to -0.0 (numerically identical for the reduction) so a
/// consumer can treat a remaining +0.0 as "not arrived yet".
template <typename T>
struct FloatTrait {};
template <>
struct FloatTrait<bf16_t> {
using type = uint16_t;
static constexpr uint16_t kNegZero = 0x8000u;
};
template <>
struct FloatTrait<fp16_t> {
using type = uint16_t;
static constexpr uint16_t kNegZero = 0x8000u;
};
template <>
struct FloatTrait<float> {
using type = uint32_t;
static constexpr uint32_t kNegZero = 0x80000000u;
};
template <typename T, uint32_t N, uint32_t kAtom = sizeof(T)>
struct LamportTrait {
static_assert(kAtom >= sizeof(T) && (kAtom == 2 || kAtom == 4 || kAtom == 8));
static_assert(kAtom % sizeof(T) == 0 && N % (kAtom / sizeof(T)) == 0);
using Packed = std::conditional_t<kAtom == 2, uint16_t, std::conditional_t<kAtom == 4, uint32_t, uint64_t>>;
static constexpr Packed kNegZero = FloatTrait<T>::kNegZero;
static constexpr uint32_t kNumPacked = N / (kAtom / sizeof(T));
SGL_DEVICE static void clear_pos_zero(void* val) {
const auto ptr = static_cast<Packed*>(val);
#pragma unroll
for (uint32_t i = 0; i < kNumPacked; ++i) {
if (ptr[i] == 0) ptr[i] = kNegZero;
}
}
SGL_DEVICE static bool has_pos_zero(const void* val) {
const auto ptr = static_cast<const Packed*>(val);
bool result = false;
#pragma unroll
for (uint32_t i = 0; i < kNumPacked; ++i) {
result |= ptr[i] == 0;
}
return result;
}
SGL_DEVICE static void fill_pos_zero(void* val) {
const auto ptr = static_cast<Packed*>(val);
#pragma unroll
for (uint32_t i = 0; i < kNumPacked; ++i) {
ptr[i] = 0;
}
}
};
template <uint32_t kWorldSize>
struct Barrier {
public:
SGL_DEVICE Barrier(Semaphore* const* semaphores, uint32_t rank, uint32_t num_arrives)
: m_counter(0), m_rank(rank), m_semaphores(semaphores) {
const auto counter = semaphores[rank][blockIdx.x].counter_ptr();
const auto signal = num_arrives * kWorldSize;
m_counter = threadIdx.x == rank ? counter->inc(signal) : 0;
}
template <bool kNeedFence>
SGL_DEVICE void arrive(uint32_t n) const {
if (const auto tx = threadIdx.x; tx < kWorldSize) {
const auto bx = blockIdx.x;
const auto semaphore = &m_semaphores[tx][bx];
const auto current = m_counter + n * kWorldSize;
if constexpr (kNeedFence) {
semaphore->put_release();
if (tx == m_rank) {
while (semaphore->get_acquire() - current < kWorldSize)
;
}
} else {
semaphore->put_relaxed();
if (tx == m_rank) {
while (semaphore->get_relaxed() - current < kWorldSize)
;
}
}
}
}
SGL_DEVICE void arrive_relaxed(uint32_t n) const {
return this->arrive<false>(n);
}
SGL_DEVICE void arrive_rel_acq(uint32_t n) const {
return this->arrive<true>(n);
}
private:
uint32_t m_counter;
uint32_t m_rank;
Semaphore* const* m_semaphores;
};
/// Picks which half of a push plane's `2 * kWorldSize` slots this round owns.
/// The halves alternate because a round leaves its pos-zero markers behind: a
/// peer still draining the previous round must not see them refilled.
template <uint32_t kWorldSize>
struct PushEpoch {
public:
SGL_DEVICE PushEpoch(Counter* counter, uint8_t* const* workspaces, uint32_t slot_bytes)
: m_counter(counter), m_workspaces(workspaces), m_slot_bytes(slot_bytes), m_epoch(m_counter[blockIdx.x].get()) {}
SGL_DEVICE PushEpoch(const PushWorkSpace<kWorldSize>& ws)
: PushEpoch(ws.counter, ws.workspaces.data(), ws.slot_bytes) {}
/// Rank `src`'s slot inside rank `dst`'s workspace, for the current epoch.
SGL_DEVICE void* slot_ptr(uint32_t dst, uint32_t src = 0) const {
const auto epoch_stride_bytes = (m_epoch & 1) * m_slot_bytes * kWorldSize;
return m_workspaces[dst] + src * m_slot_bytes + epoch_stride_bytes;
}
SGL_DEVICE uint32_t slot_offset(uint32_t src = 0) const {
const auto epoch_stride_bytes = (m_epoch & 1) * m_slot_bytes * kWorldSize;
return src * m_slot_bytes + epoch_stride_bytes;
}
SGL_DEVICE void flip() const {
if (threadIdx.x == 0) m_counter[blockIdx.x].set(m_epoch ^ 1);
}
SGL_DEVICE void unsafe_flip_at(uint32_t bx) const {
m_counter[bx].set(m_epoch ^ 1);
}
SGL_DEVICE void unsafe_flip_range(uint32_t start, uint32_t finish) const {
for (uint32_t idx = start + threadIdx.x; idx < finish; idx += blockDim.x) {
m_counter[idx].set(m_epoch ^ 1);
}
}
private:
Counter* m_counter;
uint8_t* const* m_workspaces;
uint32_t m_slot_bytes;
uint32_t m_epoch;
};
/// Same window protocol as `Barrier`, but a single `multimem.red` reaches every
/// peer's row at once instead of a `kWorldSize`-wide unicast fan-out. One
/// thread drives the whole barrier, so `world_size` need not be a constant and
/// nothing here has to be a template.
///
/// Construction only reserves the window; signalling happens in `arrive`. Keep
/// them separate at the call site: the reservation is worth hoisting above a
/// PDL wait (it takes the RMW latency off the post-wait critical path) while
/// the signal must stay after it, since it asserts the producer grid flushed.
struct McBarrier {
public:
SGL_DEVICE McBarrier(Semaphore* local, Semaphore* mc, uint32_t world_size, uint32_t num_arrives)
: m_counter(0), m_world_size(world_size), m_local(local), m_mc(mc) {
if (threadIdx.x == 0) {
m_counter = local[blockIdx.x].counter_ptr()->inc(num_arrives * world_size);
}
}
template <bool kNeedFence>
SGL_DEVICE void arrive(uint32_t n) const {
return arrive_at<kNeedFence>(m_local, m_mc, m_world_size, m_counter + n * m_world_size);
}
SGL_DEVICE void arrive_relaxed(uint32_t n) const {
return this->arrive<false>(n);
}
SGL_DEVICE void arrive_rel_acq(uint32_t n) const {
return this->arrive<true>(n);
}
/// The reserved window base, valid in thread 0. A kernel whose body is
/// register-hungry enough that keeping this object live would spill can park
/// this one word in shared memory and finish through `arrive_at` instead.
SGL_DEVICE uint32_t window() const {
return m_counter;
}
/// `arrive` against a window reserved earlier, without holding the object.
/// `window` already includes the `n * world_size` offset.
template <bool kNeedFence>
SGL_DEVICE static void arrive_at(Semaphore* local, Semaphore* mc, uint32_t world_size, uint32_t window) {
if (threadIdx.x != 0) return;
const auto bx = blockIdx.x;
const auto semaphore = &local[bx];
const auto mc_semaphore = &mc[bx];
if constexpr (kNeedFence) {
mc_semaphore->put_release_multicast();
while (semaphore->get_acquire() - window < world_size)
;
} else {
mc_semaphore->put_relaxed_multicast();
while (semaphore->get_relaxed() - window < world_size)
;
}
}
private:
uint32_t m_counter; // window base; meaningful in thread 0, the sole poller
uint32_t m_world_size;
Semaphore* m_local;
Semaphore* m_mc;
};
} // namespace device::distributed
namespace host::distributed {
using device::distributed::Counter, device::distributed::Semaphore;
using device::distributed::PullWorkSpace, device::distributed::PushWorkSpace;
using TensorView = tvm::ffi::TensorView;
template <typename T>
using Optional = tvm::ffi::Optional<T>;
inline constexpr uint32_t kMaxWorldSize = device::distributed::kMaxWorldSize;
/**
* \brief Storage plane of the custom all-reduce implementation.
*
* A thin, kernel-agnostic view over externally owned buffers: per-rank
* symmetric workspaces, synchronization primitives, and grid-size settings.
* It performs no allocation and no IPC; the Python side owns the storage
* (symmetric memory) and its lifetime.
*/
struct CommunicatorObj : public tvm::ffi::Object {
public:
using TensorView = tvm::ffi::TensorView;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.Communicator", CommunicatorObj, tvm::ffi::Object);
static constexpr bool _type_mutable = true; // config() mutates block counts
// Defined in csrc/distributed/communicator.cuh (only the registration
// module needs the implementation).
CommunicatorObj(
uint32_t rank,
uint32_t world_size,
std::vector<TensorView> push_workspaces,
std::vector<TensorView> pull_workspaces,
std::vector<TensorView> pull_semaphores,
TensorView push_counter,
std::optional<int64_t> pull_mc_workspace_ptr);
void config(std::map<std::string, uint32_t> config);
/// Multicast VAs are optional, so keep null null instead of forming
/// `nullptr + offset` when slicing a plane that has no multicast mapping.
inline uint8_t* offset_mc(uint8_t* base, int64_t offset) {
return base != nullptr ? base + offset : nullptr;
}
/// Identity shared by every plane; the constructor is the one place that
/// validates it (defined in csrc/distributed/registry.cuh).
struct BasePlane {
BasePlane(const BasePlane&) = delete;
BasePlane(uint32_t rank, uint32_t world_size);
uint32_t rank;
uint32_t world_size;
int64_t push_bytes; // per-buffer bytes; each rank holds 2 * world_size buffers
int64_t pull_bytes;
uint32_t num_push_blocks; // not configurable (bound to the counter array)
uint32_t num_pull_blocks;
uint32_t num_multicast_blocks;
std::array<uint8_t*, kMaxWorldSize> pull_workspaces; // symmetric memory
std::array<uint8_t*, kMaxWorldSize> push_workspaces; // symmetric memory
std::array<Semaphore*, kMaxWorldSize> pull_semaphores; // symmetric memory
Counter* push_counter; // local memory
uint8_t* pull_mc_workspace; // multicast address of the pull workspace (may be null)
private: // upper bounds for config()
uint32_t total_pull_blocks;
};
struct CommunicatorRef : public tvm::ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommunicatorRef, tvm::ffi::ObjectRef, CommunicatorObj);
/**
* \brief Lamport push plane: a zero-filled symmetric workspace plus a
* rank-local phase counter.
*
* Each rank owns `2 * world_size` slots of `slot_bytes` (two phases x one
* producer slot per peer). Producers store into the destination rank's slot
* and consumers poll for the pos-zero marker to clear, so the workspace MUST
* be zero-filled before first use and every kernel MUST restore the marker
* on its way out. `counter` is rank-local (never read by a peer).
*
* Holds no storage: the caller (Python) owns the tensors and their lifetime.
*/
struct PushPlaneObj : public tvm::ffi::Object, BasePlane {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.distributed.PushPlane", PushPlaneObj, tvm::ffi::Object);
// Defined in csrc/distributed/registry.cuh (only the registration module needs the implementation).
PushPlaneObj(
uint32_t rank,
uint32_t world_size,
std::vector<TensorView> workspaces, // world_size * [2 * world_size][slot_bytes]
TensorView counter, // [num_blocks]
intptr_t mc_workspace_ptr);
uint32_t num_blocks; // bound to the counter array, hence not tunable
int64_t slot_bytes; // per-slot bytes; each rank holds 2 * world_size slots
Counter* counter; // rank-local memory
std::array<uint8_t*, kMaxWorldSize> workspaces; // symmetric memory
uint8_t* mc_workspace; // multicast VA of the local workspace (may be null)
template <uint32_t N>
PushWorkSpace<N> get_workspace(int64_t size, int64_t offset = 0) const {
CHECK_HOST(N == world_size) << "Plane holds " << world_size << " ranks, asked for " << N;
CHECK_HOST(size >= 0 && offset >= 0 && offset + size <= slot_bytes)
<< "slice [" << offset << ", " << offset + size << ") escapes the " << slot_bytes << "-byte push slot";
// Device-side phase striding is 32-bit: the largest offset a kernel forms
// is `(2 * N - 1) * slot_bytes`, so the whole double-buffered plane must fit.
CHECK_HOST(2 * N * slot_bytes <= std::numeric_limits<uint32_t>::max())
<< 2 * N * slot_bytes << " bytes of push plane exceeds the 32-bit offset range";
PushWorkSpace<N> ws{{}, counter, offset_mc(mc_workspace, offset), static_cast<uint32_t>(slot_bytes)};
for (uint32_t i = 0; i < N; ++i) {
ws.workspaces[i] = workspaces[i] + offset;
}
return ws;
}
};
/**
* \brief Pull plane: symmetric per-rank buffers plus the per-block barrier
* semaphores that guard them.
*
* Either half may be absent -- pass a 0-element tensor for the one you do not
* own, which leaves the corresponding `num_bytes` / `num_blocks` at zero and
* makes any kernel needing it fail with a clear message:
*
* - workspaces only: nothing today, but the shape a caller who brings its own
* symmetric buffers would take.
* - semaphores only: the K3 fused collectives, which reduce in place on the
* caller's own symmetric input (its multicast VA arrives per call, since it
* varies with the slice) and borrow this plane purely to barrier on.
* - both: the generic custom all-reduce, whose callers hand it plain tensors,
* so it stages them through `workspaces` before reducing.
*
* Holds no storage: the caller (Python) owns the tensors and their lifetime.
*/
struct PullPlaneObj : public tvm::ffi::Object, BasePlane {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.distributed.PullPlane", PullPlaneObj, tvm::ffi::Object);
// Defined in csrc/distributed/registry.cuh (only the registration module needs the implementation).
PullPlaneObj(
uint32_t rank,
uint32_t world_size,
std::vector<TensorView> workspaces, // world_size * [num_bytes]
std::vector<TensorView> semaphores, // world_size * [num_blocks]
intptr_t mc_workspace_ptr,
intptr_t mc_semaphore_ptr);
uint32_t num_blocks; // semaphore capacity; callers clamp their grid to it
int64_t num_bytes; // per-rank workspace bytes
std::array<Semaphore*, kMaxWorldSize> semaphores; // symmetric memory
std::array<uint8_t*, kMaxWorldSize> workspaces; // symmetric memory
Semaphore* mc_semaphore; // multicast VA of the local semaphores (may be null)
uint8_t* mc_workspace; // multicast VA of the local workspace (may be null)
template <uint32_t N>
PullWorkSpace<N> get_workspace(int64_t size, int64_t offset = 0) const {
CHECK_HOST(N == world_size) << "Plane holds " << world_size << " ranks, asked for " << N;
CHECK_HOST(size >= 0 && offset >= 0 && offset + size <= num_bytes)
<< "slice [" << offset << ", " << offset + size << ") escapes the " << num_bytes << "-byte pull workspace";
PullWorkSpace<N> ws{{}, {}, mc_semaphore, offset_mc(mc_workspace, offset)};
for (uint32_t i = 0; i < N; ++i) {
ws.semaphores[i] = semaphores[i];
ws.workspaces[i] = workspaces[i] + offset;
}
return ws;
}
};
SGLANG_REGISTER_FFI_REFERENCE_CLASS(PushPlaneRef, PushPlaneObj);
SGLANG_REGISTER_FFI_REFERENCE_CLASS(PullPlaneRef, PullPlaneObj);
/**
* \brief The planes every kernel in this directory takes, plus the launch
* widths that are tuning rather than capacity.
*
* A plane is absent when the owner never uses that half (a push-only instance
* passes `pull=None`), and asking for it then fails with a clear message
* instead of silently reading a placeholder buffer.
*/
struct CommunicatorObj : public tvm::ffi::Object {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.distributed.Communicator", CommunicatorObj, tvm::ffi::Object);
static constexpr bool _type_mutable = true; // the set_pull_* knobs below
// Only the constructors live in csrc/distributed/registry.cuh: kernel
// modules are separate shared objects that see this header but never link
// that translation unit, so everything they call has to be inline here.
CommunicatorObj(Optional<PushPlaneRef> push, Optional<PullPlaneRef> pull);
uint32_t get_rank() const {
return m_push.has_value() ? m_push.value()->rank : m_pull.value()->rank;
}
uint32_t get_world_size() const {
return m_push.has_value() ? m_push.value()->world_size : m_pull.value()->world_size;
}
bool has_push() const {
return m_push.has_value();
}
bool has_pull() const {
return m_pull.has_value();
}
/// The planes as handles, for callers that hold on to one (Python). Kernels
/// want `get_*_obj()` below, which skips the refcount traffic.
Optional<PushPlaneRef> get_push() const {
return m_push;
}
Optional<PullPlaneRef> get_pull() const {
return m_pull;
}
const PushPlaneObj& get_push_obj() const {
CHECK_HOST(m_push.has_value()) << "This communicator has no push plane";
return *m_push.value().get();
}
const PullPlaneObj& get_pull_obj() const {
CHECK_HOST(m_pull.has_value()) << "This communicator has no pull plane";
return *m_pull.value().get();
}
/// Both knobs only ever narrow the grid: unset means "the plane's capacity",
/// and the multicast width additionally never exceeds the pull width, since
/// extra multicast traffic costs NVLS throughput.
void set_pull_blocks(std::optional<uint32_t> num_blocks) {
m_pull_blocks = num_blocks;
}
void set_pull_multicast_blocks(std::optional<uint32_t> num_blocks) {
m_pull_multicast_blocks = num_blocks;
}
uint32_t get_pull_blocks() const {
const auto capacity = get_pull_obj().num_blocks;
CHECK_HOST(capacity > 0) << "This pull plane has no semaphores to barrier on";
return m_pull_blocks.has_value() ? std::min(*m_pull_blocks, capacity) : capacity;
}
uint32_t get_pull_multicast_blocks() const {
const auto blocks = get_pull_blocks();
return m_pull_multicast_blocks.has_value() ? std::min(*m_pull_multicast_blocks, blocks) : blocks;
}
private:
std::optional<uint32_t> m_pull_blocks;
std::optional<uint32_t> m_pull_multicast_blocks;
Optional<PushPlaneRef> m_push;
Optional<PullPlaneRef> m_pull;
};
SGLANG_REGISTER_FFI_REFERENCE_CLASS(CommunicatorRef, CommunicatorObj);
} // namespace host::distributed
} // namespace sglang
python/sglang/kernels/jit/include/sgl_kernel/distributed/ptx.cuh
+237
View File
@@ -0,0 +1,237 @@
#pragma once
// Inline PTX the distributed collectives need and no CUDA intrinsic covers:
// system-scope flag traffic, multicast (NVLS) reductions, and the 16B
// vectorized global accesses the reduce loops are built on.
//
// One header rather than a copy per kernel: before this existed the multimem
// flag increments had three identical definitions, and the 16B helpers were
// defined at `namespace sglang` scope inside custom_all_reduce.cuh, so every
// K3 kernel picked them up transitively and broke whenever that file's
// includes changed.
//
// Names mirror their PTX mnemonic (`red.relaxed.sys.global.add.u32` ->
// `red_add_relaxed_sys`) so a call site can be checked against the ISA docs
// without indirection.
//
// These join `sglang::ptx`, the namespace sgl_kernel/mbarrier.cuh opens for
// inline PTX: one vocabulary for every family, split across headers by what
// the instructions are for.
#include <sgl_kernel/utils.h>
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <sgl_kernel/vec.cuh>
#include <cstdint>
#include <type_traits>
namespace sglang {
namespace device::ptx {
// ---------------------------------------------------------------------------
// 16B vectorized global access
//
// All six take a base pointer plus a *vector* index, and require `V` to be
// exactly one 16B vector: the reduce loops force 16B accesses to keep register
// pressure down, so the width is a static contract rather than a parameter.
// ---------------------------------------------------------------------------
template <typename V>
SGL_DEVICE void ld_global_16B(V& x, const void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
addr = static_cast<const uint8_t*>(addr) + vec_offset * sizeof(V);
uint4 val;
asm volatile("ld.global.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(addr));
x = *reinterpret_cast<const V*>(&val);
}
template <typename V>
SGL_DEVICE void st_global_16B(const V& x, void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
const uint4 val = *reinterpret_cast<const uint4*>(&x);
addr = static_cast<uint8_t*>(addr) + vec_offset * sizeof(V);
asm volatile("st.global.v4.b32 [%4], {%0, %1, %2, %3};"
: //
: "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
}
/// Peer-visible load. Relaxed: ordering comes from the surrounding barrier or
/// from the payload's own marker (lamport polling), not from here.
///
/// Every scoped access here is `.sys`, including the ones that only ever poll
/// memory on this device: load scope measures the same at `.gpu`, so one scope
/// keeps the vocabulary small. Store scope is not interchangeable -- a store a
/// peer must observe has to be `.sys`.
template <typename V>
SGL_DEVICE void ld_relaxed_16B(V& x, const void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
addr = static_cast<const uint8_t*>(addr) + vec_offset * sizeof(V);
uint4 val;
asm volatile("ld.relaxed.sys.global.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(addr));
x = *reinterpret_cast<const V*>(&val);
}
template <typename V>
SGL_DEVICE void st_relaxed_16B(const V& x, void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 16 && sizeof(V) == 16);
const uint4 val = *reinterpret_cast<const uint4*>(&x);
addr = static_cast<uint8_t*>(addr) + vec_offset * sizeof(V);
asm volatile("st.relaxed.sys.global.v4.b32 [%4], {%0, %1, %2, %3};"
: //
: "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
}
/// One load that sums the corresponding vector across every rank in the
/// multicast team (NVLS). `mc_addr` must be a multicast VA.
template <typename V>
SGL_DEVICE void ld_multimem_16B(V& x, const void* mc_addr, int64_t vec_offset) {
#if SGL_ARCH_HOPPER_OR_GREATER
static_assert(alignof(V) == 16 && sizeof(V) == 16);
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];"
: "=f"(val.x), "=f"(val.y), "=f"(val.z), "=f"(val.w)
: "l"(mc_addr));
x = *reinterpret_cast<const V*>(&val);
} else {
// Packed f16x2/bf16x2 results live in b32 registers ("=r"); .acc::f32 only
// raises the accumulation precision, not the result register type -- ptxas
// 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];"
: "=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];"
: "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
: "l"(mc_addr));
}
x = *reinterpret_cast<const V*>(&val);
}
#else
assert(false && "multimem load is only supported on Hopper or later architecture");
#endif
}
/// 8B relaxed pair, for the narrow (2 x fp32) peer accumulators the fused
/// qk-norm exchanges. Replaces `ld.volatile` / `st.volatile`, which PTX defines
/// as relaxed at system scope.
template <typename V>
SGL_DEVICE void ld_relaxed_8B(V& x, const void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 8 && sizeof(V) == 8);
addr = static_cast<const uint8_t*>(addr) + vec_offset * sizeof(V);
uint2 val;
asm volatile("ld.relaxed.sys.global.v2.b32 {%0, %1}, [%2];" : "=r"(val.x), "=r"(val.y) : "l"(addr) : "memory");
x = *reinterpret_cast<const V*>(&val);
}
template <typename V>
SGL_DEVICE void st_relaxed_8B(const V& x, void* addr, int64_t vec_offset) {
static_assert(alignof(V) == 8 && sizeof(V) == 8);
const uint2 val = *reinterpret_cast<const uint2*>(&x);
addr = static_cast<uint8_t*>(addr) + vec_offset * sizeof(V);
asm volatile("st.relaxed.sys.global.v2.b32 [%2], {%0, %1};" ::"r"(val.x), "r"(val.y), "l"(addr) : "memory");
}
/// One store fanned out to every rank in the multicast team.
template <typename V>
SGL_DEVICE void st_multimem_16B(const V& x, void* mc_addr, int64_t vec_offset) {
#if SGL_ARCH_HOPPER_OR_GREATER
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};"
:
: "f"(val.x), "f"(val.y), "f"(val.z), "f"(val.w), "l"(mc_addr));
#else
assert(false && "multimem store is only supported on Hopper or later architecture");
#endif
}
// ---------------------------------------------------------------------------
// System-scope flag traffic
//
// A flag is a u32 that a peer polls. `red` is a reduction with no result
// register, which is what a fire-and-forget increment wants; `atom` returns the
// old value for callers that need to know who arrived last.
// ---------------------------------------------------------------------------
SGL_DEVICE uint32_t load_relaxed_sys(const uint32_t* ptr) {
uint32_t val;
asm volatile("ld.relaxed.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(ptr) : "memory");
return val;
}
/// Acquire load of a peer-written flag: everything the writer released before
/// its matching store is visible to this thread afterwards.
SGL_DEVICE uint32_t load_acquire_sys(const uint32_t* ptr) {
uint32_t val;
asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(ptr) : "memory");
return val;
}
/// Peer-visible flag increment. Relaxed: ordering is established by the
/// surrounding fence, not by this store.
SGL_DEVICE void red_add_relaxed_sys(uint32_t* ptr, uint32_t val) {
asm volatile("red.relaxed.sys.global.add.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
}
/// Flag increment that also publishes every prior write to system scope.
SGL_DEVICE void red_add_release_sys(uint32_t* ptr, uint32_t val) {
asm volatile("red.release.sys.global.add.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");
}
/// Publishes every prior write to system scope. Pair with a relaxed flag store
/// so a peer's acquire load of that flag also observes the payload.
SGL_DEVICE void fence_release_sys() {
asm volatile("fence.release.sys;" ::: "memory");
}
/// Device-scope arrival counter. acq_rel so the winner of the count also
/// observes the losers' payload writes.
SGL_DEVICE uint32_t atomic_add_acq_rel_gpu(uint32_t* ptr, uint32_t val) {
uint32_t old;
asm volatile("atom.acq_rel.gpu.global.add.u32 %0, [%1], %2;" : "=r"(old) : "l"(ptr), "r"(val) : "memory");
return old;
}
// ---------------------------------------------------------------------------
// Multicast flag traffic
//
// One instruction updates the flag on every rank in the team, replacing a
// world_size-wide unicast fan-out.
// ---------------------------------------------------------------------------
SGL_DEVICE void multimem_store_relaxed(uint32_t* mc_ptr, uint32_t val) {
asm volatile("multimem.st.relaxed.sys.global.b32 [%0], %1;" : : "l"(mc_ptr), "r"(val) : "memory");
}
SGL_DEVICE void multimem_red_add_relaxed(uint32_t* mc_ptr, uint32_t val) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.relaxed.sys.global.add.u32 [%0], %1;" : : "l"(mc_ptr), "r"(val) : "memory");
#else
assert(false && "multimem red is only supported on Hopper or later architecture");
#endif
}
SGL_DEVICE void multimem_red_add_release(uint32_t* mc_ptr, uint32_t val) {
#if SGL_ARCH_HOPPER_OR_GREATER
asm volatile("multimem.red.release.sys.global.add.u32 [%0], %1;" : : "l"(mc_ptr), "r"(val) : "memory");
#else
assert(false && "multimem red is only supported on Hopper or later architecture");
#endif
}
} // namespace device::ptx
} // namespace sglang
python/sglang/kernels/jit/include/sgl_kernel/ffi.h
+17
View File
@@ -103,6 +103,23 @@ inline Tensor from_blob_like(
return from_blob(data, t.shape(), t.dtype(), t.device(), std::forward<Fn>(deleter), stride, byte_offset);
}
inline Tensor alloc_workspace_tensor(size_t required_bytes, DLDevice device) {
if (required_bytes == 0) return {};
DLDataType u8 = {kDLUInt, 8, 1};
int64_t shape[] = {static_cast<int64_t>(required_bytes)};
return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device);
}
} // namespace host::ffi
// Declare `REF`, the not-nullable ObjectRef handle for `OBJ`.
//
// Whether `ref->` hands out a mutable pointer is decided by
// `OBJ::_type_mutable` (see TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE),
// so it belongs on the object, not here.
#define SGLANG_REGISTER_FFI_REFERENCE_CLASS(REF, OBJ) \
struct REF : public tvm::ffi::ObjectRef { \
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(REF, tvm::ffi::ObjectRef, OBJ); \
}
} // namespace sglang
python/sglang/kernels/jit/include/sgl_kernel/mbarrier.cuh
+11 -18
View File
@@ -1,15 +1,8 @@
#pragma once
// `sglang::device::ptx` is the one namespace for inline PTX; both files reopen
// it to add the SM-local primitives only they need, and
// distributed/ptx.cuh adds the collective ones.
// mbarrier PTX wrappers shared by the Kimi K3 kernels that drive TMA by hand:
// kimi_k3/comm/gemm_ar.cuh and kimi_k3/attn_res/fused_tma.cuh both defined these
// with identical bodies.
//
// The enclosing namespace is the same `sglang::ptx` both files already open, so
// existing `ptx::mbar_*` call sites need no change.
//
// gemm_ar.cuh keeps mbar_arrive_cluster_release: only it uses that one.
// attention/kda_prefill.cu duplicates a different set (MMA / ldmatrix) but is
// built without a sglang include path and cannot consume this header.
#pragma once
#include <sgl_kernel/utils.cuh>
@@ -17,12 +10,12 @@
namespace sglang {
namespace ptx {
namespace device::ptx {
// Inline-PTX `.shared` instructions take a 32-bit byte offset in the shared
// window, not a generic 64-bit pointer.
template <typename T>
static SGL_DEVICE uint32_t to_shared(T* ptr) {
SGL_DEVICE uint32_t to_shared(T* ptr) {
return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
}
@@ -40,18 +33,18 @@ static SGL_DEVICE uint32_t to_shared(T* ptr) {
// has touched yet) -> 1, so the first wait is a no-op skip.
// A consumer-first wait initialized to 1 skips the producer's first signal and
// blocks forever on the second.
static SGL_DEVICE void mbar_init(uint64_t* bar, uint32_t count) {
SGL_DEVICE void mbar_init(uint64_t* bar, uint32_t count) {
asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(to_shared(bar)), "r"(count));
}
static SGL_DEVICE uint64_t mbar_arrive(uint64_t* bar) {
SGL_DEVICE uint64_t mbar_arrive(uint64_t* bar) {
uint64_t state;
asm volatile("mbarrier.arrive.shared.b64 %0, [%1];" : "=l"(state) : "r"(to_shared(bar)));
return state;
}
// Combined arrive + set tx-count, for TMA-load completion.
static SGL_DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) {
SGL_DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) {
asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(to_shared(bar)), "r"(bytes));
}
@@ -59,7 +52,7 @@ static SGL_DEVICE void mbar_arrive_expect_tx(uint64_t* bar, uint32_t bytes) {
// early wakeups. Default `.acquire` semantics mean prior `cp.async.bulk` writes
// tracked by this mbarrier are visible to later generic-proxy reads on this
// thread with no `fence.proxy.async` (spec §9.7.13.15.16 point 3).
static SGL_DEVICE void mbar_wait_parity(uint64_t* bar, uint32_t parity) {
SGL_DEVICE void mbar_wait_parity(uint64_t* bar, uint32_t parity) {
asm volatile(
"{\n\t.reg .pred p;\n\t"
"WAIT_%=: mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n\t"
@@ -67,6 +60,6 @@ static SGL_DEVICE void mbar_wait_parity(uint64_t* bar, uint32_t parity) {
"r"(parity));
}
} // namespace ptx
} // namespace device::ptx
} // namespace sglang
python/sglang/kernels/jit/include/sgl_kernel/utils.cuh
-8
View File
@@ -15,7 +15,6 @@
#pragma once
#include <sgl_kernel/ffi.h>
#include <sgl_kernel/utils.h>
#include <dlpack/dlpack.h>
@@ -246,13 +245,6 @@ inline void RuntimeDeviceCheck(DebugInfo location = {}) {
return RuntimeDeviceCheck(::cudaGetLastError(), location);
}
inline auto alloc_workspace_tensor(size_t required_bytes, DLDevice device) -> tvm::ffi::Tensor {
if (required_bytes == 0) return {};
DLDataType u8 = {kDLUInt, 8, 1};
int64_t shape[] = {static_cast<int64_t>(required_bytes)};
return ffi::empty(tvm::ffi::ShapeView(shape, 1), u8, device);
}
/**
* \brief Kernel launcher with automatic stream resolution and PDL support.
*
python/sglang/kernels/jit/include/sgl_kernel/vec.cuh
+39
View File
@@ -7,6 +7,7 @@
/// 256 bits (32 bytes), matching CUDA's widest vector load.
#pragma once
#include <sgl_kernel/type.cuh>
#include <sgl_kernel/utils.cuh>
#include <cstddef>
@@ -117,6 +118,44 @@ struct AlignedVector {
storage_t m_storage;
};
/// Sum `M` vectors element-wise into one, accumulating in fp32 regardless of
/// the packed element type. Used by every collective that reduces peer
/// contributions in registers.
template <bool kFP32Acc = true, typename T2, size_t N, size_t M>
SGL_DEVICE auto reduce_vec(device::AlignedVector<T2, N> (&vec)[M]) -> device::AlignedVector<T2, N> {
static_assert(DTypeTrait<T2>::kVecSize == 2, "reduce_vec only supports 2-element vectors for now");
static_assert(M > 0, "reduce_vec requires at least one vector to reduce");
if constexpr (kFP32Acc) {
fp32x2_t acc[N];
#pragma unroll
for (size_t i = 0; i < M; ++i) {
#pragma unroll
for (size_t j = 0; j < N; ++j) {
const auto [x, y] = cast<fp32x2_t>(vec[i][j]);
acc[j].x = i == 0 ? x : acc[j].x + x;
acc[j].y = i == 0 ? y : acc[j].y + y;
}
}
device::AlignedVector<T2, N> out_vec;
#pragma unroll
for (size_t j = 0; j < N; ++j) {
out_vec[j] = cast<T2>(acc[j]);
}
return out_vec;
} else {
using SumOp = ReductionTrait<ReductionOp::SUM, T2>;
device::AlignedVector<T2, N> out_vec = vec[0];
#pragma unroll
for (size_t i = 1; i < M; ++i) {
#pragma unroll
for (size_t j = 0; j < N; ++j) {
out_vec[j] = SumOp::apply(out_vec[j], vec[i][j]);
}
}
return out_vec;
}
}
} // namespace device
} // namespace sglang
python/sglang/kernels/ops/communication/all_reduce.py
+121 -53
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
import enum
from typing import TYPE_CHECKING, List, Tuple, Union
from typing import TYPE_CHECKING, List, Tuple
import torch
import tvm_ffi
@@ -9,6 +9,7 @@ from tvm_ffi import Module
from sglang.kernels.jit.utils import (
cache_once,
empty_sentinel,
is_arch_support_pdl,
lazy_register_class,
load_jit,
@@ -36,82 +37,146 @@ _ALGO_NAMES = {
AllReduceAlgo.TWO_SHOT_PULL: "2shot_pull",
}
# ``pull_arg`` of the all-reduce kernel: a row of the graph-params pointer
# table selects graph mode; a plain bool selects multicast (True) / eager.
PullArg = Union[torch.Tensor, bool]
if TYPE_CHECKING:
# (cudaIpcMemHandle bytes, offset-in-allocation) for one device pointer
IPC_HANDLE_PAIR = Tuple[List[int], int]
@cache_once
def _init_communicator() -> None:
module = load_jit(
"communicator",
cuda_files=["distributed/communicator.cuh"],
cuda_wrappers=[("register_once", "register_communicator")],
cuda_files=["distributed/registry.cuh"],
cuda_wrappers=[("register_communicator", "register_communicator")],
)
module.register_once()
module.register_communicator()
@lazy_register_class("sgl.Communicator", _init_communicator)
class Communicator(tvm_ffi.Object):
"""Storage plane of the custom all-reduce: a thin pointer holder.
@lazy_register_class("sgl.distributed.PushPlane", _init_communicator)
class PushPlane(tvm_ffi.Object):
"""Lamport push plane: a zero-filled symmetric workspace + a local counter.
All buffers are owned by the caller (symmetric-memory tensor views plus
a local push counter); this object only validates and records them.
All buffers are owned by the caller; this object only validates and
records them.
"""
# C++ interface
if TYPE_CHECKING:
# C++ interface
rank: int
world_size: int
def _config(self, kwargs: dict) -> None: ...
def __init__(
self,
rank: int,
world_size: int,
push_workspaces: List[torch.Tensor],
pull_workspaces: List[torch.Tensor],
pull_semaphores: List[torch.Tensor],
push_counter: torch.Tensor,
pull_mc_workspace: int | None,
*,
workspaces: List[torch.Tensor],
counter: torch.Tensor,
mc_workspace: int | None = None,
) -> None:
"""
:param push_workspaces: per-rank ``[2 * world_size, push_bytes]``
uint8 views of symmetric memory.
:param pull_workspaces: per-rank ``[pull_bytes]`` uint8 views of
symmetric memory.
:param pull_semaphores: per-rank ``[num_pull_blocks, 128]`` uint8
views of symmetric memory.
:param push_counter: local ``[num_push_blocks, 4]`` uint8 tensor.
:param pull_mc_workspace: multicast address of the pull workspace,
or None when multicast is unavailable.
:param workspaces: per-rank ``[2 * world_size, slot_bytes]`` uint8
views of symmetric memory. The local rank's view
MUST be zero-filled before first use -- the
kernels poll for a pos-zero marker.
:param counter: local ``[num_blocks, 4]`` uint8 tensor, zero-filled.
:param mc_workspace: multicast VA of the local workspace, or None.
"""
self.__ffi_init__(rank, world_size, workspaces, counter, mc_workspace or 0)
@lazy_register_class("sgl.distributed.PullPlane", _init_communicator)
class PullPlane(tvm_ffi.Object):
"""Symmetric per-rank buffers plus the per-block semaphores guarding them.
Either half may be omitted; the plane then holds a 0-element tensor in its
place and any kernel needing that half fails with a clear message. The K3
fused collectives take semaphores only -- they reduce in place on the
caller's own symmetric input -- while the generic all-reduce takes both,
since it stages plain tensors through the buffers before reducing.
"""
# C++ interface
if TYPE_CHECKING:
rank: int
world_size: int
def __init__(
self,
rank: int,
world_size: int,
*,
workspaces: List[torch.Tensor] | None = None,
semaphores: List[torch.Tensor] | None = None,
mc_workspace: int | None = None,
mc_semaphore: int | None = None,
) -> None:
"""
:param workspaces: per-rank ``[num_bytes]`` uint8 views of symmetric
memory, or None when the caller brings its own.
:param semaphores: per-rank ``[num_blocks, 128]`` uint8 views of
symmetric memory, zero-filled before first use, or
None when the caller never barriers on this plane.
:param mc_workspace: multicast VA of the local workspace, or None.
:param mc_semaphore: multicast VA of the local semaphores, or None.
"""
if workspaces is None:
device = torch.device("cuda", torch.cuda.current_device())
sentinel = empty_sentinel(device, torch.uint8).view(-1)
workspaces = [sentinel for _ in range(world_size)]
if semaphores is None:
device = torch.device("cuda", torch.cuda.current_device())
sentinel = empty_sentinel(device, torch.uint8).view(-1, 128)
semaphores = [sentinel for _ in range(world_size)]
mc_workspace = mc_workspace or 0
mc_semaphore = mc_semaphore or 0
self.__ffi_init__(
rank,
world_size,
push_workspaces,
pull_workspaces,
pull_semaphores,
push_counter,
pull_mc_workspace,
rank, world_size, workspaces, semaphores, mc_workspace, mc_semaphore
)
def config(
@lazy_register_class("sgl.distributed.Communicator", _init_communicator)
class Communicator(tvm_ffi.Object):
"""The planes shared by every kernel in ``kernels.ops.communication``.
Pass ``None`` for a plane the owner never uses; kernels that need it then
fail with a clear message instead of reading a placeholder buffer.
"""
if TYPE_CHECKING:
# C++ interface
def get_rank(self) -> int: ...
def get_world_size(self) -> int: ...
def get_push(self) -> PushPlane | None: ...
def get_pull(self) -> PullPlane | None: ...
def set_pull_blocks(self, num_blocks: int | None) -> None: ...
def set_pull_multicast_blocks(self, num_blocks: int | None) -> None: ...
def __init__(
self,
num_pull_blocks: int | None = None,
num_multicast_blocks: int | None = None,
) -> Communicator:
kwargs = {}
if num_pull_blocks is not None:
kwargs["num_pull_blocks"] = num_pull_blocks
if num_multicast_blocks is not None:
kwargs["num_multicast_blocks"] = num_multicast_blocks
self._config(kwargs)
return self
push: PushPlane | None = None,
pull: PullPlane | None = None,
) -> None:
self.__ffi_init__(push, pull)
@property
def rank(self) -> int:
return self.get_rank()
@property
def world_size(self) -> int:
return self.get_world_size()
@property
def push(self) -> PushPlane | None:
"""The push plane, or None for a pull-only communicator."""
return self.get_push()
@property
def pull(self) -> PullPlane | None:
"""The pull plane, or None for a push-only communicator."""
return self.get_pull()
def _init_ipc_manager() -> None:
@@ -126,7 +191,7 @@ def _init_ipc_manager() -> None:
@lazy_register_class("sgl.IPCManager", _init_ipc_manager)
class IPCManager(tvm_ffi.Object):
"""Batched cudaIpc handle exchange for CUDA-graph input pointers."""
"""Batched cudaIPC handle exchange for CUDA-graph input pointers."""
if TYPE_CHECKING:
# C++ interface
@@ -145,7 +210,7 @@ def get_all_reduce_module(dtype: torch.dtype, world_size: int) -> Module:
"custom_all_reduce",
*args,
cuda_files=["distributed/custom_all_reduce.cuh"],
cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
cuda_wrappers=[("all_reduce", f"AllReduceKernel<{args}>::run")],
)
@@ -154,10 +219,13 @@ def custom_all_reduce(
comm: Communicator,
input: torch.Tensor,
algo: AllReduceAlgo,
pull_arg: PullArg,
) -> tvm_ffi.Tensor:
*,
graph_params: torch.Tensor | None = None,
use_multicast: bool = False,
) -> torch.Tensor:
module = get_all_reduce_module(input.dtype, comm.world_size)
return module.all_reduce(comm, input, algo.algo_name, pull_arg)
result = module.all_reduce(comm, input, algo.algo_name, graph_params, use_multicast)
return torch.from_dlpack(result)
@cache_once
python/sglang/kernels/ops/kimi_k3/all_reduce.py
+23 -47
View File
@@ -12,7 +12,7 @@ pull (2shot) :func:`all_reduce_pull_res` :func:`all_reduce_pull_norm`
* **push** — 1shot multicast-push. Works on ANY contiguous bf16 tensor
(input is read and written in place); reuses the CustomAllReduceV2 push
workspace, so the caller passes the workspace slab's multicast base.
plane, whose multicast base the plane itself carries.
Best for small messages. Needs :func:`register_comm`.
* **pull** — low-SM NVLS 2shot ON the input, which must be allocated from
multicast-bound symmetric memory (the caller passes its multicast VA):
@@ -20,10 +20,8 @@ pull (2shot) :func:`all_reduce_pull_res` :func:`all_reduce_pull_norm`
Launch geometry defaults to :data:`RES_TUNING` /
:data:`NORM_TUNING` and can be overridden per call via ``num_blocks`` /
``unroll`` (``num_blocks`` must be uniform across ranks per call).
Barriers reuse the CustomAllReduceV2 pull semaphores (same reservation
protocol as the generic pull kernels, signaled via one multicast red), so
:func:`register_comm` additionally needs the semaphore region's multicast
VA (``CustomAllReduceV2.pull_sem_mc_ptr``).
Barriers reuse the CustomAllReduceV2 barrier plane (same reservation
protocol as the generic pull kernels, signaled via one multicast red).
Epilogue contracts: the ``res`` residual must be identical on every rank (a
fully reduced tensor such as the attn-res prefix sum) or absent; the
@@ -74,20 +72,15 @@ def _jit_module(world_size: int) -> Module:
# Storage plane: the CustomAllReduceV2 Communicator
class _CommEntry(NamedTuple):
obj: Communicator # sgl.Communicator
pull_sem_mc_ptr: int
_COMM_MAP: dict[int, Communicator] = {}
_COMM_MAP: dict[int, _CommEntry] = {}
def register_comm(comm: Communicator) -> None:
"""Register the CustomAllReduceV2 communication planes.
def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None:
"""Register the CustomAllReduceV2 storage plane.
The push kernels only need ``comm``; the pull kernels additionally need
``pull_sem_mc_ptr`` (``CustomAllReduceV2.pull_sem_mc_ptr``), the
multicast VA of the pull-semaphore region their barriers reuse.
The push kernels use the push plane, the pull kernels the barrier plane;
both carry their own multicast base, so nothing else has to be threaded
through here.
"""
# world_size is the whole key, so at most one communicator per size can be
# registered in a process. That matches how these ops are called -- the custom
@@ -97,12 +90,12 @@ def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None:
# it instead of letting the overwrite happen; widening to per-group handles
# means changing the custom-op signatures, which is a separate change.
prev = _COMM_MAP.get(comm.world_size)
assert prev is None or prev.obj is comm, (
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] = _CommEntry(obj=comm, pull_sem_mc_ptr=pull_sem_mc_ptr)
_COMM_MAP[comm.world_size] = comm
class PullTuning(NamedTuple):
@@ -164,10 +157,8 @@ def _push_res_op(
world_size: int,
x: torch.Tensor,
residual: Optional[torch.Tensor],
ws_mc_base: int,
) -> None:
comm = _COMM_MAP[world_size].obj
_jit_module(world_size).push_res(comm, x.view(-1), residual, ws_mc_base)
_jit_module(world_size).push_res(_COMM_MAP[world_size], x.view(-1), residual)
@register_custom_op(mutates_args=["x"])
@@ -177,11 +168,9 @@ def _push_norm_op(
weight: torch.Tensor,
eps: float,
num_norm_rows: int,
ws_mc_base: int,
) -> None:
comm = _COMM_MAP[world_size].obj
_jit_module(world_size).push_norm(
comm, x.view(-1), weight, eps, num_norm_rows, ws_mc_base
_COMM_MAP[world_size], x.view(-1), weight, eps, num_norm_rows
)
@@ -194,18 +183,15 @@ def _finalize_push_norm_op(
expert_weights: torch.Tensor,
weight: torch.Tensor,
eps: float,
ws_mc_base: int,
) -> None:
comm = _COMM_MAP[world_size].obj
_jit_module(world_size).finalize_push_norm(
comm,
_COMM_MAP[world_size],
out.view(-1),
gemm2_out,
expanded_idx_to_permuted_idx,
expert_weights,
weight,
eps,
ws_mc_base,
)
@@ -218,13 +204,11 @@ def _pull_res_op(
num_blocks: int,
unroll: int,
) -> None:
entry = _COMM_MAP[world_size]
_jit_module(world_size).pull_res(
entry.obj,
_COMM_MAP[world_size],
x.view(-1),
residual,
input_mc_ptr,
entry.pull_sem_mc_ptr,
num_blocks,
unroll,
)
@@ -241,15 +225,13 @@ def _pull_norm_op(
num_blocks: int,
unroll: int,
) -> None:
entry = _COMM_MAP[world_size]
_jit_module(world_size).pull_norm(
entry.obj,
_COMM_MAP[world_size],
x.view(-1),
weight,
eps,
num_norm_rows,
input_mc_ptr,
entry.pull_sem_mc_ptr,
num_blocks,
unroll,
)
@@ -259,17 +241,14 @@ def all_reduce_push_res(
world_size: int,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
*,
ws_mc_base: int,
) -> torch.Tensor:
"""In-place ``x = allreduce(x) [+ residual]`` via 1shot multicast push.
``x`` may be any contiguous bf16 CUDA tensor whose byte size fits the
registered push workspace. ``ws_mc_base`` is the multicast VA of the v2
workspace slab base. Call :func:`register_comm` once beforehand.
``x`` may be any contiguous bf16 CUDA tensor whose byte size fits a slot
of the registered push plane. Call :func:`register_comm` once beforehand.
"""
residual_ = residual.view(-1) if residual is not None else None
_push_res_op(world_size, x, residual_, ws_mc_base)
_push_res_op(world_size, x, residual_)
return x
@@ -280,11 +259,10 @@ def all_reduce_push_norm(
eps: float,
*,
num_norm_rows: int,
ws_mc_base: int,
) -> torch.Tensor:
"""In-place allreduce via 1shot multicast push + RMSNorm over the first
``num_norm_rows`` rows of ``x`` viewed as [numel / 3584, 3584]."""
_push_norm_op(world_size, x, weight, eps, num_norm_rows, ws_mc_base)
_push_norm_op(world_size, x, weight, eps, num_norm_rows)
return x
@@ -296,8 +274,6 @@ def finalize_all_reduce_push_norm(
expert_weights: torch.Tensor,
weight: torch.Tensor,
eps: float,
*,
ws_mc_base: int,
) -> torch.Tensor:
"""Deferred MoE finalize + 1shot push all-reduce + RMSNorm on EVERY row.
@@ -314,7 +290,6 @@ def finalize_all_reduce_push_norm(
expert_weights,
weight,
eps,
ws_mc_base,
)
return out
@@ -331,8 +306,9 @@ def all_reduce_pull_res(
"""In-place ``x = allreduce(x) [+ residual]`` via low-SM NVLS 2shot.
``x`` MUST be allocated from multicast-bound symmetric memory and
``input_mc_ptr`` must be its multicast VA. Call :func:`register_comm`
(with ``pull_sem_mc_ptr``) once beforehand.
``input_mc_ptr`` must be its multicast VA (it varies per call, unlike the
barrier plane's own multicast base). Call :func:`register_comm` once
beforehand.
"""
tuning = _resolve_tuning(
RES_TUNING,
python/sglang/kernels/ops/kimi_k3/attn_res.py
+1 -5
View File
@@ -99,10 +99,9 @@ def _tuning(nvb: int, num_tokens: int) -> tuple[int, int, int]:
_COMM_MAP: dict[int, Communicator] = {}
_PULL_SEM_MC_MAP: dict[int, int] = {}
def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int) -> None:
def register_comm(comm: Communicator) -> None:
# One communicator per world_size per process -- see the note in
# kimi_k3/all_reduce.py::register_comm. The ops key only on world_size, so an
# overwrite here would hand the old group's callers the new group's peer
@@ -113,7 +112,6 @@ def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int) -> None:
f"{comm.world_size}"
)
_COMM_MAP[comm.world_size] = comm
_PULL_SEM_MC_MAP[comm.world_size] = pull_sem_mc_ptr
@register_custom_op(mutates_args=["out", "prefix_out"])
@@ -146,7 +144,6 @@ def _attn_res_fused_pull_rs_op(
nvb,
eps,
input_mc_ptr,
_PULL_SEM_MC_MAP[world_size],
max_blocks,
)
@@ -218,7 +215,6 @@ def _attn_res_fused_direct_ag_op(
nvb,
eps,
output_mc_ptr,
_PULL_SEM_MC_MAP[world_size],
max_blocks,
write_prefix,
)
python/sglang/kernels/ops/kimi_k3/gemm_ag.py
+3 -8
View File
@@ -61,10 +61,8 @@ def _gemm_ag_op(
b: torch.Tensor,
c: Optional[torch.Tensor],
out: torch.Tensor,
ws_mc_base: int,
) -> None:
comm = _COMM_MAP[world_size].obj
_jit_module().run(comm, x, weight, b, c, out, ws_mc_base)
_jit_module().run(_COMM_MAP[world_size], x, weight, b, c, out)
def gemm_ag_up_proj(
@@ -74,15 +72,12 @@ def gemm_ag_up_proj(
b: torch.Tensor,
c: Optional[torch.Tensor],
out: torch.Tensor,
*,
ws_mc_base: int,
) -> torch.Tensor:
"""``out = x @ weight.T (allgathered) + b (+ c)``, all bf16.
``x`` is [M, 3584] with M in [1, MAX_TOKENS]; ``weight`` is the FULL
replicated [7168, 3584] up_proj weight (each rank reads only its own
row block); ``b`` / ``c`` / ``out`` are [M, 7168] (``out`` is
output-only). ``ws_mc_base`` is the multicast VA of the v2 workspace
slab base (``comm.mc_base_ptr``)."""
_gemm_ag_op(world_size, x, weight, b, c, out, ws_mc_base)
output-only)."""
_gemm_ag_op(world_size, x, weight, b, c, out)
return out
python/sglang/kernels/ops/kimi_k3/sp_collective.py
+1 -9
View File
@@ -155,10 +155,9 @@ def _jit_module(world_size: int) -> Module:
_COMM_MAP: dict[int, Communicator] = {}
_PULL_SEM_MC_MAP: dict[int, int] = {}
def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None:
def register_comm(comm: Communicator) -> None:
# One communicator per world_size per process -- see the note in
# kimi_k3/all_reduce.py::register_comm. The ops key only on world_size, so an
# overwrite here would hand the old group's callers the new group's peer
@@ -169,7 +168,6 @@ def register_comm(comm: Communicator, *, pull_sem_mc_ptr: int = 0) -> None:
f"{comm.world_size}"
)
_COMM_MAP[comm.world_size] = comm
_PULL_SEM_MC_MAP[comm.world_size] = pull_sem_mc_ptr
@register_custom_op(mutates_args=["output"])
@@ -211,7 +209,6 @@ def _reduce_scatter_pull_op(
None if residual is None else residual.view(-1),
residual_is_local,
input_mc_ptr,
_PULL_SEM_MC_MAP[world_size],
num_blocks,
block_size,
)
@@ -222,7 +219,6 @@ def _all_gather_op(
world_size: int,
input: torch.Tensor,
output: torch.Tensor,
ws_mc_base: int,
num_blocks: int,
block_size: int,
) -> None:
@@ -230,7 +226,6 @@ def _all_gather_op(
_COMM_MAP[world_size],
input.view(-1),
output.view(-1),
ws_mc_base,
num_blocks,
block_size,
)
@@ -250,7 +245,6 @@ def _all_gather_direct_op(
input.view(-1),
output.view(-1),
output_mc_ptr,
_PULL_SEM_MC_MAP[world_size],
num_blocks,
block_size,
)
@@ -305,14 +299,12 @@ def all_gather(
input: torch.Tensor,
output: torch.Tensor,
*,
ws_mc_base: int,
tuning: Tuning = DEFAULT_TUNING,
) -> torch.Tensor:
_all_gather_op(
world_size,
input,
output,
ws_mc_base,
tuning.num_blocks,
tuning.block_size,
)
python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
+123 -99
View File
@@ -1,14 +1,22 @@
"""JIT custom all-reduce (v2) over a decoupled storage plane.
The CUDA side is split into two independent pieces:
The CUDA side is split into independent pieces:
- ``Communicator``: a thin pointer holder over symmetric-memory workspaces
(push buffers, pull buffer, semaphores) plus a local push counter. All
storage is allocated and owned here, in Python.
- the all-reduce kernel: a pure function of ``(input, Communicator, algo,
pull_arg)`` with three algorithms (1shot_push / 1shot_pull / 2shot_pull)
and three pull data sources (eager workspace / CUDA-graph pointer table /
multicast address).
- ``PushPlane``: the zero-filled symmetric push buffers plus a rank-local
phase counter.
- ``PullPlane``: the symmetric pull buffers plus the per-block semaphores
guarding them. The buffers exist only because this class hands the
all-reduce tensors that are not symmetric memory, so it stages them
through; the K3 fused collectives bring their own symmetric input and
borrow the semaphores alone.
- ``Communicator``: the two planes above (either may be absent), the handle
every kernel takes, plus the pull launch widths.
- the all-reduce kernel: a pure function of ``(Communicator, input, algo,
graph_params, use_multicast)`` with three algorithms (1shot_push /
1shot_pull / 2shot_pull) and three pull data sources (eager pull buffer /
CUDA-graph pointer table / multicast address).
All storage is allocated and owned here, in Python.
CUDA-graph inputs are exchanged from Python after capture (cudaIpc handles
for cudaMalloc-backed pointers, fabric/posix-fd VMM mapping for expandable
@@ -16,10 +24,9 @@ segments) and written into a device-side pointer table (``graph_params``);
the kernel captured in the graph dereferences its row at replay time.
"""
import enum
import logging
from contextlib import contextmanager
from typing import List, Optional, Tuple
from typing import List, NamedTuple, Optional, Tuple
import torch
import torch.distributed as dist
@@ -29,6 +36,8 @@ from sglang.kernels.ops.communication.all_reduce import (
AllReduceAlgo,
Communicator,
IPCManager,
PullPlane,
PushPlane,
custom_all_reduce,
)
from sglang.srt.distributed.parallel_state import in_the_same_node_as
@@ -67,12 +76,6 @@ _FORCE_PULL_SIZE_KB = envs.SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB.get()
_FORCE_PUSH_SIZE_KB = envs.SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB.get()
class _PullMode(enum.Enum):
EAGER = enum.auto() # pull_arg = False (also used for 1shot_push)
MULTICAST = enum.auto() # pull_arg = True
GRAPH = enum.auto() # pull_arg = a graph_params row
def _ceil_align(nbytes: int, align: int) -> int:
return (nbytes + align - 1) // align * align
@@ -95,6 +98,12 @@ def _allocate_symmetric_memory(nbytes: int, device: torch.device, group: Process
return tensor, symm_mem
class AllReduceConfig(NamedTuple):
algo: AllReduceAlgo
use_graph: bool = False
use_multicast: bool = False
class CustomAllReduceV2:
def __init__(
self,
@@ -112,10 +121,13 @@ class CustomAllReduceV2:
sized to what the tuned config wants, clipped to
this bound. Defaults to
``SGLANG_CUSTOM_ALL_REDUCE_V2_MAX_SIZE_KB`` (16 MB).
:param max_pull_size: explicit pull workspace size; overrides both
the tuned size and ``max_size``.
:param max_push_size: explicit per-buffer push workspace size;
:param max_pull_size: explicit pull buffer size; overrides both
the tuned size and ``max_size``. ``0`` builds a
push-only instance.
:param max_push_size: explicit per-slot push workspace size;
overrides both the tuned size and ``max_size``.
:param max_pull_blocks: cap on the barrier plane's block count; ``0``
builds a push-only instance.
``SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PULL_SIZE_KB`` /
``SGLANG_FORCE_CUSTOM_ALL_REDUCE_V2_PUSH_SIZE_KB`` take the highest
@@ -153,14 +165,22 @@ class CustomAllReduceV2:
graph=force_thresholds(base_config.graph),
eager=force_thresholds(base_config.eager),
)
# a minimal workspace keeps the Communicator valid even when a caller
# only uses one direction (e.g. push-only fused qk-norm instances)
self.max_pull_size = _ceil_align(max(max_pull_size, _ALIGN_BYTES), _ALIGN_BYTES)
# Zero on either pull knob opts out of the pull half entirely: no
# pull plane at all, and every pull algo disabled below by clipping
# its threshold to 0. Push-only callers (the fused qk-norm instances)
# take this path rather than allocating a placeholder buffer just to
# satisfy a constructor.
self.pull_enabled = max_pull_blocks != 0 and max_pull_size > 0
self.max_pull_size = (
_ceil_align(max(max_pull_size, _ALIGN_BYTES), _ALIGN_BYTES)
if self.pull_enabled
else 0
)
self.max_push_size = _ceil_align(max(max_push_size, _ALIGN_BYTES), _ALIGN_BYTES)
self.max_size = max(self.max_pull_size, self.max_push_size)
num_pull_blocks = base_config.num_pull_blocks
num_push_blocks = base_config.num_push_blocks
if max_pull_blocks is not None:
if max_pull_blocks:
num_pull_blocks = max(min(num_pull_blocks, max_pull_blocks), 1)
if max_push_blocks is not None:
num_push_blocks = max(max_push_blocks, 1)
@@ -195,84 +215,86 @@ class CustomAllReduceV2:
self.disabled = False
def _init_workspace(self) -> None:
"""Slice one symmetric-memory allocation into all shared buffers.
"""Slice one symmetric-memory allocation into every shared buffer.
Layout per rank: ``[2 * world_size push buffers | pull buffer |
pull semaphores]``. The push counter is rank-local, so it lives in a
plain CUDA tensor instead.
Layout per rank: ``[2 * world_size push slots | pull buffer | pull
semaphores]``. One allocation rather than three keeps it to a single
rendezvous and a single multicast base to offset from. The push
counter is rank-local, so it lives in a plain CUDA tensor.
"""
cfg = self.config
push_num_bufs = 2 * self.world_size # 2 phases x world_size peers
push_ws_bytes = push_num_bufs * self.max_push_size
pull_ws_bytes = self.max_pull_size
pull_sem_bytes = _SEMAPHORE_BYTES * cfg.num_pull_blocks
total_bytes = push_ws_bytes + pull_ws_bytes + pull_sem_bytes
pull_ws_offset = push_ws_bytes
pull_sem_offset = push_ws_bytes + pull_ws_bytes
push_num_slots = 2 * self.world_size # 2 phases x world_size peers
push_bytes = push_num_slots * self.max_push_size
pull_bytes = self.max_pull_size # 0 when the pull half is disabled
sem_bytes = _SEMAPHORE_BYTES * cfg.num_pull_blocks if self.pull_enabled else 0
total_bytes = push_bytes + pull_bytes + sem_bytes
pull_offset = push_bytes
sem_offset = push_bytes + pull_bytes
self._symm_tensor, symm_mem = _allocate_symmetric_memory(
total_bytes, device=self.device, group=self.group
)
workspaces = [
slabs = [
symm_mem.get_buffer(i, [total_bytes], torch.uint8)
for i in range(self.world_size)
]
workspaces[self.rank].zero_()
# The push slots (lamport pos-zero markers) and the semaphores must
# start zeroed; the pull buffer need not, but it rides along in the
# one-shot memset.
slabs[self.rank].zero_()
torch.cuda.synchronize()
dist.barrier(group=self.group)
def slice_ws(rank: int, shape: List[int], offset: int) -> torch.Tensor:
def slice_all(shape: List[int], offset: int) -> List[torch.Tensor]:
"""The same sub-range of every rank's slab, one view per rank."""
nbytes = 1
for s in shape:
nbytes *= s
for dim in shape:
nbytes *= dim
assert offset + nbytes <= total_bytes
return workspaces[rank][offset : offset + nbytes].view(shape)
return [slab[offset : offset + nbytes].view(shape) for slab in slabs]
multicast_ptr = int(symm_mem.multicast_ptr)
self.has_multicast = multicast_ptr != 0
def mc_at(offset: int) -> Optional[int]:
return multicast_ptr + offset if self.has_multicast else None
push_workspaces = [
slice_ws(i, [push_num_bufs, self.max_push_size], 0)
for i in range(self.world_size)
]
pull_workspaces = [
slice_ws(i, [pull_ws_bytes], pull_ws_offset) for i in range(self.world_size)
]
pull_semaphores = [
slice_ws(i, [cfg.num_pull_blocks, _SEMAPHORE_BYTES], pull_sem_offset)
for i in range(self.world_size)
]
self._push_counter = torch.zeros(
(cfg.num_push_blocks,), dtype=torch.uint32, device=self.device
)
multicast_ptr = int(symm_mem.multicast_ptr)
can_multicast = multicast_ptr != 0
# multicast VA of the slab base (== the push workspace, at offset 0);
# consumed by the K3 all_reduce push kernel
self.mc_base_ptr = multicast_ptr if can_multicast else 0
# multicast VA of the pull-semaphore region; the K3 pull kernels reuse
# these semaphores (same reservation protocol, multicast-signaled)
self.pull_sem_mc_ptr = multicast_ptr + pull_sem_offset if can_multicast else 0
pull_mc_workspace = multicast_ptr + pull_ws_offset if can_multicast else None
if not can_multicast or cfg.num_mc_blocks is None:
push_plane = PushPlane(
self.rank,
self.world_size,
workspaces=slice_all([push_num_slots, self.max_push_size], 0),
counter=self._push_counter.view(-1, 1).view(torch.uint8),
mc_workspace=mc_at(0),
)
pull_plane = None
if self.pull_enabled:
pull_plane = PullPlane(
self.rank,
self.world_size,
workspaces=slice_all([pull_bytes], pull_offset),
semaphores=slice_all(
[cfg.num_pull_blocks, _SEMAPHORE_BYTES], sem_offset
),
mc_workspace=mc_at(pull_offset),
mc_semaphore=mc_at(sem_offset),
)
if not self.has_multicast or not self.pull_enabled:
self.config = self.config._replace(num_mc_blocks=None)
self.obj = Communicator(
rank=self.rank,
world_size=self.world_size,
push_workspaces=push_workspaces,
pull_workspaces=pull_workspaces,
pull_semaphores=pull_semaphores,
push_counter=self._push_counter.view(-1, 1).view(torch.uint8),
pull_mc_workspace=pull_mc_workspace,
)
self.obj = Communicator(push=push_plane, pull=pull_plane)
if self.config.num_mc_blocks is not None:
self.obj.config(num_multicast_blocks=self.config.num_mc_blocks)
self.obj.set_pull_multicast_blocks(self.config.num_mc_blocks)
if self.rank == 0:
logger.info(
"All Reduce config: symmetric_memory = %.2f MB, "
"local_buffer = %.2f MB, multicast = %s",
"local_buffer = %.2f MB, multicast = %s, pull = %s",
total_bytes / MB,
(self.graph_params.nbytes + self._push_counter.nbytes) / MB,
self.config.num_mc_blocks is not None,
self.pull_enabled,
)
dist.barrier(group=self.group)
@@ -288,6 +310,9 @@ class CustomAllReduceV2:
custom-AR path can lift that cap up to ``max_pull_size``.
"""
if not self.pull_enabled:
return
def uncap(heuristic):
return heuristic._replace(two_shot_pull_threshold=self.max_pull_size)
@@ -307,21 +332,19 @@ class CustomAllReduceV2:
and torch.cuda.is_current_stream_capturing()
)
def _pick_algo(
self, nbytes: int, can_use_graph: bool
) -> Tuple[Optional[AllReduceAlgo], _PullMode]:
def _pick_config(self, nbytes: int, can_use_graph: bool) -> AllReduceConfig | None:
# TODO: refactor this along with the config file
heuristic = self.config.graph if can_use_graph else self.config.eager
default_mode = _PullMode.GRAPH if can_use_graph else _PullMode.EAGER
use_multicast = self.config.num_mc_blocks is not None
can_use_multicast = self.config.num_mc_blocks is not None
if nbytes <= heuristic.one_shot_push_threshold:
return AllReduceAlgo.ONE_SHOT_PUSH, _PullMode.EAGER
return AllReduceConfig(AllReduceAlgo.ONE_SHOT_PUSH)
if nbytes <= heuristic.one_shot_pull_threshold:
return AllReduceAlgo.ONE_SHOT_PULL, default_mode
if use_multicast and heuristic.mc.contains(nbytes):
return AllReduceAlgo.TWO_SHOT_PULL, _PullMode.MULTICAST
return AllReduceConfig(AllReduceAlgo.ONE_SHOT_PULL, use_graph=can_use_graph)
if can_use_multicast and heuristic.mc.contains(nbytes):
return AllReduceConfig(AllReduceAlgo.TWO_SHOT_PULL, use_multicast=True)
if nbytes <= heuristic.two_shot_pull_threshold:
return AllReduceAlgo.TWO_SHOT_PULL, default_mode
return None, _PullMode.EAGER
return AllReduceConfig(AllReduceAlgo.TWO_SHOT_PULL, use_graph=can_use_graph)
return None
def should_custom_ar(self, inp: torch.Tensor) -> bool:
"""Check if the input tensor is suitable for custom all-reduce."""
@@ -335,8 +358,7 @@ class CustomAllReduceV2:
return False
if self.override_algo is not None:
return inp_size <= self.max_size
algo, _ = self._pick_algo(inp_size, can_use_graph=self._can_use_graph())
return algo is not None
return self._pick_config(inp_size, self._can_use_graph()) is not None
# ------------------------------------------------------------------
# All-reduce
@@ -344,25 +366,27 @@ class CustomAllReduceV2:
def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
nbytes = input.numel() * input.element_size()
can_use_graph = self._can_use_graph()
if self.override_algo is not None:
# TODO: enhance this override pattern
algo = self.override_algo
use_graph = can_use_graph and not algo.is_push()
mode = _PullMode.GRAPH if use_graph else _PullMode.EAGER
use_graph = self._can_use_graph() and not algo.is_push()
use_multicast = False
else:
algo, mode = self._pick_algo(nbytes, can_use_graph=can_use_graph)
assert algo is not None, f"No algo for {nbytes} bytes"
if mode == _PullMode.GRAPH:
pull_arg: torch.Tensor | bool = self._allocate_graph_row(input, nbytes)
else:
pull_arg = mode == _PullMode.MULTICAST
return torch.from_dlpack(custom_all_reduce(self.obj, input, algo, pull_arg))
config = self._pick_config(nbytes, self._can_use_graph())
assert config is not None, f"No config for {nbytes = }"
algo, use_graph, use_multicast = config
graph_params = self._allocate_graph_row(input, nbytes) if use_graph else None
return custom_all_reduce(
self.obj,
input,
algo=algo,
graph_params=graph_params,
use_multicast=use_multicast,
)
def _allocate_graph_row(self, input: torch.Tensor, nbytes: int) -> torch.Tensor:
index = self._graph_counter + len(self._graph_inputs)
assert (
index < _MAX_GRAPH_INPUTS
), "Graph input table overflow, increase _MAX_GRAPH_INPUTS!"
assert index < _MAX_GRAPH_INPUTS, "Graph input table overflow"
self._graph_inputs.append((input.data_ptr(), nbytes))
return self.graph_params[index]
python/sglang/srt/layers/k3_ar_fusion.py
+3 -8
View File
@@ -93,7 +93,7 @@ def _get_state() -> Optional[_State]:
if (
not isinstance(comm, CustomAllReduceV2)
or comm.disabled
or comm.mc_base_ptr == 0
or not comm.has_multicast
):
if explicit:
logger.warning(
@@ -108,7 +108,7 @@ def _get_state() -> Optional[_State]:
return None
from sglang.kernels.ops.kimi_k3 import all_reduce as mod
mod.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
mod.register_comm(comm.obj)
logger.info("K3 all-reduce fusion enabled (world_size=%d)", comm.world_size)
return _State(comm, comm.world_size, group.cpu_group.group_name)
@@ -238,9 +238,7 @@ def all_reduce(
return x if residual is None else x + residual
nbytes = x.numel() * 2
if nbytes <= min(_PUSH_MAX_BYTES, state.comm.max_push_size):
return mod.all_reduce_push_res(
state.world_size, x, residual, ws_mc_base=state.comm.mc_base_ptr
)
return mod.all_reduce_push_res(state.world_size, x, residual)
return mod.all_reduce_pull_res(
state.world_size, x, residual, input_mc_ptr=get_mc_ptr(x)
)
@@ -297,7 +295,6 @@ def all_reduce_norm(
weight,
eps,
num_norm_rows=num_tokens,
ws_mc_base=state.comm.mc_base_ptr,
)
return mod.all_reduce_pull_norm(
state.world_size,
@@ -357,7 +354,6 @@ def gemm_ag_up_proj(
b,
c,
torch.empty_like(b),
ws_mc_base=state.comm.mc_base_ptr,
)
@@ -383,5 +379,4 @@ def finalize_all_reduce_push_norm(
expert_weights,
weight,
eps,
ws_mc_base=state.comm.mc_base_ptr,
)
python/sglang/srt/layers/k3_sp_collective.py
+3 -4
View File
@@ -76,7 +76,7 @@ def _init_state() -> Optional[_State]:
or a2a not in ("megamoe", "deepep")
or not isinstance(comm, CustomAllReduceV2)
or comm.disabled
or comm.mc_base_ptr == 0
or not comm.has_multicast
):
message = (
"K3 SP collective requires SM103, TP4/TP8, MegaMoE/DeepEP, and "
@@ -104,8 +104,8 @@ def _init_state() -> Optional[_State]:
)
return None
sp_collective.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
attn_res.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
sp_collective.register_comm(comm.obj)
attn_res.register_comm(comm.obj)
_STATE = _State(group, comm)
logger.info(
"K3 SP collective enabled (TP%d, fused RS residual + AG)",
@@ -393,7 +393,6 @@ def all_gather(tensor: torch.Tensor) -> Optional[torch.Tensor]:
state.group.world_size,
tensor,
output,
ws_mc_base=state.comm.mc_base_ptr,
tuning=dispatch.tuning,
)
if dispatch.strategy == "direct":
python/sglang/srt/models/minimax_m2.py
+1
View File
@@ -442,6 +442,7 @@ class MiniMaxM2QKRMSNorm:
comm = CustomAllReduceV2(
group=get_parallel().attn_tp_group.cpu_group,
device=device,
# push-only: no barrier plane and no staging buffer
max_pull_size=0,
max_pull_blocks=0,
max_push_size=max_size,
test/registered/kernels/ops/kimi_k3/test_ar_fusion.py
+9 -11
View File
@@ -102,9 +102,9 @@ def _init_comm() -> CustomAllReduceV2:
comm = CustomAllReduceV2(
cpu_group, device, max_pull_size=1 * MB, max_push_size=2 * MB
)
if comm.disabled or comm.mc_base_ptr == 0:
if comm.disabled or not comm.has_multicast:
raise RuntimeError("ar_fusion requires CustomAllReduceV2 with multicast")
all_reduce.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
all_reduce.register_comm(comm.obj)
register_comm_cleanup(comm)
return comm
@@ -165,7 +165,7 @@ def test_ar_fusion_push(bs: int, use_residual: bool):
x = _int_input(n, bs, per_rank=True)
residual = _int_input(n, bs + 7, per_rank=False) if use_residual else None
ref = _nccl_ref(x, residual)
all_reduce.all_reduce_push_res(world, x, residual, ws_mc_base=comm.mc_base_ptr)
all_reduce.all_reduce_push_res(world, x, residual)
torch.cuda.synchronize()
torch.testing.assert_close(x, ref, atol=0, rtol=0)
@@ -240,9 +240,7 @@ def test_ar_fusion_push_norm(num_tokens: int, rows_per_token: int):
x = _int_input(n, num_tokens + 41 + rows_per_token, per_rank=True)
weight = _int_input(NORM_DIM, 43, per_rank=False) + 1
ref = _norm_ref(_nccl_ref(x, None), num_tokens, weight, eps=1e-6)
all_reduce.all_reduce_push_norm(
world, x, weight, 1e-6, num_norm_rows=num_tokens, ws_mc_base=comm.mc_base_ptr
)
all_reduce.all_reduce_push_norm(world, x, weight, 1e-6, num_norm_rows=num_tokens)
torch.cuda.synchronize()
_assert_norm_close(x, ref, num_tokens)
@@ -310,7 +308,7 @@ def test_ar_fusion_finalize_push_norm(bs: int):
ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps)
out = torch.empty(bs, NORM_DIM, dtype=torch.bfloat16, device=_device())
all_reduce.finalize_all_reduce_push_norm(
world, out, gemm2, idx, weights, norm_w, eps, ws_mc_base=comm.mc_base_ptr
world, out, gemm2, idx, weights, norm_w, eps
)
torch.cuda.synchronize()
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
@@ -331,12 +329,12 @@ def test_ar_fusion_finalize_push_norm_stress():
ref = _finalize_norm_ref(gemm2, idx, weights, norm_w, eps)
out = torch.empty(bs, NORM_DIM, dtype=torch.bfloat16, device=_device())
all_reduce.finalize_all_reduce_push_norm(
world, out, gemm2, idx, weights, norm_w, eps, ws_mc_base=comm.mc_base_ptr
world, out, gemm2, idx, weights, norm_w, eps
)
torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
x = _int_input(bs * H, 8000 + it, per_rank=True)
ref2 = _nccl_ref(x, None)
all_reduce.all_reduce_push_res(world, x, None, ws_mc_base=comm.mc_base_ptr)
all_reduce.all_reduce_push_res(world, x, None)
torch.testing.assert_close(x, ref2, atol=0, rtol=0)
@@ -382,7 +380,7 @@ def test_ar_fusion_stress_mixed():
num_blocks = (1, 2, 4, 8)[it % 4]
x = _int_input(n, 3000 + it, per_rank=True)
ref = _nccl_ref(x, None)
all_reduce.all_reduce_push_res(world, x, None, ws_mc_base=comm.mc_base_ptr)
all_reduce.all_reduce_push_res(world, x, None)
torch.testing.assert_close(x, ref, atol=0, rtol=0)
y = buf[:n]
y.copy_(_int_input(n, 4000 + it, per_rank=True))
@@ -407,7 +405,7 @@ def test_ar_fusion_graph_capture():
gz, mc_z = buf[n : 2 * n], mc + n * buf.element_size()
def _run_all():
all_reduce.all_reduce_push_res(world, gx, gres, ws_mc_base=comm.mc_base_ptr)
all_reduce.all_reduce_push_res(world, gx, gres)
all_reduce.all_reduce_pull_res(world, gy, gres, input_mc_ptr=mc_y)
all_reduce.all_reduce_pull_res(world, gz, gres, input_mc_ptr=mc_z)
test/registered/kernels/ops/kimi_k3/test_collectives.py
+4 -7
View File
@@ -69,11 +69,11 @@ def _init_comm():
max_pull_size=4 * _MB,
max_push_size=4 * _MB,
)
if comm.disabled or comm.mc_base_ptr == 0:
if comm.disabled or not comm.has_multicast:
raise RuntimeError("Kimi K3 collectives require multicast symmetric memory")
all_reduce.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
sp_collective.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
attn_res.register_comm(comm.obj, pull_sem_mc_ptr=comm.pull_sem_mc_ptr)
all_reduce.register_comm(comm.obj)
sp_collective.register_comm(comm.obj)
attn_res.register_comm(comm.obj)
register_comm_cleanup(comm)
return comm
@@ -138,7 +138,6 @@ def test_all_reduce_push():
comm.world_size,
x,
residual,
ws_mc_base=comm.mc_base_ptr,
)
torch.cuda.synchronize()
torch.testing.assert_close(x, expected, rtol=0, atol=0)
@@ -204,7 +203,6 @@ def test_sequence_parallel_collectives():
world_size,
gather_input,
gather_output,
ws_mc_base=comm.mc_base_ptr,
tuning=_SP_TUNING,
)
torch.cuda.synchronize()
@@ -243,7 +241,6 @@ def test_gemm_all_gather():
bias,
None,
output,
ws_mc_base=comm.mc_base_ptr,
)
torch.cuda.synchronize()
torch.testing.assert_close(output, expected, rtol=3e-2, atol=3e-2)